LeakGFN: Robust Molecular Generation in Generative Flow Networks via Flow Decomposition

Motivation

• Known: GFlowNets는 reward function에 비례하여 다양한 후보를 샘플링하는 강력한 프레임워크로, flow matching, trajectory balance, detailed balance, subtrajectory balance 등 다양한 training objective와 exploration 전략(ϵ-greedy, Thompson sampling, local search, genetic operator)이 연구되어 왔으며 pocket-conditioned 및 fragment-based 분자 생성으로 확장되어 왔다.
• Gap: 방대한 chemical space에서 계산 효율성을 위해 trajectory length를 Lmax로 truncation할 수밖에 없는데, 기존 방법들은 truncation boundary에 도달한 미완성 fragment(forced terminal)를 valid molecule(valid terminal)과 동일하게 terminal state로 취급하여 학습된 분포가 화학적으로 무의미한 상태에 확률 질량을 할당하는 target space mismatch 문제를 초래하며, 기존 연구들은 training objective, exploration, active learning 개선에 집중했을 뿐 이 근본적 문제를 직접 다루지 않았다.
• Why: Truncation으로 인한 flow leakage 문제는 특히 복잡한 약물 scaffold와 같이 큰 분자를 다룰 때 더욱 심각해지며, 정확한 target 분포를 학습하지 못하면 유효한 chemical space를 탐색할 유효 용량이 줄어들어 약물 발견을 위한 다양한 고보상 후보 탐색이라는 GFlowNet의 핵심 강점이 저해되기 때문에 이를 이론적으로 해결하는 것이 중요하다.
• Approach: Flow를 chemical head(전체 chemical space에 대한 flow 모델링)와 valid head(truncation boundary 내에서 valid molecule에 도달하는 flow 비율 추정)로 분해하는 dual-head 아키텍처(Fvalid = Fchem · λ)를 제안하여, valid head가 명시적 supervision 없이 molecular reachability를 암묵적으로 학습하도록 한다.

Achievement

Figure 2. Evolution of learned valid fraction λ during training on JNK3 (Lmax = 8). The model progressively learns to as

1. Flow Leakage의 최초 정형화: Trajectory truncation이 GFlowNet으로 하여금 valid molecule과 incomplete fragment가 conflate된 잘못된 target space를 학습하게 만드는 현상을 Definition 3.1로 최초로 공식 정의하였다.
2. 이론적 보장을 갖는 dual-head 해법 제안: LEAKGFN이 mild assumption 하에서 accessible molecule에 대한 정확한 분포를 복원함을 증명하였고, chemical/valid flow의 gradient 신호 분리를 통해 truncation boundary 근처의 training dynamics를 개선함을 보였다.
3. 광범위한 실증 검증: 5개 분자 최적화 벤치마크 중 4개에서 state-of-the-art 성능을 달성하고 hyperparameter(Lmax) 선택에 강건함을 확인했으며, ablation을 통해 성능 향상이 exploration strategy가 아닌 flow decomposition 자체에 기인함을 입증했다.
4. Plug-and-play 모듈로서의 범용성 입증: Pocket-conditioned 생성과 multi-objective 생성 task에 LEAKGFN 모듈을 통합하여 기존 프레임워크의 성능을 향상시켰다.

How

Figure 3. Top-6 molecules generated by LEAKGFN for JNK3 kinase inhibition. Each molecule is shown with its predicted act

• GFlowNet의 flow conservation 원리와 forced terminal(Tforced) vs valid terminal(Tvalid)의 구분을 formal하게 정의하고 flow leakage(Lleak)를 수식으로 정량화함(Definition 3.1).
• 각 state에서의 flow를 chemical flow(Fchem, 전체 chemical space에 대한 flow)와 valid flow(Fvalid, valid molecule로 이어지는 flow)로 분해하고, Fvalid = Fchem · λ 관계를 통해 valid head가 reachability 비율 λ를 예측하도록 함.
• Chemical head와 valid head가 별도의 gradient pathway를 가지도록 하여 flow magnitude 학습과 reachability 추정 학습 간의 충돌을 해소함.
• Mild assumption 하에서 이 decomposition이 valid, reachable molecule에 대한 정확한 distribution을 복원함을 이론적으로 증명함.
• JNK3, GSK3β, DRD2 등을 포함한 5개 분자 최적화 task와 pocket-conditioned/multi-objective 생성 task에서 flow matching baseline과 비교 실험 및 ablation study(exploration strategy와 분리하여 decomposition 자체의 기여도 검증)를 수행함.

Originality

• Trajectory truncation으로 인한 flow leakage 현상을 최초로 formal하게 정의하고 이론적으로 분석한 점이 독창적이다.
• 기존 연구들이 training objective, exploration, active learning 등 간접적 개선에 집중한 반면, LeakGFN은 truncation이 야기하는 incorrect target space 문제를 flow decomposition이라는 직접적이고 원리적인 방식으로 해결한다.
• Valid head가 explicit supervision 없이 molecular reachability를 암묵적으로 학습하도록 하는 asymmetric terminal treatment 설계가 독창적이며, 이는 기존 dual-head 또는 auxiliary network 설계와 차별화된다.
• 기존 GFlowNet training objective(FM, TB, DB, subTB)와 orthogonal한 plug-and-play 모듈로 설계되어 범용성을 갖춘다.

Limitation & Further Study

• 논문 발췌에서는 이론적 증명의 "mild assumption"이 실제 분자 생성 환경에서 얼마나 현실적인지에 대한 구체적 논의가 제한적으로 보이며, 이 가정이 위배되는 경우의 성능 저하 가능성에 대한 분석이 필요하다.
• 5개 벤치마크 중 1개 task에서는 state-of-the-art를 달성하지 못했는데, 이 실패 사례에 대한 원인 분석이 충분히 제시되지 않은 것으로 보인다.
• Dual-head 구조 도입으로 인한 추가 계산 비용 및 파라미터 증가에 대한 정량적 분석이 부족할 수 있다.
• 후속 연구로 더 다양한 fragment-based/atom-level 생성 방식, 더 큰 Lmax 범위, 그리고 다른 도메인(예: 재료 설계)으로의 확장 가능성을 검토할 필요가 있다.

Evaluation

총평: GFlowNet 기반 분자 생성에서 그동안 간과되어 온 truncation 유발 flow leakage 문제를 이론적으로 정형화하고 이를 해결하는 간결하면서도 원리적인 dual-head 구조를 제안한 점에서 완성도 높은 연구이며, plug-and-play 특성 덕분에 실용적 파급력도 기대된다.