Bayesian Flow Is All You Need to Sample Out-of-Distribution Chemical Spaces

이 논문은 반자기회귀 전략과 강화학습을 도입한 베이지안 플로우 네트워크 (ChemBFN) 가 훈련 데이터 분포를 벗어나는 고품질 분자 생성을 가능하게 하여 기존 최첨단 모델을 능가함을 보여줍니다.

핵심

이 논문은 **"새로운 약을 만들기 위해, AI 가 기존에 알려지지 않은 완전히 새로운 분자를 어떻게 찾아낼 수 있는가?"**라는 질문에 대한 답을 제시합니다.

기존의 AI 모델들은 주로 "배운 데이터와 비슷한 것"을 만드는 데 능숙했지만, "배운 것보다 훨씬 더 좋은 성질을 가진 새로운 것"을 찾는 데는 서툴렀습니다. 이 연구는 **베이지안 플로우 네트워크 (Bayesian Flow Network, BFN)**라는 새로운 기술을 활용해 이 문제를 해결했다고 말합니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 문제 상황: "유리창을 닦는 청소부" vs "새로운 세계를 탐험하는 모험가"

기존의 AI 모델 (확산 모델 등) 은 마치 유리창을 닦는 청소부와 같습니다.

• 방식: 이미 있는 유리창 (기존 분자 데이터) 을 최대한 깨끗하게, 원본과 똑같이 닦아내려고 노력합니다.
• 한계: 청소부는 유리창을 더 깨끗하게 만들 수는 있어도, 유리창을 부수고 새로운 창을 만들거나, 유리창이 아닌 '벽'을 만들 수는 없습니다. 즉, **배운 데이터 범위 밖 (Out-of-Distribution)**의 새로운 것을 창조하는 데는 한계가 있습니다.

반면, 이 논문에서 제안한 ChemBFN은 새로운 세계를 탐험하는 모험가입니다.

• 방식: 기존 지도 (데이터) 를 보지만, 그 지도에 없는 새로운 길을 찾아 나섭니다.
• 목표: 단순히 기존과 비슷한 분자를 만드는 게 아니라, 기존보다 훨씬 더 좋은 약효를 가진 '새로운 분자'를 찾아내는 것입니다.

2. 해결책: "마법의 나침반"과 "스마트한 지도 읽기"

이 모험가 (ChemBFN) 가 성공할 수 있었던 이유는 두 가지 핵심 전략 덕분입니다.

① "보상 게임" (강화 학습)

AI 가 분자를 만들 때, 단순히 "틀리지 않게"만 하는 게 아니라 **"유효한 분자 (약이 될 수 있는 분자) 를 만들면 점수를 주는 게임"**을 시켰습니다.

• 비유: 아이가 레고로 집을 지을 때, "무너지지 않는 집"을 지으면 사탕을 주는 것과 같습니다. 이렇게 하면 AI 는 무의미한 글자 나열이 아니라, 실제로 쓸모 있는 분자 구조를 빠르게 찾아내게 됩니다.

② "스마트한 속도 조절" (ODE 방식)

기존 방식은 분자를 하나씩 천천히 만들어내느라 시간이 매우 오래 걸렸습니다. 이 연구는 이를 고속도로로 바꿨습니다.

• 비유: 흙을 한 줌씩 퍼서 성을 쌓는 대신, 성형수술처럼 한 번에 형태를 잡아주는 방식입니다. 이렇게 하면 기존에 1,000 단계 걸리던 일을 10 단계 만에 끝낼 수 있어, 일반 노트북에서도 빠르게 실행할 수 있게 되었습니다.

3. 핵심 기술: "반자동 운전 모드" (Semi-Autoregressive)

가장 혁신적인 부분은 **'반자동 운전 모드 (SAR)'**를 도입했다는 점입니다.

• 기존 방식 (완전 자동): AI 가 분자의 첫 글자부터 끝까지 한 번에 다 예측하려고 하다가, 앞뒤가 맞지 않아 엉뚱한 분자를 만들어내는 경우가 많았습니다. (비유: 눈가리고 운전)
• 이 연구의 방식 (반자동): AI 는 분자의 앞부분을 보고 뒷부분을 예측하되, 앞으로만 보고 뒤는 보지 않는 규칙을 적용했습니다.
• 효과: 마치 나침반을 들고 길을 찾는 것처럼, AI 는 분자의 앞부분을 기준으로 뒤를 자연스럽게 이어가면서도, 기존 데이터에 갇히지 않고 **새로운 영역 (OOD)**으로 과감하게 뻗어나갈 수 있게 되었습니다.

