기술 요약: KRONOS – 3D 분자 생성을 위한 자기회귀적 잠재 확산 모델 (Autoregressive Latent Diffusion)

문제 정의

3차원(3D) 분자 생성은 신약 개발 및 분자 설계에서 매우 중요한 과제입니다. 현재의 최첨단 접근 방식은 주로 **확산 모델(diffusion models)**과 **플로우 매칭(flow matching)**에 의존합니다. 이러한 방법들은 높은 생성 품질을 달성하지만, 일반적으로 분자의 크기(원자 수)를 사전에 지정해야 하므로 가변 길이 생성에서의 유연성이 제한됩니다.

반면, 자기회귀(autoregressive, AR) 모델은 가변 길이 생성과 부분적인 분자 구조(예: 알려진 스캐폴드의 정교화)에 대한 조건부 생성을 자연스럽게 지원하므로 의약 화학 분야에서 매우 가치가 높습니다. 그러나 기존의 AR 모델은 두 가지 주요 과제에 직면해 있습니다:

1. 성능 격차: 역사적으로 AR 모델은 생성 품질, 특히 기하학적 충실도(geometric fidelity)와 화학적 타당성(chemical validity) 측면에서 확산 모델에 비해 뒤처져 왔습니다.
2. 목적의 균형: 단일 아키텍처 내에서 무조건적(unconditional) 생성(de novo 설계)과 파편 조건부(fragment-conditioned) 생성(분자 완성) 모두에서 강력한 성능을 달방하는 것은 여전히 어렵습니다. 많은 AR 모델은 아키텍처 수정이나 성능 저하 없이 두 패러다임 사이를 일반화하는 데 어려움을 겪습니다.

또한, 기존의 잠재 분자 생성기(latent molecular generators)는 거의 예외 없이 확산 모델에 의존해 왔으며, **잠재 공간에서의 자기회귀적 생성(autoregressive generation in latent space)**의 잠재력은 거의 탐구되지 않은 상태로 남아 있습니다.

방법론: KRONOS

저자들은 사전 학습된 통합 오토인코더(Unified Autoencoder, UAE)의 잠재 공간에서 직접 3D 분자를 생성하는 잠재 자기회귀 확산 프레임워크인 KRONOS를 소개합니다. 이 방법은 세 가지 핵심 구성 요소를 결합합니다:

1. 잠재 자기회귀 생성 (Latent Autoregressive Generation)

KRONOS는 원시 데이터 공간(이산적인 원자 유형/결합과 연속적인 좌표를 동시에 예측)에서 작동하는 대신, 사전 학습된 UAE(Luo et al., 2025)의 연속적인 잠재 공간에서 작동합니다.

• 표현(Representation): 3D 분자는 원자의 정체성, 결합 환경 및 3D 기하 구조를 공동으로 인코딩하는 일련의 잠재 토큰 $Z = \{z_1, \dots, z_n\}$으로 인코딩됩니다.
• 인수분해(Factorization): 모델은 좌측에서 우측으로 진행되는 자기회귀 분포를 학습합니다:
  $$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$$
• 가변 길이(Variable Length): 잠재 공간에는 자연스러운 문장 종료 토큰(end-of-sequence token)이 없으므로, KRONOS는 시퀀스 종료를 동적으로 예측하기 위해 생성 모델과 함께 학습되는 전용 중단 분류기(stop classifier)(경량 MLP)를 채택하여, 사전 정의된 원자 수 없이도 가변 길이를 구현합니다.

2. 확산 기반 토큰 예측 (Diffusion-Based Token Prediction)

연속적인 잠재 토큰의 조건부 분포 $p_\theta(z_i | z_{<i})$를 모델링하기 위해, KRONOS는 확산 손실(Diffusion Loss) 공식(Li et al., 2024)을 활용합니다.

• 인과적 트랜스포머(causal transformer)가 이전 토큰들을 컨텍스트 표현 $c_i$로 인코딩합니다.
• 확산 디노이징 네트워크 $\epsilon_\theta$는 $c_i$와 확산 타임스텝(timestep)에 조건화되어 현재 잠재 토큰 $z_i$에 추가된 노이즈를 예측합니다.
• 이 접근 방식은 이산적인 토큰 예측을 디노이징 목적 함수로 대체하여, 이산화 과정 없이 연속 변수를 유연하게 모델링할 수 있게 합니다.

3. 혼합 학습 전략 (파편 조건부 생성)

단일 모델 내에서 무조건적 생성과 파편 조건부 생성을 모두 지원하기 위해, 저자들은 FIM(Fill-in-the-Middle) 패러다임에서 영감을 얻은 학습 전략을 도입했습니다.

