Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

기술 요약: 분자 설계를 위한 순차적 미세 조정을 통한 제약 조건이 있는 흐름 최적화

1. 문제 정의

본 논문은 흐름(flow) 및 확산(diffusion) 모델과 같은 생성형 파운데이션 모델을 분자 설계 및 단백질 공학 같은 실제 과학적 발견 응용 분야에 적응시키는 과정에서의 과제를 다룹니다. 이러한 모델들은 복잡한 데이터 분포를 포착하는 데는 뛰어나지만, 특정 보상(예: 결합 친화도, 쌍극자 모멘트)을 최대화하는 동시에 도메인 특화 제약 조건(예: 합성 가능성, 독성 경계 또는 에너지 안정성)을 엄격히 준수해야 하는 발견 작업에서의 표준 샘플링 방식은 불충분합니다.

RL(강화 학습) 또는 제어 이론에 기반한 기존의 미세 조정 방법들은 분포 드리프트(distributional drift)를 방지하기 위해 KL 정규화를 적용하여 보상 함수를 최적화하지만, 하드 제약 조건(hard constraints)을 신뢰성 있게 처리하는 데 어려움을 겪습니다. 보상을 목적 함수 내에 가중치 항으로 포함시키는 단순한 접근 방식은 보상과 제약 사이의 최적 가중치가 작업 및 훈련 단계마다 달라지기 때문에 비효령적인 시행착오 튜닝을 요구하므로 신뢰할 수 없다고 간주됩니다. 또한, 이러한 방법들은 생성된 샘로가 명시적인 제약 조건을 반드시 만족하도록 보장하지 못하며, 종종 높은 보상을 얻더라도 중요한 물리적 또는 화학적 특성을 위반하는 샘플을 생성합니다.

핵심 질문은 다음과 같습니다: 사전 훈련된 흐름 또는 확산 모델을 어떻게 신뢰할 수 있고 예측 가능하게 보상 최적화와 제약 조건 만족 사이의 트레이드오프를 수행하도록 미세 조정할 수 있는가?

2. 방법론

저자들은 이 문제를 해결하기 위한 공식적인 프레임워크와 알고리즘 솔루션을 제안합니다.

제약 조건이 있는 생성 최적화 프레임워크:
이 문제는 사전 훈련된 모델 $\pi_{pre}$에 대한 KL 정규화를 유지하면서, 기대 제약 조건 $c(x) \leq B$를 만족하면서 기대 보상 $r(x)$를 최대화하는 정책 $\pi^*$를 찾는 것으로 공식화됩니다. 목적 함수는 다음과 같습니다:
$$\max_{\pi} \mathbb{E}_{x \sim p^{\pi}_1}[r(x)] - \alpha D_{KL}(p^{\pi}_1 || p^{pre}_1) \quad \text{s.t.} \quad \mathbb{E}_{x \sim p^{\pi}_1}[c(x)] \leq B$$
이 공식화는 보상 유도 최적화(reward-guided optimization)와 제약 조건이 있는 생성(보상이 상수인 경우)을 모두 일반화합니다.

제약 조건이 있는 흐름 최적화 (CFO):
이를 해결하기 위해 저자들은 증강 라그랑주(Augmented Lagrangian, AL) 기법에 기반한 알고리즘인 **제약 조건이 있는 흐름 최적화(Constrained Flow Optimization, CFO)**를 도입합니다. CFO는 제약 조건이 있는 문제를 일련의 무제한 엔트로피 정규화 미세 조정 하위 문제로 변환합니다. 알고리즘은 다음 단계와 함께 반복적으로 작동합니다:

1. 증강 목적 함수 구축: 각 반복 단계 $k$에서, 증강 목적 함수 $f_k(x)$가 구축됩니다. 이 함수는 보상 $r(x)$와 라그랑주 승수 $\lambda_k$ 및 패널티 파라미터 $\rho_k$에 의해 조정된 제약 위반 유래 패널티 항을 결래합니다:
   $$f_k(x) = r(x) - \frac{\rho_k}{2} \left[ \max\left(0, c(x) - B - \frac{\lambda_k}{\rho_k}\right) \right]^2$$
2. 미세 조정 단계: 사전 훈련된 모델에 대한 KL 발산을 정규화 요소로 하여, 증강 목적 함수 $f_k$에 대해 정책 $\pi_k$를 최적화하기 위해 표준 미세 조정 솔버(예: Adjoint Matching, DiffusionNFT)가 사용됩니다.
3. 쌍 변수 업데이트:
   • 라그랑주 승수 ($\lambda$): 경험적 제약 갭 $G_k = \mathbb{E}[c(x)] - B$를 기반으로 업데이트됩니다. 제약 조건이 위반되면($G_k > 0$), 패널티를 강화하기 위해 $\lambda$를 감소시키고, 만족되면 완화합니다.
   • 패널티 파라미터 ($\rho$): 수축 통계량 $V_k$를 기반으로 업데이트됩니다. 제약 갭이 충분히 수축되지 않으면, 더 엄격한 실행 가능성을 강제하기 위해 $\rho$를 증가시킵니다.

이러한 순차적 접근 방식은 수동적인 가중치 튜닝 없이도 알고리즘이 보상 최대화와 제약 조건 만족 사이의 균형을 자동으로 맞출 수 있게 합니다.

3. 주요 기여

본 논문은 네 가지 주요 기여를 합니다:

