Active Flow Expansion for Out-of-Distribution Discovery: from Theory to Molecules

Motivation

• Known: flow matching과 diffusion model은 관측 데이터 분포 pdata를 근사하도록 학습되며(distribution matching), 이는 화학, 생물학, 로보틱스 등 다양한 도메인에서 고품질 샘플 생성에 성공적으로 활용되어 왔다. 최근에는 test-time reward-tilting 기법들이 등장하여 분자 특성 최적화나 효소 설계 등 과학적 발견에 생성 모델을 활용하는 흐름이 확대되고 있다.
• Gap: 기존 pre-training은 distribution matching을 목표로 하므로 관측 데이터가 커버하지 못하는 valid design space의 넓은 영역, 즉 new-to-nature designs를 근본적으로 생성할 수 없다. out-of-distribution discovery의 목표는 pdata에서 확률이 거의 0인 영역을 탐색하는 것이므로 표준 생성 모델링 패러다임과 근본적으로 상충되며, 이를 task-agnostic하게 해결하는 방법이 부재했다.
• Why: 과학적 발견(신약, 신규 단백질, 신규 분자 등)의 핵심은 자연이나 기존 데이터베이스에 존재하지 않는 새로운 valid design을 찾는 것인데, 현재 가장 강력한 생성 모델들이 구조적으로 이러한 발견에 취약하다는 근본 문제를 이론과 실증 양쪽에서 다룬다는 점에서 중요하다.
• Approach: 생성 모델을 데이터 분포 근사기가 아닌 valid design space 근사기로 재해석하여 generable set이라는 개념을 도입하고, verifier feedback과 active exploration을 이용해 pre-trained flow model의 generable set을 반복적으로 확장하는 continued pre-training 알고리즘 ACTFLOW를 제안한다.

Achievement

Figure 1. (left) Pre-trained flow model µθ0 generates with sufficient probability a set Ωτ

1. Generable set 개념 도입: τ-level generable set Ωτ_θ = {x : p^θ_1(x) ≥ τ}를 정의하여 모델이 non-negligible probability로 커버하는 영역을 수학적으로 포착하고, 표준 pre-training의 근본 한계(Ωτ_θpre ⊂ Ω, Vol(Ωτ_θpre) ≪ Vol(Ω))를 formal하게 규명함.
2. 새로운 학습 원리 제시: distribution matching(pθ_1 ≈ pdata) 대신 generable set expansion(Ωτ_θ ≈ Ω*)이라는 OOD flow modeling을 위한 학습 목표를 새롭게 정식화함.
3. ACTFLOW 알고리즘: verifier feedback과 학습된 diffusion representation의 중간 noise level에서의 active exploration을 활용하여 self-generated synthetic data로 반복적으로 모델을 적응시키는 continued pre-training 방법을 제안함.
4. 이론적 보장: energy-based modeling abstraction에 기반해 OOD flow modeling에 대한 최초의 statistical learning guarantee를 수립하고, generable set expansion을 학습된 표현 공간에서의 local-to-global reachability process로 분석함.
5. 폭넓은 실증 검증: QM9 소분자, GEOM-Drugs 약물유사분자, therapeutic peptides, protein sequence design 등 네 가지 도메인에서 domain-agnostic OOD 생성 모델링 지표(coverage, diversity)로 평가하여 REC-NF, REC-F 등 기존 recursive synthetic pre-training baseline 대비 유의미한 성능 향상을 입증함(Table 1).

How

Figure 1. (left) Pre-trained flow model µθ0 generates with sufficient probability a set Ωτ

• 설계 공간 X 위에서 flow model을 velocity field uθ와 ODE dxt/dt = uθ(xt,t)로 정의하고 base distribution p0에서 terminal density pθ_1로 흘러가는 구조를 사용
• validity indicator v(x)와 valid design space Ω* = {x : v(x)=1}을 정의하고, generable set Ωτ_θ를 통해 커버리지를 정량화
• 표준 flow matching pre-training(conditional interpolation paths, vector-field regression loss)으로부터 시작하여 pre-trained model θ0을 얻음
• black-box verifier(예: 분자 유효성 검사기, protein foldability 예측기)를 통해 self-generated synthetic data의 validity를 확인
• 학습된 diffusion representation의 intermediate noise level에서 active exploration을 수행하여 새로운 valid 영역을 탐색하고, 이를 반복적(iterative)으로 모델에 재학습시켜 θ0 → θ1 → ... → θT로 generable set을 점진적으로 확장(Ωτ_θ0 ⊆ Ωτ_θ1 ⊆ ... ⊆ Ωτ_θT ≈ Ω*)
• energy-based modeling 관점에서 이 과정을 local-to-global reachability process로 이론화하여 statistical learning guarantee 도출
• coverage(valid clusters 수)와 diversity(Vendi score) 지표로 QM9, GEOM-Drugs, peptides, protein sequences 도메인에서 REC-NF, REC-F 등 baseline과 비교 실험

Originality

• 생성 모델을 data distribution approximator가 아닌 valid design space approximator로 재개념화하는 관점 전환이 참신함
• generable set이라는 새로운 수학적 대상을 정의하여 OOD 생성 모델링을 이론적으로 다룰 수 있는 프레임워크를 최초로 제시
• flow matching의 표준 목적함수(distribution matching)를 대체하는 generable set expansion이라는 새로운 학습 원리를 도입
• active exploration을 diffusion representation의 intermediate noise level에서 수행한다는 아이디어로, 순수 데이터 공간이 아닌 학습된 표현 공간의 기하학적 특성을 활용한 탐색 전략을 제안
• OOD flow modeling에 대한 최초의 statistical learning guarantee를 제공한다고 주장하는 이론적 기여

Limitation & Further Study

• verifier가 black-box이며 정확하다는 가정에 의존하는데, 실제 도메인에서는 verifier 자체의 오류나 편향이 generable set expansion의 방향을 왜곡시킬 수 있음
• iterative recursive self-training 방식이므로 초기 반복에서 잘못된 방향으로 확장되면 이후 반복에서 오류가 누적(model collapse, distributional drift)될 위험이 있으나 이에 대한 심층 분석이 부족해 보임
• 실험이 워크숍 논문 형식으로 제한적이며(coverage, diversity 지표 중심), 실제 downstream task에서의 유효성(예: wet-lab validation)에 대한 검증은 부재
• 이론적 보장이 energy-based modeling abstraction 등 여러 단순화 가정에 기반하므로 실제 고차원 분자/단백질 설계 공간에 대한 이론과 실제 간 격차 존재
• coverage 지표(valid clusters 수)의 정의와 threshold τ 선택이 결과에 미치는 민감도 분석이 충분히 제시되지 않음
• 후속 연구로 verifier의 불확실성을 명시적으로 모델링하거나, 더 큰 규모의 실제 discovery task(예: 실험적 검증)로의 확장이 필요해 보임

Evaluation

총평: 생성 모델링의 근본 패러다임인 distribution matching에서 벗어나 valid design space coverage를 직접 목표로 하는 관점 전환은 매우 신선하며, 이를 뒷받침하는 이론과 다양한 도메인(분자, 펩타이드, 단백질)에서의 실증 결과가 설득력 있게 제시된 우수한 워크숍 논문이다.