• 파편 분해(Fragment Decomposition): 분자는 BRICS 기반 스킴을 사용하여 화학적으로 의미 있는 파편들로 분해됩니다.
• 시퀀스 구성: 학습 과정에서 각 분자는 다음 중 하나로 변환됩니다:
  • 기본 순서 시퀀스(default-order sequence): 무조건적 생성을 위함.
  • 파편 조건부 시퀀스: 무작위로 선택된 파편 집합이 "시드(seed)" 역할을 하며, 나머지 파편들이 생성 대상이 됩니다.
• 메커니즘: 모델은 시드 파편들에 조건화되어 대상 파편들을 자기회귀적으로 생성하도록 학습됩니다. 이는 모든 순열을 마진화(marginalizing)하는 대신(Order-Agnostic AR 모델과 달리), 파편 및 그 내부 원자들에 대해 단일 결정론적 순서를 할당함으로써 달성됩니다. 이를 통해 모델은 아키텍처 변경 없이도 분자 완성 능력을 학습할 수 있습니다.

주요 기여

1. KRONOS 프레임워크: 통합된 잠재 공간 내에서 분자 그래프 토폴로지와 기하 구조를 공동으로 모델링하는 새로운 잠재 자기회귀 확산 프레임워크를 제안합니다.
2. 통합 학습 전략: FIM 패러다임에서 영감을 얻어, 단일 좌측-우측 자기회귀 모델 내에서 무조건적 생성과 파편 조건부 생성을 모두 가능하게 하는 혼합 학습 접근 방식을 제공합니다.
3. 가변 길이 생성: 중단 분류기를 통해 사전에 크기를 지정할 필요 없이 가변 길이를 자연스럽게 지원하는 통합 프레임워크를 구축했습니다.
4. 성능: 잠재 자기회귀 생성이 사전 정의된 크기 지정 없이도 확산 모델과 경쟁 가능한 수준의 품질을 유지하면서, 자기회귀 방식 중에서는 선도적인 성능을 달성할 수 있음을 입증했습니다.

실험 결과

저자들은 QM9 및 GEOM-Drugs 데이터셋을 통해 KRONOS를 평가했습니다.

무조건적 생성 (Unconditional Generation)

• QM9: KRONOS는 자기회귀 방식 중 가장 높은 **3D 타당성(97.3%)**과 가장 강력한 전반적 성능을 달성했습니다. 또한 3D 원자 안정성, 결합 길이, 결합 각도 MMD를 포함한 지표에서 선도적인 확산 모델(예: JODO, EDM)과 대등한 경쟁력을 보였습니다.
• GEOM-Drugs: KRONOS는 평가된 모든 베이스라인(자기회귀 및 확산 모델 모두) 중에서 가장 강력한 전반적 성능을 기록했습니다.
• 물리적 타당성: 반경험적 기하 구조 최적화(GFN2-xTB)를 사용했을 때, KRONOS가 생성한 컨포메이션(conformations)은 최적화된 구조와 비교하여 낮은 RMSD 및 에너지 차이($\Delta E$)를 보였으며, 이는 생성된 구조가 물리적으로 타당한 영역 근처에 위치함을 나타냅전합니다.

파편 조건부 생성 (Fragment-Conditioned Generation)

• 약 15k개의 분자-파편 쌍 데이터셋에 대해, KRONOS는 타당성 및 안정성 측면의 Mean@3 및 Best@3 지표 모두에서 자기회귀 베이스라인인 QUETZAL을 크게 앞질렀습니다.
• 결정적으로, 파편 조건부 생성을 활성화하더라도 무조건적 생성 성능에 미치는 영향이 미미함을 보여주었으며, 이는 단일 아키텍처 내에서 두 패러다임을 모두 지원할 수 있음을 입증합니다.

의의 및 주장

본 논문은 KRONOS가 기존의 데이터 공간 자기회귀 및 확산 기반 접근 방식에 대한 유망한 대안을 제공한다고 주장합니다. 본 연구의 의의는 다음과 같습니다:

• 격차 해소: 자기회귀 모델과 확산 모델 사이의 성능 격차를 실질적으로 좁혔으며, 확산 베이스라인과 대등한 수준을 유지하면서 자기회귀 방식 중 선도적인 결과를 달성했습니다.
• 유연성: 자기회귀 모델의 강점(가변 길이, 부분 조건부 생성)과 확산 모델의 고품질 생성 능력을 단일 잠재 공간 공식 내에서 통합했습니다.
• 실용적 유용성: 아키텍처 변경 없이 파편 조건부 생성을 지원함으로써, 분자가 종종 알려진 스캐폴드를 확장하며 설계되는 의약 화학 작업에 실질적인 도구를 제공합니다.

저자들은 사전 학습된 잠재 공간의 표현력에 대한 의존성과 고정된 BRICS 분해 방식 사용을 한계점으로 언급하며, 향후 연구에서 파편 추상화를 생성 모델과 함께 공동 학습하거나 속성 가이드 설계를 확장할 수 있음을 시사했습니다.