• 공식 프레임워크: 임의의 제약 조건 하에서 보상 유도 적응이라는 실질적인 작업을 포착하여, 흐름 미세 조정을 통한 제약 조건이 있는 생성 최적화를 공식화했습니다.
• CFO 알고리즘: 순차적 미세 조정을 통해 제약 문제를 증명 가능하게 해결하는 증강 라그랑주 기반 방법인 CFO를 도입했습니다.
• 이론적 보장: 근사 솔버에 대한 현실적인 가정과 특정 오라클 조건 하에서, 저자들은 수렴 보장을 제공합니다. 그들은 CFO가 불가능성(infeasibility)을 최소화하는 정책을 반환하며, 더 강력한 조건(근사 오차 소멸) 하에서는 전역 최적성을 달성함을 증명합니다.
• 실증적 검증: 제안된 방법은 저차원 예시 설정과 고차원 분자 설계 작업 모두에서 평가되었으며, 보상과 제약 사이의 트레이드오프를 효과적으로 수행하는 능력을 입증했습니다.

4. 실험 결과

저자들은 두 가지 주요 측면에서 CFO를 평가했습니다.

저차원 예시 설정:

• 제약 조건이 있는 최적화: 가우시안 혼합 모델을 사용하여, CFO는 유효한 제약 영역(삼각형) 내에 머물면서 모델을 고보상 영역(타겟 지점 근처)으로 성공적으로 유도했습니다. 제약 조건을 무시하는 표준 Adjoint Matching과 비교했을 때, CFO는 높은 보상을 유지하면서도 제약 위반을 크게 감소시켰습니다 (2.47에서 0.12로).
• 제약 조건이 있는 생성: 보상을 상수로 설정함으로써, CFO는 제약 조건이 있는 생성 능력을 회복했습니다. 제약 경계 $B$를 변화시킴으로써 제약의 완화를 제어할 수 있음을 보여주었으며, 이는 방법의 유연성을 입증합니다.
• 솔버 불가지론(Solver Agnosticism): CFO는 1차 솔버(Adjoint Matching)와 그래디언트 프리 솔버(DiffusionNFT) 모두와 결합되었을 때 일관되게 성능을 발휘하여 모듈성을 확인했습니다.

분자 설계 작업:

• 에너지 제약 쌍극자 극대화: GEOM Drugs 데이터셋에 대해 FlowMol 모델을 사용하여, CFO는 총 에너지(제약, 안정성의 대리 지표)에 대한 상한선을 준수하면서 쌍극자 모멘트(보상)를 최적화했습니다.
  • CFO는 쌍극자 모멘트를 사전 훈련된 모델의 6.55 D에서 8.39 D로 증가시키는 동시에 에너지 제약(기대 에너지 -82.28 Ha vs. 경계 -80 Ha)을 충족했습니다.
  • 반면, 제약이 없는 미세 조정(Adjoint Matching)은 유사한 보상을 달성했지만 에너지 제약을 위반했습니다 (기대 에너지 -78.31 Ha).
• 베이스라인과의 비교: CFO는 수동으로 튜닝된 고정 가중치 패널티 베이스라인보다 우수한 성능을 보였습니다. 수동 튜닝은 CFO의 성능과 일치하는 단일 구성을 찾기 위해 18개 차수의 가중치 범위를 탐색해야 했으며, 이는 수동 튜닝의 비효율성을 강조합니다.
• 유효성 및 효율성: 미세 조정이 화학적 유효 분자의 비율을 감소시키는 것(생성 설계의 알려진 과제)은 사실이지만, CFO는 유사한 보상 수준에서 제약이 없는 베이스라인보다 더 높은 유효성을 유지했습니다. 또한 CFO는 추론 시 제약 조건 방법인 DiffOpt의 높은 샘플링 비용과 비교하여 단 한 번의 미세 조정 패스만 필요하므로 계산 효율성을 입증했습니다.

5. 의의 및 주장

본 논문은 CFO가 제약 조건이 있는 생성 최적화를 위한 신뢰할 수 있고 예측 가능한 메커니즘을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 자동화된 트레이드오프: 현재의 생성 설계 워크플로에서 주요 병목 현상인 제약 가중치의 수동적이고 비효율적인 시행착오 튜닝을 제거합니다.
• 이론적 엄밀성: 제약 조건 만족과 보상 최적화 모두에 대한 수렴 보장을 제공하여, 방법을 확립된 최적화 이론(증강 라그랑주)에 근거하게 합니다.
• 실용적 유용성: 고차원 분자 설계 작업에 성공적으로 적용함으로써, 저자들은 제약 조건(화학적 안정성과 같은)이 타협 불가능한 실제 과학적 문제에 이 방법이 확장 가능하다는 것을 보여주었습니다.
• 모듈성: 프레임워크는 기초가 되는 미세 조정 솔버에 구애받지 않으므로 다양한 최적화 전략(그래디언트 기반 또는 그래디언트 프리)과 결합될 수 있으며, 다양한 데이터 모달리티(연속형, 이산형 또는 혼합형)에 적용 가능합니다.

저자들은 결론적으로 CFO가 사전 훈련된 흐름 모델을 임의의 제약 조건을 준수하면서 기대 보상을 최대화하도록 적응시키는 것을 가능하게 하며, 과학적 발견 응용 분야를 위한 강력한 도구를 제공한다고 설명합니다.