4. 결과: "기존의 한계를 넘어서다"

이 방법을 실험해 본 결과 놀라운 성과가 나왔습니다.

1. 더 좋은 약 후보: 기존 데이터에 없던, 하지만 약효는 훨씬 좋고 합성하기 쉬운 새로운 분자들을 찾아냈습니다.
2. 단백질 설계: 작은 분자뿐만 아니라, 인체 내에서 중요한 역할을 하는 단백질을 설계하는 데도 성공했습니다.
3. 속도: 기존에 하루 걸리던 작업을 몇 분 만에 끝낼 수 있게 되어, 신약 개발 속도를 획기적으로 높였습니다.

5. 결론: 왜 이 연구가 중요한가?

이 논문은 **"AI 가 단순히 과거를 모방하는 것을 넘어, 미래를 창조할 수 있다"**는 것을 증명했습니다.

• 기존: "우리가 아는 것 중에서 가장 좋은 것을 고르자."
• 이 연구: "우리가 아직 모르는, 하지만 훨씬 더 좋은 새로운 세계를 찾아보자."

이 기술은 새로운 약을 개발하거나, 새로운 소재를 발견하는 과정에서 과학자들이 겪는 수많은 시행착오를 줄여주며, 인류의 건강과 과학 발전에 큰 도움을 줄 것으로 기대됩니다.

---

한 줄 요약:

"기존 AI 는 배운 것을 모방하는 '청소부'였지만, 이 연구는 **보상 게임과 나침반을 들고 새로운 세상을 탐험하는 '모험가'**를 만들어, 약효가 뛰어난 완전히 새로운 분자를 빠르게 찾아냈습니다."

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 신약 개발 및 재료 발견을 위해서는 훈련 데이터의 분포를 벗어난 Out-of-Distribution (OOD) 영역에서 높은 특성을 가진 새로운 분자를 생성하는 것이 필수적입니다.
• 문제점: 기존의 확산 모델 (Diffusion Models, DM) 과 같은 분포 학습 기반 모델들은 훈련 데이터의 분포에 최대한 가깝게 적합하도록 설계되었습니다. 이로 인해 다음과 같은 한계가 존재합니다.
  1. 훈련 데이터와 매우 다른 고유의 (Novel) 특성을 가진 샘플 생성 능력이 부족함.
  2. 다목적 최적화 (Multi-objective optimization) 가 어려움.
  3. 과도한 가이드 (Overconfident guidance) 시 잘못된 영역 (False-positive) 으로 샘플링이 급변할 수 있음.
• 목표: 훈련 데이터보다 더 높은 성능 (예: 높은 약물 유사성, 낮은 합성 난이도, 높은 도킹 점수) 을 가진 분자를 효율적으로 생성할 수 있는 OOD 생성 모델 개발.

2. 제안된 방법론 (Methodology)

이 논문은 베이지안 플로우 네트워크 (Bayesian Flow Networks, BFN), 특히 화학 분자 생성에 특화된 ChemBFN 모델을 기반으로 한 새로운 접근법을 제시합니다.

A. 핵심 모델: ChemBFN

• 기존 확산 모델과 달리 BFN 은 노이즈 분포를 정의하거나 역확산 과정을 학습할 필요가 없습니다. 대신 분포의 매개변수를 더 정보력 있는 방향으로 직접 최적화합니다.
• 이는 연속, 이산, 그리고 이산화된 데이터 모두에 적용 가능하며, 1 차원 분자 표현 (SMILES, SELFIES) 처리에 Encoder-only Transformer 아키텍처를 사용합니다.

B. 주요 기술적 개선 사항

1. 효율적인 샘플링 (Fast Sampling):
   • 보조 강화 학습 (Auxiliary RL) 항: 생성된 분자의 유효성 (Validity) 을 높이기 위해 REINFORCE 알고리즘에서 영감을 받은 RL 손실 항을 훈련 손실에 추가합니다.
   • ODE 유사 생성 과정: 확률 미분 방정식 (SDE) 대신 잠은 공간 (Latent space) 에서 결정론적인 상미분 방정식 (ODE) 솔버와 유사한 과정을 도입하여 샘플링 단계를 대폭 줄이고 속도를 높입니다. 온도 계수 ($\tau$) 를 통해 유효성과 다양성 사이의 균형을 조절합니다.
