RxnFlow: Generative Flows on Synthetic Pathway for Drug Design

Created on August 19, 2025

2025   ·   GFlowNet   Synthesizable Molecular Generation   Structure-based Drug Design   Action Space Subsampling   ·   Paper Review

Introduction

Structure-based drug discovery (SBDD)는 protein structure prediction tool의 발전과 고해상도 crystallography 데이터 축적에 힘입어 초기 신약 개발의 핵심 paradigm으로 자리 잡았다. 논문에서는 기존의 brute-force virtual screening이 계산 비용이 크다는 한계가 있다고 이야기하며, 이를 우회하기 위해 deep generative model이 활발히 연구되었다고 설명하고 있다.

그러나 저자들은 대부분의 generative model이 synthesizability를 고려하지 않아 wet-lab validation 단계에서의 활용성이 제한된다고 주장한다. 합성 가능성을 보조 함수로 multi-objective optimization에 포함하는 접근은 복잡한 합성 원리를 충분히 반영하지 못하고, 생성된 분자를 synthesizable space로 projection하는 접근은 최적화된 특성을 훼손할 수 있다는 한계가 있다고 이야기한다.

이러한 문제를 해결하기 위해 저자들은 RxnFlow를 제안한다. RxnFlow는 building block과 reaction template을 기반으로 synthetic pathway를 sequential하게 구성하고, generative flow networks (GFlowNets) objective로 학습하여 reward에 비례한 분포로 다양한 molecule을 sampling한다고 설명하고 있다.

Figure 1. Overview of RxnFlow. (a) Synthetic action space which is represented in a continuous action space. Each colored box corresponds to a reaction template and the molecules in the box are reactant blocks. (b) Policy estimation using the action space subsampling in a manner of importance sampling. (c) Molecular generation process and model training.

Contribution

논문에서는 RxnFlow의 주요 기여를 다음과 같이 정리하고 있다.

  • 1.2M개의 building block과 71개의 reaction template으로 구성된 거대한 action space 위에서 generative flow를 학습할 수 있는 synthesis-oriented framework를 제안한다.
  • Large action space에서의 memory 및 computational overhead를 완화하기 위해 importance sampling 기반의 action space subsampling을 도입한다.
  • Non-hierarchical MDP와 action embedding을 결합하여, 재학습 없이도 새로운 building block library나 추가 objective로 확장 가능한 구조를 제시한다.
  • LIT-PCBA 기반 pocket-specific optimization과 CrossDocked2020 기반 pocket-conditional generation 모두에서 기존 reaction-based, atom-based, fragment-based baseline을 능가하는 결과를 보고한다.

Method

GFlowNet Preliminaries

논문에서는 GFlowNet을 reward $R(x)$에 비례한 확률 $p(x) \propto R(x)$로 object $x$를 sampling하는 generative model로 정의하고 있다. Object $x$는 state transition $s \to s’$의 sequence를 따라 trajectory $\tau = (s_0 \to \cdots \to s_n = x)$로 점진적으로 구성되며, 전체 trajectory의 집합은 directed acyclic graph $G = (S, A)$로 표현된다.

Trajectory flow는 terminal state의 reward로 정의되고, edge flow $F(s \to s’)$는 해당 edge를 지나는 trajectory들의 flow 합으로 정의된다.

