베이지안 플로우로 OOD 화학공간을 샘플링하는 혁신적 접근

초록

본 논문은 베이지안 플로우 네트워크(BFN) 기반 ChemBFN 모델에 강화학습과 ODE‑유사 샘플링을 결합하고, 반자동 회귀(SAR) 전략을 도입함으로써 기존 확산 모델이 어려워하는 훈련 데이터 분포 밖(Out‑of‑Distribution, OOD) 고품질 분자를 효율적으로 생성할 수 있음을 입증한다. 실험 결과 MOSES, GuacaMol, 단백질 서열 데이터셋 등에서 유효성·다양성·신규성 모두에서 최첨단 모델을 능가한다.

상세 분석

이 연구는 베이지안 플로우 네트워크(BFN)의 본질적 특성을 OOD 샘플링에 활용한다는 점에서 의미가 크다. BFN은 확산 모델과 달리 명시적인 노이즈 스케줄을 정의하지 않고, 연속적인 파라미터 공간(θ) 대신 잠재 변수 공간(z)에서 역방향 SDE를 근사한다. 이러한 구조는 분포를 “정밀하게” 맞추는 것이 아니라, 파라미터를 더 정보‑풍부한 방향으로 최적화하게 하여 훈련 데이터와의 거리(예: Fréchet ChemNet Distance)를 자연스럽게 확대한다.

논문은 세 가지 핵심 개선을 제시한다. 첫째, 강화학습(RL) 보상 항을 KL 손실에 추가한다. 여기서는 “시간 t에서 출력 분포가 유효한 분자를 생성하는가”를 이진 기준(c)으로 삼아, η=0.01의 스케일링으로 보상을 부여한다. 이는 샘플링 단계가 적을 때도 유효 SMILES 비율을 크게 끌어올린다.

둘째, ODE‑like 샘플링 알고리즘을 도입한다. β(t) 정확도 스케줄을 이용해 잠재 변수 z를 단계별로 업데이트하고, 온도 파라미터 τ>0으로 노이즈를 조절한다. τ=0.5(또는 데이터에 따라 0.05)로 설정하면 유효성은 99% 이상이면서 다양성 손실을 최소화한다.

셋째, 반자동 회귀(SAR) 전략을 적용한다. 기존 ChemBFN은 양방향 토큰 업데이트를 사용했으나, 주의 행렬에서 대각선에서 멀어질수록 값이 거의 0임을 관찰하고, 인과적 마스크(causal mask)를 삽입해 토큰을 블록 단위로 동시에 업데이트하면서도 미래 토큰을 사용하지 않는다. SAR은 훈련과 추론 모두에 적용 가능하며, 네 가지 조합(전통/normal × SAR 적용 여부) 중 전략 4(SAR + SAR)와 RL·ODE 결합이 가장 높은 OOD 성능을 보였다.

실험에서는 MOSES와 GuacaMol 벤치마크에서 샘플링 스텝을 1 000→10으로 감소시켜도 유효성·신규성·다양성이 크게 저하되지 않았다. 특히 ChemBFN+RL+ODE는 MOSES에서 유효성 0.999, 독창성 0.998, FCD 0.797을 달성했으며, GuacaMol에서도 유효성 0.863, 독창성 0.980을 기록했다. 단백질 서열 생성 실험에서는 베타‑시트 비율과 SASA를 목표로 삼아, 기존 모델 대비 OOD 다중 목표 최적화에서 우수한 결과를 보여준다.

이러한 결과는 베이지안 플로우가 “분포 학습”을 넘어 “분포 탐색” 도구로 활용될 수 있음을 증명한다. 강화학습 보상과 온도 조절 ODE 샘플링, 그리고 SAR이라는 구조적 변형이 결합돼, 적은 샘플링 비용으로도 훈련 데이터 밖의 고품질 화합물을 효율적으로 탐색한다는 점에서 신약 설계와 재료 발견에 실용적 가치를 제공한다.