Amortized Tree Generation for Bottom-Up Synthesis Planning and Synthesizable Molecular Design
Introduction
이 글은 Gao, Mercado, Coley가 ICLR 2022에서 발표한 “Amortized Tree Generation for Bottom-Up Synthesis Planning and Synthesizable Molecular Design” 논문을 정리한다.
논문에서는 molecular design과 synthesis planning을 분리된 task로 다루는 기존 pipeline의 한계를 지적하고, 두 문제를 conditional synthetic pathway generation이라는 하나의 task로 통합하여 해결하는 framework를 제안한다.
저자들은 target molecular embedding에 conditioned된 Markov decision process (MDP)로 synthetic tree 생성을 모델링한다.
이를 통해 bottom-up 방식의 synthesis planning과 optimized conditional code로부터의 synthesizable molecule design이 동시에 가능하다고 주장한다.
Contribution
논문에서 제시하는 주요 contribution은 다음과 같다.
- Synthetic tree 생성을 MDP로 정식화하여 multi-step 및 convergent (nonlinear) synthetic pathway 생성을 가능하게 했다고 설명하고 있다.
- Reaction template과 purchasable building block으로 정의된 discrete action space 위에서 rapid bottom-up synthesis planning과 constrained molecular optimization을 동시에 수행하는 model을 제안한다고 이야기한다.
- 복잡한 organic molecule에 대해 amortized multi-step synthesis planning을 성공적으로 수행한 최초의 시도라고 주장한다.
- Drug discovery 관련 oracle function에 대한 de novo molecular optimization에서 유의미한 결과를 보였다고 보고한다.
Method
Problem Definition
논문에서는 synthetic pathway를 synthetic tree 구조로 모델링한다.
Valid synthetic tree는 하나의 root node (최종 product molecule)를 가지며, purchasable building block으로부터 시작하여 reaction template에 의해 정의된 feasible reaction을 통해 root node에 도달한다.
위 figure는 synthetic tree와 reaction template의 예시를 보여준다. (A)는 COVID-19 치료제 remdesivir의 synthetic tree이고, (B)와 (C)는 각각 uni-molecular 및 bi-molecular reaction template의 SMARTS 표현 예시이다.
논문에서는 두 가지 task를 다음과 같이 정의한다.
- Synthesis Planning: 주어진 target molecule \(M_\text{target}\)에 대해, root node의 product molecule \(M_\text{product}\)가 \(M_\text{target}\)과 일치하는 synthetic tree \(T\)를 생성하는 문제이다.
- Synthesizable Molecular Design: Product molecule \(M_\text{product}\)의 properties가 oracle function 기준으로 최적화되도록 synthetic tree \(T\)의 구조를 최적화하는 문제이다.
Synthetic Tree Generation as a Markov Decision Process
저자들은 synthetic tree 생성 과정이 Markov property를 자연스럽게 만족한다고 설명한다.
특정 intermediate compound가 주어지면, 이후 reaction step은 해당 compound와 target molecule만으로 결정되며 이전 pathway에는 의존하지 않는다고 본다.
Markov property는 다음과 같이 표현된다.
\[p(S^{(t+1)} | S^{(t)}, \ldots, S^{(0)}) = p(S^{(t+1)} | S^{(t)})\]위 식에서 \(S^{(t)}\)는 step \(t\)에서의 state로, intermediate synthetic tree \(T^{(t)}\)의 root molecule(s)로 정의된다.
논문에서는 최대 두 개의 disconnected sub-tree만 동시에 존재할 수 있도록 제약을 두고, reverse depth-first order로 generation이 진행되도록 강제한다.
위 figure는 iterative generation procedure를 보여준다. Model은 building block으로부터 시작하여 점진적으로 복잡한 molecule을 만들어가는 bottom-up 방식으로 synthetic tree를 구성한다.
Action Space는 네 가지 component로 분해된다고 설명하고 있다.
- \(a_\text{act}\): action type으로, “Add”, “Expand”, “Merge”, “End” 중에서 sampling한다.
- \(a_\text{rt1}\): first reactant로, building block set \(\mathcal{C}\) 또는 \(M_\text{most\_recent}\)에서 sampling한다.
- \(a_\text{rxn}\): reaction template로, template set \(\mathcal{R}\)에서 sampling한다.
- \(a_\text{rt2}\): second reactant로, \(\mathcal{C}\)에서 sampling한다.
각 action type의 역할은 다음과 같이 설명된다.
- Add: 새로운 reactant node를 추가하고 product node를 생성한다. 항상 synthetic tree 구축의 첫 action으로 사용된다.
- Expand: \(M_\text{most\_recent}\)을 first reactant로 사용하여 새로운 reaction step을 진행한다.
- Merge: 두 root node를 bi-molecular reaction의 reactant로 사용하여 sub-tree를 합친다.
- End: synthetic tree를 종료하고 마지막 product를 \(M_\text{product}\)로 확정한다.
State Transition Dynamics는 valid action이 선택되면 deterministic하다고 설명한다.
Template을 따르지 않는 infeasible action은 rejection되며, template 적용 결과 생성된 structure가 명시적으로 새로운 state에 포함된다.
저자들은 Bradshaw et al. (2020)의 RNN 기반 model이 reaction predictor의 dynamics를 implicit하게 학습해야 했던 것과 비교하여, 이 점이 이 model의 장점이라고 주장한다.
Conditional Generation Model
논문에서는 MDP를 해결하기 위해 네 개의 module로 구성된 model을 제안한다.
위 figure는 model의 전체 구조를 보여준다. 각 step에서 최대 두 개의 root molecule embedding과 target molecule embedding이 state를 정의하며, network들이 순차적으로 action type, first reactant, reaction template, second reactant를 예측한다.
각 module은 다음과 같은 확률 분포를 예측한다.
\[a^{(t)}_\text{act} \sim f_\text{act}(S^{(t)}, M_\text{target}) = \sigma(\text{MLP}_\text{act}(z^{(t)}_\text{state} \oplus z_\text{target}))\] \[a^{(t)}_\text{rt1} \sim f_\text{rt1}(S^{(t)}, M_\text{target}) = \text{k-NN}_\mathcal{C}(\text{MLP}_\text{rt1}(z^{(t)}_\text{state} \oplus z_\text{target}))\] \[a^{(t)}_\text{rxn} \sim f_\text{rxn}(S^{(t)}, a^{(t)}_\text{rt1}, M_\text{target}) = \sigma(\text{MLP}_\text{rxn}(z^{(t)}_\text{state} \oplus z_\text{target} \oplus z^{(t)}_\text{rt1}))\] \[a^{(t)}_\text{rt2} \sim f_\text{rt2}(S^{(t)}, a^{(t)}_\text{rt1}, a^{(t)}_\text{rxn}, M_\text{target}) = \text{k-NN}_{\mathcal{C}'}(\text{MLP}_\text{rt2}(z^{(t)}_\text{state} \oplus z_\text{target} \oplus z^{(t)}_\text{rt1} \oplus z^{(t)}_\text{rxn}))\]위 수식에서 \(\oplus\)는 concatenation을 의미하고, \(z_\ast\)는 해당 entity의 embedding을 나타낸다. \(\mathcal{C}'\)은 선택된 template \(a^{(t)}_\text{rxn}\)과 호환되지 않는 reactant가 masking된 \(\mathcal{C}\)의 subset이다. \(\sigma\)는 invalid action을 masking하는 softmax 함수이다.
\(f_\text{act}\)와 \(f_\text{rxn}\)은 classification network로 학습되고, \(f_\text{rt1}\)과 \(f_\text{rt2}\)는 first 및 second reactant의 embedding을 regression으로 예측한 뒤 k-NN search를 통해 후보 molecule을 찾는다.
Synthetic Tree Generation Algorithm
위 algorithm은 synthetic tree generation의 전체 흐름을 보여준다.
- Line 1–4에서는 input (reaction template, building block, target molecule, network)을 정의하고, state \(S\), tree \(T\), \(M_\text{most\_recent}\)를 초기화한다고 설명하고 있다.
- Line 5–6에서는 generation loop를 시작하며 action type을 예측하고 sampling한다고 볼 수 있다.
- Line 7에서는 first reactant embedding을 예측한다.
- Line 8–16에서는 action type에 따라 first reactant를 결정하는 부분이라고 설명한다. End action이면 loop를 종료하고, Add action이면 k-NN으로 building block에서 reactant를 선택하며, 그 외의 경우 \(M_\text{most\_recent}\)를 first reactant로 사용한다.
- Line 17에서는 state, target, first reactant embedding을 기반으로 reaction template을 예측한다고 이야기한다.
- Line 18–27에서는 bi-molecular reaction인 경우 second reactant를 결정하는 과정이다. Merge action이면 다른 root molecule을, 그 외에는 network로 예측된 embedding을 k-NN search하여 building block에서 선택한다.
- Line 28–30에서는 선택된 template을 reactant에 적용하여 product를 생성하고, tree와 state, \(M_\text{most\_recent}\)를 업데이트한다고 설명하고 있다.
Genetic Algorithm for Molecular Optimization
논문에서는 synthesizable molecular design을 위해 target embedding \(z_\text{target}\)을 numerical optimization하는 방식으로 접근한다.
저자들은 reinforcement learning보다 절차적으로 단순한 genetic algorithm (GA)을 채택했다고 설명한다.
위 figure는 genetic algorithm의 동작 과정을 보여준다. Conditional synthetic tree generator가 decoder 역할을 하며, fingerprint pool 위에서 crossover와 mutation을 수행하여 molecule을 implicit하게 최적화한다.
각 generation에서 mating pool에서 두 vector를 sampling하여 crossover로 offspring을 만들고, 작은 확률로 mutation을 적용한 뒤 fitness function으로 평가하여 top performer를 다음 mating pool로 유지한다.
Experiment
Experiment Setup
논문에서 사용한 실험 설정은 다음과 같다.
- Reaction Templates: Hartenfeller et al. (2011)과 Button et al. (2019)의 template set을 결합하여 91개 reaction template을 사용한다. 13개는 uni-molecular, 78개는 bi-molecular이다.
- Purchasable Building Blocks: Enamine Building Blocks 중 template에 매칭되는 147,505개의 molecule을 사용한다.
- Dataset: random policy로 생성한 synthetic tree 중 QED filtering을 거쳐 training 208,644개, validation 69,548개, test 69,548개를 확보했다고 설명한다.
- Representations: MLP input에는 길이 4096, radius 2 of Morgan fingerprint를, k-NN module에는 길이 256, radius 2 of Morgan fingerprint를 사용한다.
- Optimization Oracle: QED, LogP, JNK3, GSK3β, DRD2 surrogate model 및 AutoDock Vina를 이용한 docking score (DRD3, SARS-Cov-2 Mpro)를 사용한다.
Result
Synthesis Planning Results
논문에서는 reachable (test set)과 unreachable (ChEMBL random sample) target에 대해 reconstruction 성능을 평가한다.
저자들은 reachable molecule의 51%를 reconstruct하는 데 성공했다고 보고한다.
기존 top-down retrosynthesis가 분당 단위의 시간을 요구하는 데 비해, 이 bottom-up 접근은 \(k=1\) greedy construction 기준 약 1초 만에 tree를 생성한다고 강조한다.
ChEMBL recovery rate는 4.5%로 낮지만, template과 starting material을 확장하면 architecture 변경 없이 더 높은 recovery가 가능하다고 설명한다.
위 figure는 ChEMBL을 target으로 한 conditional generation의 세 가지 결과를 보여준다. (A)는 training set과의 유사도가 낮은데도 성공적으로 reconstruction된 사례, (B)는 reconstruct는 실패했지만 synthesizable analog로 해석할 수 있는 사례, (C)는 target과 유사한 부분 구조를 가진 synthetic route를 제안하여 chemist에게 inspiration을 줄 수 있는 사례이다.
네 module 중 가장 큰 error는 first reactant network \(f_\text{rt1}\)에서 발생하며 validation accuracy가 30.8%에 그친다고 보고한다.
반면 action network는 99% 이상, reaction network는 85.8%의 validation accuracy를 달성했다고 이야기한다.
Synthesizable Analog Recommendation Results
논문에서는 unrecovered case의 product를 synthesizable structural analog로 활용할 수 있다고 주장한다.
위 figure는 target과 product molecule의 property correlation을 보여준다. SA Score, LogP, molecular weight, QED 모두 양의 correlation을 보이며, QED가 가장 낮은 correlation을 보이는 이유는 structure-activity landscape index (SALI) 관점에서 QED가 구조 변화에 매우 민감하기 때문이라고 해석한다.
저자들은 desired property가 분자 구조와 강하게 correlated될수록 model이 합리적인 analog를 추천할 수 있다고 설명하고 있다.
Synthesizable Molecular Optimization Results
논문에서는 heuristic oracle function에 대한 optimization 성능을 baseline과 비교한다.
저자들은 synthesizability constraint가 없는 baseline 대비 점수는 약간 낮을 수 있지만, model이 GCPN과 MolDQN을 일관되게 상회하며 GA+D, MARS와 비교 가능한 수준이라고 보고한다.
위 figure는 GSK3β bioactivity optimization과 SARS-Cov-2 Mpro docking score optimization 결과를 비교한다. (A)에서 이 model이 제안한 GSK3β inhibitor는 DST나 MARS보다 구조적으로 훨씬 단순하며, 단일 Suzuki reaction step으로 합성 가능한 actionable 구조라고 강조한다. (B)에서는 known inhibitor (Vina score -8.96 kcal/mol) 대비 더 강한 predicted binding affinity를 가지는 multiple molecule을 제안했다고 보고한다.
저자들은 TDC DRD3 docking benchmark에서 이 model이 낮은 SA_Score와 높은 Pass rate를 동시에 달성한 유일한 method라고 주장한다.
Limitation
논문에서는 다음과 같은 한계를 명시적으로 언급한다.
- Initial Reactant Selection: First reactant selection이 target molecule recovery의 주된 bottleneck이라고 설명한다. Input이 target molecule뿐이고 action mask가 적용되지 않으며, action space가 너무 크다는 점이 원인이라고 본다. 또한 depth-first generation이 implicit canonical ordering을 도입하여 학습을 어렵게 할 수 있다고 이야기한다.
- Reaction Templates and Purchasable Compound Selection: 91개 template은 NameRxn의 약 1,300개, SYNTHIA의 수만 개, ASKCOS의 수십만 개에 비해 매우 적은 수준이라고 인정한다. 더 큰 template set을 사용하면 탐색 가능한 chemical space가 확장될 것이라고 설명한다.
- Molecular Representation and Similarity: Boolean Morgan fingerprint가 repeating unit이나 symmetry를 구분하지 못하는 한계가 있다고 이야기한다. Count 기반 representation과 그에 맞는 similarity measure 개발이 필요하다고 본다.
- Beam Search for Decoding: 현재 greedy 또는 single-path sampling 방식을 사용하지만, beam search를 도입하면 bottom-up synthesis planning 성능이 향상될 수 있다고 제안한다.
다만 이 model은 synthesizability constraint가 없는 baseline과 비교했을 때 절대적인 oracle score는 다소 낮을 수 있는데, 이는 동일한 trade-off 관점에서 추가 검증이 필요하다고 볼 수 있다.
Summary
논문에서는 synthetic tree generation을 conditional MDP로 정식화하여 synthesis planning과 synthesizable molecular design을 하나의 framework로 통합하는 SynNet을 제안한다.
저자들은 bottom-up 방식이 기존 retrosynthesis 기반 top-down 방법보다 훨씬 빠른 amortized inference를 가능하게 한다고 주장한다.
또한 GA를 이용한 latent embedding optimization을 통해 synthesizable analog recommendation과 de novo design을 동시에 수행할 수 있음을 실험적으로 보였다고 보고한다.
향후 first reactant selection 개선, 더 큰 template/building block set 활용, count-based fingerprint 도입, beam search decoding 등의 방향이 남아 있다고 정리한다.
References
Gao, Wenhao, Rocío Mercado, and Connor W. Coley. “Amortized tree generation for bottom-up synthesis planning and synthesizable molecular design.” arXiv preprint arXiv:2110.06389 (2021).