\[F(s) = \sum_{(s'' \to s) \in A} F(s'' \to s) = \sum_{(s \to s') \in A} F(s \to s')\]

위 수식에서 $F(s)$는 state $s$를 지나는 trajectory들의 total flow를 의미한다. Intermediate state에서는 들어오는 edge flow의 합과 나가는 edge flow의 합이 같아야 한다는 flow matching condition을 만족해야 한다.

Forward policy와 backward policy는 각각 다음과 같이 정의한다.

\[P_F(s' \mid s) = \frac{F(s \to s')}{F(s)}, \quad P_B(s \mid s') = \frac{F(s \to s')}{F(s')}\]

여기서 $P_F$는 forward transition probability를, $P_B$는 backward transition probability를 의미한다.

Action Space for Synthetic Pathway Generation

저자들은 chemical reaction을 forward transition으로, synthetic pathway를 trajectory로 취급한다. 초기 state $s_0$에서는 항상 AddFirstReactant를 통해 building block set $\mathcal{B}$에서 starting molecule을 선택하고, 이후 state $s$에서는 ReactUni, ReactBi, Stop 중 하나의 action을 선택한다고 설명하고 있다.

State $s$에서의 allowable action space $\mathcal{A}(s)$는 다음과 같이 정의한다.

\[\mathcal{A}(s) = \begin{cases} \mathcal{B} & \text{if } s = s_0 \\\ \{\text{Stop}\} \cup \mathcal{R}_1 \cup \{(r, b) \mid r \in \mathcal{R}_2, b \in \mathcal{B}_r\} & \text{otherwise} \end{cases}\]

위 수식에서 $\mathcal{R}_1$은 uni-molecular reaction template 집합, $\mathcal{R}_2$는 bi-molecular reaction template 집합, $\mathcal{B}_r \subseteq \mathcal{B}$는 reaction template $r$에 사용 가능한 reactant block 집합을 의미한다. 현재 분자에 적용 불가능한 reaction template은 mask 처리한다고 이야기한다.

Action Space Subsampling

논문에서는 거대한 action space 문제를 해결하기 위해 action space subsampling을 제안한다. 이는 전체 outgoing edge flow가 아닌 일부 subset에서 state flow $\hat{F}_\theta(s)$를 추정하는 memory-efficient 기법이라고 설명하고 있다.

Subsampling policy \(\mathcal{P}(\mathcal{A})\)는 initial state에서는 uniform sampling을, 이후 state에서는 importance sampling을 수행한다. 희귀 action을 보존하기 위해 Stop과 ReactUni는 모두 사용하고, ReactBi의 reactant block만 uniform subset \(\mathcal{B}_r^{\ast} \subseteq \mathcal{B}_r\)로 sampling한다고 이야기한다.

\[\mathcal{A}^*(s) = \begin{cases} \mathcal{B}^{\ast} & \text{if } s = s_0 \\\ \{\text{Stop}\} \cup \mathcal{R}_1 \cup \{(r, b) \mid r \in \mathcal{R}_2, b \in \mathcal{B}_r^{\ast}\} & \text{otherwise} \end{cases}\]

여기서 $\mathcal{B}^{\ast}$는 initial state에서 sampling된 building block subset, $\mathcal{B}_r^{\ast}$는 bi-molecular reaction $r$에 대해 sampling된 reactant block subset을 의미한다.

Importance sampling 보정을 위해 edge flow에는 subsampling ratio에 따른 weight $w_a$를 곱한다.

\[w_a = \begin{cases} \lvert\mathcal{B}\rvert / \lvert\mathcal{B}^*\rvert & \text{if } a \in \mathcal{B} \\\ \lvert\mathcal{B}_r\rvert / \lvert\mathcal{B}_r^*\rvert & \text{if } a \text{ is } (r, b), r \in \mathcal{R}_2 \\\ 1 & \text{if } a \text{ is Stop or } a \in \mathcal{R}_1 \end{cases}\]

위 weight를 적용하여 state flow를 다음과 같이 추정한다.

\[\hat{F}_\theta(s; \mathcal{A}^*) = \sum_{(s \to s') \in \mathcal{A}^*(s)} w_{(s \to s')} F_\theta(s \to s')\]

여기서 \(\hat{F}_\theta(s; \mathcal{A}^{\ast})\)는 partial action space \(\mathcal{A}^{\ast}\)로부터 추정된 state flow를 의미한다. 추정된 forward policy는 \(\hat{P}_F(s' \mid s; \mathcal{A}^{\ast}; \theta) = F_\theta(s \to s')/\hat{F}_\theta(s; \mathcal{A}^{\ast})\)로 정의된다고 설명한다.

논문에서는 이 추정이 unbiased estimator임을 부록 A에서 증명하고 있으며, 계산 복잡도를 $O(\lvert\mathcal{B}\rvert \lvert\mathcal{R}_2\rvert)$에서 $O(\lvert\mathcal{B}^{\ast}\rvert \lvert\mathcal{R}_2\rvert)$로 감소시킬 수 있다고 주장한다.

Action Embedding

거대한 action space는 action-specific parameter $\theta_a$가 많아지는 문제를 야기한다. 이를 해결하기 위해 building block $b$를 molecular fingerprint 등 structural information으로 continuous action space에 embedding하는 별도의 network $\phi_\theta^{\text{block}}$을 사용한다고 설명하고 있다.

\[F_\theta(s_0, b) = \phi_\theta^{\text{flow}}(\phi_\theta^{\text{state}}(s), \phi_\theta^{\text{block}}(b))\] \[F_\theta(s, (r, b)) = \phi_\theta^{\text{flow}}(\phi_\theta^{\text{state}}(s), \delta(r), \phi_\theta^{\text{block}}(b))\]

위 수식에서 $\phi_\theta^{\text{state}}$는 state $s$의 representation을 생성하는 graph transformer를 의미하고, $\delta(r)$은 bi-molecular reaction template $r$에 대한 one-hot encoding을 의미한다. $\phi_\theta^{\text{block}}$은 building block $b$를 continuous embedding으로 변환하는 MLP를 나타낸다.

Joint Selection of Templates and Blocks

기존의 synthesis-oriented method는 reaction template $r$을 먼저 선택한 후 reactant block $b$를 선택하는 hierarchical MDP를 사용한다고 이야기한다. 그러나 hierarchical MDP에서는 학습 이후 reaction template의 선택 확률이 고정되어, building block library가 변경되면 policy 추정에 bias가 생긴다고 주장한다.

Figure 2. Comparison of using modified building block library for the generation: (a) a hierarchical MDP, and (b) a non-hierarchical MDP. More details are in Figure 8.

이를 해결하기 위해 저자들은 reaction template과 reactant block을 한 번에 joint하게 선택하는 non-hierarchical MDP를 제안한다.

\[P_F(T(s, (r, b)) \mid s; \theta) = \frac{F_\theta(s, (r, b))}{\sum_{r' \in \mathcal{R}_1 \cup \{\text{Stop}\}} F_\theta(s, r') + \sum_{r' \in \mathcal{R}_2} \sum_{b' \in \mathcal{B}_{r'}} F_\theta(s, (r', b'))}\]

위 수식에서 분모는 모든 가능한 action들에 대한 edge flow의 합을 의미한다. 이 구조에서는 일부 reactant block이 제거되어도 동일한 reaction template을 공유하는 다른 action에는 영향을 주지 않으므로, modified library에서도 consistent한 forward policy 추정이 가능하다고 설명하고 있다.

Figure 8. Illustration of a situation where an additional objective is introduced: excluding building blocks (reactants) containing the functional group −NR3. The Gray dashed lines mean masked actions. (a) GFlowNet on hierarchical MDP. (b) GFlowNet on non-hierarchical MDP.

Training Objective

학습 objective로는 trajectory balance (TB)를 사용한다고 설명하고 있다.

\[\hat{\mathcal{L}}_{\text{TB}}(\tau) = \left( \log \frac{Z_\theta \prod_{t=1}^{n} \hat{P}_F(s_t \mid s_{t-1}; \mathcal{A}_{t-1}^*; \theta)}{R(x) \prod_{t=1}^{n} P_B(s_{t-1} \mid s_t)} \right)^2\]

위 수식에서 $Z_\theta$는 partition function을 의미하고, $\hat{P}_F$는 action space subsampling으로 추정한 forward policy를, $P_B$는 fixed backward policy를 의미한다. $R(x)$는 terminal state $x$의 reward를 나타낸다.

Online training 시 sampling policy $\pi_\theta$는 estimated forward policy에 비례한 형태로 정의한다.

\[\pi_\theta(s' \mid s; \mathcal{A}^*) = \frac{w_{(s \to s')} F_\theta(s \to s')}{\sum_{(s \to s'') \in \mathcal{A}^*(s)} w_{(s \to s'')} F_\theta(s \to s'')}\]

Training Algorithm

Algorithm 1. Training GFlowNets with action space subsampling.
  • Line 1에서는 전체 action space $\mathcal{A}$와 maximum trajectory length $N$을 입력으로 받는다고 설명한다.
  • Line 3에서는 각 transition step별로 subsampling policy $\mathcal{P}(\mathcal{A})$를 통해 partial action space $\mathcal{A}_t^{\ast}$를 sampling한다고 이야기한다.
  • Line 4에서는 partial action space로 정의된 sampling policy $\pi_\theta$를 따라 trajectory $\tau$를 sampling하는 단계라고 볼 수 있다.
  • Line 5에서는 trajectory balance loss $\hat{\mathcal{L}}_{\text{TB}}$의 gradient로 model parameter를 업데이트한다고 설명한다.
  • Line 6에서는 model이 수렴할 때까지 위 과정을 반복한다고 이야기한다.

Experiments

Setup

논문에서는 Cretu et al. (2024)의 reaction template set (uni-molecular 13개, bi-molecular 58개)을 사용하고, Enamine comprehensive catalog의 1.2M building block을 사용한다고 설명한다. Synthetic pathway의 maximum reaction step은 Enamine REAL Space의 관행에 따라 3으로 설정하고, subsampling ratio는 1%로 둔다고 이야기한다.

Synthesizability 평가에는 retrosynthesis tool인 AiZynthFinder를 사용하며, USPTO reaction과 Enamine building block으로 합성 경로를 찾을 수 있을 때만 synthesizable로 인정한다고 설명한다.

Pocket-specific Optimization

LIT-PCBA dataset의 15개 protein 각각에 대해 64,000개의 분자를 생성하고, QED > 0.5 제약을 만족하는 분자 중 docking score 기준 상위 100개를 Tanimoto distance 0.5로 다양성을 유지하며 선정한다고 설명한다. Reward는 $R(x) = w_1 \text{QED}(x) + w_2 \widetilde{\text{Vina}}(x)$로 정의하고, 여기서 $\widetilde{\text{Vina}}$는 normalized docking score를 의미한다.

Table 1. Hit ratio. Mean and standard deviation over 4 runs. The best results are in bold.
Table 2. Vina. Mean and standard deviation over 4 runs. The best results are in bold.
Table 3. Synthesizability. Mean and standard deviation over 4 runs. The best results are in bold.
Table 4. Synthetic complexity. Mean and standard deviation over 4 runs. The best results are in bold.

논문에서는 RxnFlow가 모든 target에서 hit ratio, Vina score, synthesizability, synthetic complexity 측면에서 baseline을 능가한다고 보고한다. 특히 SynFlowNet과 RGFN처럼 동일한 reaction을 사용하는 GFlowNet baseline보다도 synthesizability가 높았으며, 이는 더 적은 reaction step과 더 큰 building block을 활용한 결과라고 해석한다.

또한 FragGFN+SA처럼 단순한 SA score를 reward에 포함시키는 방식은 synthesizability를 의미 있게 향상시키지 못한다고 보고한다. 저자들은 이 결과를 cheap한 synthesizability estimator가 suboptimal이라는 근거로 제시한다.

흥미롭게도 synthesizability를 고려하지 않는 FragGFN보다 RxnFlow의 docking score가 더 좋았는데, 저자들은 (1) Enamine library가 drug-like chemical space에 한정되어 optimization 복잡도가 줄어들고, (2) RxnFlow가 큰 building block을 사용하므로 trajectory가 짧아져 trajectory balance gradient variance가 감소한다는 두 가지 요인을 그 원인으로 설명하고 있다.

Pocket-conditional Generation

CrossDocked2020 benchmark에서 100개의 test pocket 각각에 대해 100개의 분자를 zero-shot 생성하고 Vina, QED, synthesizability, diversity, generation success, sampling time을 비교한다고 설명한다.

Table 5. CrossDocked2020 benchmark. We report the average and median values over the average properties for each test pocket. The best results are in bold and the second ones are in underlined. We denote the Generation Success as Succ., Synthesizability as Synthesiz., and Diversity as Div. Reference means the known binding active molecules in the test set. MolCRAFT-large (Qu et al., 2024) is the result when generating with more atoms than the reference ligands.

논문에서는 RxnFlow가 average Vina score -8.85 kcal/mol, synthesizability 34.8%로 state-of-the-art를 달성했다고 보고한다. 특히 SA score를 reward로 사용한 fragment-based TacoGFN (1.3%)이나 synthesizable drug set에서 학습된 distribution-based model들보다 훨씬 높은 synthesizability를 보였다고 강조한다.

또한 structural diversity가 0.81로 높게 유지되어, optimization power와 diversity 사이의 trade-off를 효과적으로 완화한다고 주장한다. 이는 chemically diverse한 building block 집합을 사용하기 때문이라고 해석한다.

Figure 3. Visualization of generated molecules in a zero-shot manner. (a-b) Docking results of generated molecules and known reference ligands of TBK1 (PDB Id: 1FV, SU6). (c) Generative trajectory, which is the generated synthetic pathway of the left molecule in (a).

Introducing Additional Objective without Retraining

논문에서는 non-hierarchical MDP 덕분에 재학습 없이 추가 objective를 도입할 수 있음을 보여준다. KRAS-G12C target에 대해 사전 학습된 model에서, building block을 TPSA 하위 15%로 제한하여 hydrophobic molecule을 생성하는 시나리오를 실험하였다.

Figure 4. Property distribution of sampled molecules with "all" building blocks and "low"-TPSA building blocks. Vina score was calculated against the KRAS-G12C target.

저자들은 TPSA가 낮은 building block만 사용해도 reward distribution이 크게 훼손되지 않으며, hydrophobicity를 효과적으로 증가시킬 수 있다고 보고한다. QED 분포에는 약간의 변화가 있었는데, 이는 TPSA와 QED 사이의 상관관계 때문이라고 해석한다.

Scaling Action Space without Retraining

신약 개발용 building block library는 2020년 약 6만 개에서 현재 120만 개 이상으로 빠르게 확장되고 있다고 이야기한다. 기존 model은 새로운 block을 수용하기 위해 재학습이 필요하지만, RxnFlow는 action embedding을 통해 unseen block까지 일반화할 수 있다고 주장한다.

전체 1M block library를 학습용 500K (seen)과 학습에 사용되지 않은 500K (unseen)으로 나눈 뒤, seen으로만 학습한 model이 unseen 및 all에 대해서도 유사한 reward distribution을 생성함을 보인다.

Figure 5. QED reward distribution of generated molecules for each of the "seen", "unseen", and "all" blocks. Additional results are in Figure 14.

Ablation Study: Effect of Building Block Library Size

Kappa-opioid receptor (PDB Id: 6b73)를 대상으로 building block library size에 따른 optimization power, diversity, generation time을 분석한다고 설명한다.

Figure 6. Optimization power, diversity, and generation time according to building block library size. (a-c) Average and standard deviation of properties of the top-1000 high-affinity molecules over 4 runs on pocket-specific generation. (a) Average docking score. (b) The uniqueness of Bemis-Murcko scaffolds. (c) Average Tanimoto distance. (d) Average runtime to generate 100 molecules over the CrossDocked2020 test pockets in a zero-shot manner.

논문에서는 library size가 커질수록 docking score, Bemis-Murcko scaffold uniqueness, Tanimoto diversity가 모두 향상된다고 보고한다. 동시에 100배 더 큰 action space에서도 generation time은 약 2배만 증가하여 효율성을 입증하였다고 주장한다.

Towards Further Development

저자들은 향후 개선 방향으로 3D interaction modeling을 제시한다. Beta-2 adrenergic receptor에 대해 docking conformation을 활용한 3D interaction modeling을 적용했을 때, 2D만 사용한 경우보다 candidate molecule 발견 성능이 향상됨을 관찰하였다고 보고한다.

Figure 7. Affect of 3D interaction modeling. Mean and standard error over 3 seeds.

또한 현재의 action space subsampling은 uniform sampling으로 인해 exploration이 강제되는 측면이 있으므로, building block prioritization을 도입해 exploitation을 강화할 수 있다고 이야기한다.

Limitation

논문은 RxnFlow의 성능을 다양한 task에서 강하게 입증하였으나, 다음과 같은 측면에서 추가 검증이 부족하다고 볼 수 있다.

  • Synthesizability 평가가 AiZynthFinder의 retrosynthetic analysis에 의존하므로, 실제 wet-lab synthesis success rate와의 일치 여부는 별도로 검증되지 않았다.
  • Action space subsampling이 unbiased estimator임을 이론적으로 보였지만, large-scale benchmark에서 subsampling ratio가 bias-variance trade-off에 미치는 영향에 대한 충분한 ablation은 제한적이다.
  • 3D interaction modeling은 한 개의 target에 대해 maximum reaction step 1의 단순화된 setting에서만 평가되어, 일반적인 multi-step generation에서의 효과는 추가 연구가 필요하다.

Conclusion

논문에서는 RxnFlow를 합성 경로 위에서 generative flow를 학습하는 synthesis-oriented framework로 제안하고 있다. Action space subsampling을 통해 1.2M 규모의 building block과 71개의 reaction template을 효율적으로 다루며, non-hierarchical MDP와 action embedding을 결합해 재학습 없이도 modified building block library와 추가 objective에 적응할 수 있다고 주장한다.

실험 결과 RxnFlow는 LIT-PCBA pocket-specific optimization과 CrossDocked2020 pocket-conditional generation 모두에서 state-of-the-art 성능을 보였으며, synthesizability와 diversity를 동시에 확보할 수 있는 실용적인 방법임을 보여준다고 해석한다.

Reference

  • Seo, Seonghwan, et al. “Generative flows on synthetic pathway for drug design.” International Conference on Learning Representations. Vol. 2025. 2025.