Constrained Flow Optimization via Sequential Fine-Tuning for Molecular Design

Motivation

• Known: 최근 연구들은 KL-regularization을 활용한 reinforcement learning 또는 control theoretic 기법을 통해 사전학습된 diffusion/flow model을 특정 reward function에 맞게 파인튜닝하는 scalable한 방법들을 제시해왔다(예: Domingo-Enrich et al., 2025; Uehara et al., 2024b; Tang & Zhou, 2024).
• Gap: 기존 방법들은 사전학습 모델로부터의 KL 정규화를 통해 분포 이동을 제한할 뿐, synthesizability나 toxicity 같은 domain-specific한 hard constraint의 만족을 보장(certify)할 수 없으며, constraint를 reward에 수동으로 가중합해 넣는 방식은 task와 학습 단계마다 적절한 가중치가 달라 비효율적인 trial-and-error가 필요하고 reward 극대화 과정에서 예기치 않게 constraint 위반을 유발할 수 있다는 근본적 한계가 있다.
• Why: binding affinity 최적화와 molecular synthesizability 같은 constraint 만족을 동시에 신뢰성 있게 달성하는 것은 약물 설계 및 단백질 엔지니어링 등 실제 과학적 발견 응용에서 생성 foundation model을 채택하기 위한 핵심 전제조건이기 때문이다.
• Approach: Constrained Generative Optimization이라는 최적화 관점의 formal framework를 도입하고, augmented Lagrangian scheme에 기반해 제약이 있는 목적함수를 일련의 ordinary generative optimization 하위 문제(sequential fine-tuning)로 환원하는 CFO 알고리즘을 제안하며, 이에 대한 수렴 보장을 제공한다.

Achievement

Figure 3. Energy-constrained dipole moment maximization of FlowMol (Dunn & Koes, 2024) on GEOM Drugs (Axelrod & Gomez-

1. Constrained Generative Optimization 프레임워크 정식화: reward function이 상수인 특수 사례로 constrained generation을 포함하는, flow model fine-tuning을 통한 constrained optimization 문제를 최적화 관점에서 엄밀하게 정의하였다.
2. CFO 알고리즘 제안: augmented Lagrangian 기반 dual 접근법으로, KL-regularized fine-tuning(augmented reward 최대화)과 생성 샘플의 constraint violation 추정치를 이용한 파라미터 업데이트를 번갈아 수행하여 reward maximization과 constraint satisfaction 간 trade-off를 자동으로 조정한다.
3. 이론적 수렴 보장: approximate solver라는 현실적 가정 하에서 constraint satisfaction을, 더 이상화된 세팅에서 reward maximization을 보장하는 augmented Lagrangian 이론 기반 증명을 제시하였다.
4. 실험적 검증: 시각적으로 해석 가능한 synthetic setting과 quantum mechanical property의 constrained optimization을 포함한 고차원 molecular design task(FlowMol 등)에서 CFO가 일관되게 reward를 증가시키면서 높은 constraint satisfaction을 유지함을 입증하였다.

How

Figure 2. Top Constrained Generative Optimization: Samples from the pre-trained policy (2a) and policies fine-tuned with

• Flow model과 Flow Matching, continuous-time RL을 배경 이론으로 사용하여 constrained generative optimization 문제를 formal하게 정의함.
• Augmented Lagrangian scheme(Birgin & Martínez, 2014)을 도입하여 원 문제를 augmented reward function을 최대화하는 KL-regularized fine-tuning 하위 문제들의 시퀀스로 환원함.
• 각 iteration에서 (1) 현재 augmented reward를 최대화하는 fine-tuning 단계와 (2) 생성된 샘플들에서 추정된 constraint violation을 바탕으로 augmented reward의 penalty 파라미터(듀얼 변수)를 업데이트하는 단계를 반복(alternate)함.
• Approximate solver oracle 가정 하에서 constraint satisfaction 보장을, 보다 이상화된 oracle 가정 하에서 reward maximization 보장을 도출함.
• 시각적으로 해석 가능한 synthetic 실험(2D density 이동 등)과 FlowMol 기반 molecular design(에너지 제약 하 dipole moment 최대화, drug-like molecule 생성) 실험을 통해 CFO와 fixed-µ baseline들을 Pareto 비교함.

Originality

• Reward-guided fine-tuning과 constraint satisfaction을 별도의 heuristic한 가중합이 아니라, augmented Lagrangian 기반의 sequential fine-tuning으로 환원하여 자동으로 균형을 맞추는 최초의 접근으로 보임.
• Constrained generation(제약된 생성모델링)을 reward가 상수인 특수 사례로 포섭하는 통합적 Constrained Generative Optimization 프레임워크를 제시하여, 기존의 개별적인 constrained generation 연구와 reward-guided fine-tuning 연구를 하나의 이론적 틀로 연결함.
• Mirror descent, chance constraint, bilevel optimization 등 고전적 최적화 아이디어를 flow/diffusion model fine-tuning에 접목하는 최근 흐름 속에서, augmented Lagrangian이라는 검증된 제약최적화 도구를 생성모델 파인튜닝에 처음으로 체계적으로 적용함.

Limitation & Further Study

• 이론적 수렴 보장이 approximate solver 또는 idealized oracle이라는 가정에 의존하고 있어, 실제 유한 샘플·유한 스텝 학습에서 이러한 가정이 얼마나 근사적으로 성립하는지에 대한 추가 검증이 필요함.
• 평가가 synthetic setting과 상대적으로 제한된 molecular design task(quantum mechanical property, drug-likeness 등)에 국한되어 있어, 단백질 엔지니어링 등 더 크고 복잡한 실제 과학적 discovery task로의 확장성 및 scalability 검증이 후속 연구로 필요함.
• Constraint violation 추정을 위한 oracle(또는 classifier)의 정확도에 성능이 크게 좌우될 수 있는데, 이러한 추정 오차가 CFO의 이론적 보장에 미치는 영향에 대한 정량적 분석이 부족함.
• Augmented Lagrangian iteration의 계산 비용(반복적인 fine-tuning 단계)이 단일 파인튜닝 대비 얼마나 증가하는지에 대한 효율성 논의가 제한적임.

Evaluation

총평: 고전적 제약 최적화 이론(augmented Lagrangian)을 생성 foundation model의 fine-tuning 문제에 엄밀하고 실용적으로 접목한 견고한 연구로, reward-constraint trade-off를 수동 튜닝 없이 자동화했다는 점에서 분자 설계 등 과학적 발견 응용에 실질적 기여를 할 것으로 기대된다.