2. 반자율적 (Semi-Autoregressive, SAR) 전략:
   • 기존 BFN 은 양방향 (Bidirectional) 어텐션을 사용하지만, 저자들은 어텐션 맵에서 주대각선에서 먼 값이 0 에 수렴하는 경향을 발견했습니다.
   • 이를 바탕으로 ** causal mask(인과적 마스크)** 를 적용하여, 현재 토큰이 미래 토큰을 참조하지 못하도록 제한하는 SAR 방식을 도입했습니다.
   • 4 가지 전략 조합: 훈련 (Training) 과 샘플링 (Sampling) 시 SAR 유무에 따라 4 가지 전략 (Normal/SAR 조합) 을 정의하고 실험했습니다.

3. 주요 기여 (Key Contributions)

• 천연 OOD 샘플러로서의 BFN: BFN 이 복잡한 수정 없이도 훈련 데이터와 거리가 먼 OOD 영역을 자연스럽게 탐색할 수 있음을 증명했습니다.
• SOTA 모델 초월: SAR 전략과 RL/ODE 기법을 결합한 ChemBFN 은 기존 최첨단 (SOTA) 모델들 (REINVENT, MOOD, FREED 등) 을 OOD 다목적 최적화 과제에서 압도적으로 능가했습니다.
• 이론적 분석: BFN 과 SAR 의 결합이 왜 OOD 생성에 유리한지 이론적으로 분석했습니다. BFN 의 학습 목적 함수와 정확도 스케줄 (Accuracy Schedule) 이 모델로 하여금 국소성 (Locality) 을 학습하게 하여, 분자를 작은 서브구조로 분해하고 이를 새로운 방식으로 결합할 수 있게 한다는 점을 규명했습니다.
• 범용성: 작은 분자 (Small molecules) 뿐만 아니라 단백질 서열 (Protein sequences) 생성에도 성공적으로 적용되었습니다.

4. 실험 결과 (Results)

A. 소분자 생성 (Small Molecules)

• 샘플링 효율: 기존 ChemBFN 은 1,000 단계가 필요했으나, 제안된 RL + ODE 방식을 적용하면 10~100 단계로도 99% 이상의 유효한 SMILES 를 생성할 수 있게 되었습니다.
• OOD 성능 (ZINC250k 벤치마크):
  • Novel Hit Ratio (새로운 히트 비율): SAR 훈련 + SAR 샘플링 (Strategy 4) 전략이 5 개 타겟 단백질 중 4 개에서 최상위 성능을 보였습니다.
  • Novel Top 5% DS (도킹 점수): 모든 5 개 타겟에서 기존 SOTA 모델들보다 훨씬 낮은 (더 좋은) 도킹 점수를 기록했습니다. 특히 Strategy 3 (SAR 훈련 + Normal 샘플링) 이 가장 높은 점수를 보였습니다.
  • SELFIES 적용: SMILES 대신 SELFIES 형식을 사용할 때 유효성 비율이 6% 미만에서 25% 이상으로 크게 향상되었으며, 모든 SOTA 모델을 능가했습니다.
• 시각화: UMAP 분석 결과, 제안된 모델은 훈련 데이터 공간과 멀리 떨어진 새로운 화학 공간을 탐색하면서도 화학적으로 의미 있는 구조를 생성함을 확인했습니다.

B. 단백질 생성 (Protein Sequences)

• 베타 시트 비율과 용매 접근 표면적 (SASA) 을 최적화하는 라벨이 있는 데이터셋으로 훈련했습니다.
• 훈련 데이터보다 더 높은 목표 값을 가진 단백질 서열을 생성하면서도 자연스러운 단백질과 유사한 자연스러움 (Naturalness, ProtGPT2 로우 가능도 기준) 을 유지했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 신약 설계의 패러다임 전환: 이 연구는 분자 생성 모델이 단순히 훈련 데이터의 복제에 그치지 않고, 훈련 데이터의 분포를 벗어난 (OOD) 고성능 영역을 탐색할 수 있음을 입증했습니다.
• 실용적 가치: 생성된 분자들은 높은 약물 유사성 (QED), 낮은 합성 난이도 (SA), 그리고 우수한 표적 결합력 (DS) 을 동시에 만족하여, 실제 신약 개발 (De novo drug design) 에 즉시 활용 가능한 강력한 도구임을 보여줍니다.
• 계산 효율성: ODE 기반 샘플링과 SAR 전략을 통해 추론 시간을 단축하여, 고사양 GPU 없이도 노트북 수준에서 대규모 샘플링이 가능해졌습니다.

요약하자면, 이 논문은 ChemBFN에 RL, ODE 샘플링, 그리고 SAR 전략을 결합함으로써, 기존 생성 모델의 한계를 극복하고 고성능의 새로운 분자 및 단백질을 효율적으로 설계할 수 있는 새로운 표준을 제시했습니다.