기술 요약: SinAE – 분자, 결정, 단백질을 아우르는 단일 아키텍처 기반 플로우 매칭 오토인코더

1. 문제 정의

3D 원자 구조의 생성은 신약 개발, 재료 설계 및 단백질 공학에 있어 매우 중요합니다. 그러나 현재의 생성 파이프라인은 세 가지 서로 다른 도메인—소분자, 결정, 단백질—에 따라 파편화되어 있습니다. 각 도메인은 고유한 구조적 사전 지식(prior)에 맞춤화된 전문 아키텍처를 사용합니다:

• 분자: 결합 연결성과 원자가 규칙을 처리하기 위해 일반적으로 등변(equivariant) 메시 패싱 네트워크를 사용합니다.
• 결정: 주기적 경계 조건과 격자 기저 불변성(lattice-basis invariance)을 처리해야 합니다.
• 단백질: 잔기 수준의 파라미터화(예: SE(3) 프레임) 또는 트라이앵글 어텐션(triangle attention) 메커니즘에 의존하는 경우가 많습니다.

이러한 파편화는 데이터 부족 문제를 심화시킵니다. 각 도메인별 데이터셋(예: QM9의 134k 분자, MP-20의 45k 결정)은 현대적 기준에서 매우 적은 양이기 때문입니다. 교차 도메인 학습을 통해 이를 완화할 수 있지만, 기존 방식은 두 가지 주요 장애물에 직면해 있습니다:

1. 직접 생성의 한계: 3D 좌표 공간에서 직접 생성하는 방식은 하나의 모델이 유효한 조성, 정밀한 기하학, 그리고 이질적인 구조적 사전 지식(주기성, 잔기 백본, 원자가 등)을 동시에 학습하도록 강제하며, 이는 복잡하고 도메인 특화된 토큰화나 단계적 생성을 요구하게 됩니다.
2. 잠재 오토인코더(Latent Autoencoder)의 한계: 기존의 교차 도메인 잠재 오토인코더는 재구성 충실도(reconstruction fidelity) 문제(손실이 발생하는 디코더가 사전 지식으로도 교정할 수 없는 기하학적 오류를 유발함)와 학습 복잡성(도메인별 손실 조합 및 가중치 튜닝 필요) 문제를 겪습니다. 기존의 어떤 방법도 세 도메인 모두에서 동시에 손실이 없고, 아키텍처가 일반적이며, 통합된 목적 함수를 가진 오토인코더를 구현하지 못했습니다.

2. 방법론: SinAE

저자들은 분자, 결정, 단백질을 위한 단일 아키텍처 플로우 매칭 오토인코더인 SinAE를 소개합니다.

아키텍처

• 통합 백본: SinAE는 인코더와 디코더 모두에 **바닐라 트랜스포머(vanilla Transformer)**를 사용합니다. 이 모델은 등변 연산자, 그래프 메시 패싱, 또는 도메인 특화 모듈(예: 격자 인식 레이어 또는 트라이앵글 어텐션)을 명시적으로 배제합니다.
• 입력 특징 추출(Featurization):
  • 좌표: 불변 기술자(예: 쌍 거리)나 정준 프레임(canonical frame) 구축 없이 가공되지 않은 카테시안 좌표를 사용합니다.
  • 주기성: 결정의 경우, 세 개의 격자 벡터를 시퀀스에 "가상 원자 토큰"으로 추가하여 모델이 격자 행렬을 원자 좌표와 동일하게 처리할 수 있도록 합니다.
  • 등변성: 아키텍처 차원의 등변성 대신, 모델은 훈련 중 **무작위 SO(3) 회전 및 이동 증강(augmentation)**에 의존하여 방향성 강건성을 확보합니다.
• 토큰화: 원자는 학습된 원소 임베딩과 좌표의 선형 투영의 합으로 표현됩니다. 단백질의 경우, 잔기는 $C_\alpha$ 위치에서의 단일 토큰으로 취급됩니다.

훈련 목적 함수

모델은 두 단계로 훈련됩니다:

1. 플로우 매칭 오토인코더:
   • 인코더: 확률적 트랜스포머가 입력 구조 $G=(X, A)$를 각 토큰별 잠재 코드 $z$로 매핑합니다.
   • 디코더: 비등변 트랜스포머가 **반복적 플로우 매칭(iterative flow matching)**을 사용하여 $z$로부터 구조를 재구성합니다. 단일 패스 VAE 디코더와 달리, 디코더는 여러 ODE 단계를 거쳐 구조를 정교화합니다.
   • 손실 함수: 연속적 조건부 플로우 매칭을 통한 좌표 손실($L_{coord}$)과 이산 플로우 매칭을 통한 원자 유형 손실($L_{atom}$), 그리고 가벼운 KL 정규화 항($L_{KL}$)을 결합한 통합 목적 함수를 사용합니다. 도메인별 손실 가중치나 보조 손실은 필요하지 않습니다.
2. 잠재 플로우 매칭 사전 확률(Latent Flow Matching Prior):
   • 바닐라 **디퓨전 트랜스포머(Diffusion Transformer, DiT)**가 인코더 출력에 대한 분포 $p(z)$를 학습합니다.
   • KL 항이 사후 확률을 표준 가우시안으로 밀어내기 때문에, 사전 확률은 등변 제약이 없는 매끄러운 근사 가우시안 잠재 공간 위에서 작동합니다.
   • 생성 과정에서 도메인 레이블은 분류기 없는 가이드(classifier-free guidance)를 통해 DiT가 분자, 결정 또는 단백질로 향하도록 유도하는 조건으로 사용됩니다.

적용된 주요 기여

1. 단일 아키텍처 토큰화: SinAE는 등변 또는 도메인 특화 연산자 없이 단일 바닐라 트랜스포머를 사용하여 분자, 결정, 단백질 백본을 처리하는 최초의 아키텍처입니다.
2. 근사 무손실 재구성: 재구성의 부담을 반복적 플로우 매칭 디코더로 전환함으로써, SinAE는 모든 도메인에서 서브 밀리-옹스트롬(sub-milli-Ångström) 단위의 RMSD를 달로 달성했습니다:
   • 분자: QM9에서 0.0002 Å, GEOM-Drugs에서 0.0006 Å.
   • 결정: MP-20에서 0.0017 Å.
   • 단백질: CASP15에서 약 0.013 Å의 백본 $C_\alpha$-RMSD.
   • 이는 기존 잠재 베이스라인 대비 10~100배 향상된 재구성 충실도를 나타내며, 분자와 결정에 대해 100% 구조 일치율을 보입니다.
3. 긍정적 교차 도메인 전이: QM9과 MP-20에 대한 공동 학습은 단일 도메인 학습에 비해 두 도메인 모두에서 재구성 및 생성 성능을 엄격하게 향상시키며, 이는 공유된 원자 잠재 공간이 전이 가능한 기하학적 규칙성(예: 결합 길이 분포, 배위 패턴)을 포착한다는 증거를 제공합니다.

4. 실험 결과

저자들은 표준 벤치마크(QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF)를 통해 SinAE를 평가했습니다.

• 재구성 충실도: SinAE는 RMSD 및 구조 일치율 측면에서 모든 베이스라인(GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE 포함)을 능가합니다. 반복적 디코더는 원자 배치를 조기에 해결하고 결합 기하학/입체 장애를 나중에 정교화함으로써, 사후 완화(post-hoc relaxation) 없이도 높은 정밀도를 달성합니다.
• 분자 생성: QM9 및 GEOM-Drugs에서, 잠재 DiT 사전 확률은 높은 유효성(PB Valid 점수 각각 0.97 및 0.93)을 달성하여 전문화된 등변 생성기(예: EQGAT-diff, SemlaFlow)와 대등하거나 이를 능가합니다.
• 결정 생성: MP-20에서, 공동 학습된 SinAE 모델은 격자 불변 아키텍처를 사용하는 방법들을 포함한 비교 대상들 중 가장 높은 S.U.N. (Stable, Unique, Novel) 및 M.S.U.N. 점수를 기록했습니다. 이는 일반적인 트랜스포머를 사용했음에도 불구하고 우수한 열역학적 안정성(준안정 구조)을 보여줍니다.
• 단백질 생성: SinAE는 높은 설계 가능성(0.96 scRMSD $\le$ 2.0 Å)과 다양성을 동시에 달성하며 ProteinAE 및 LatentDiff를 능가합니다. 서열-구조 공동 설계 작업에서, 길이 100 및 200에 대해 경쟁력 있는 scTM 및 scRMSD 점수를 달성했습니다.

5. 의의 및 주장

본 논문은 교차 도메인 원자 생성이 충실도를 희생하지 않으면서도 하나의 단순한 아키텍처 아래 통합될 수 있음을 SinAE가 입증한다고 주장합니다.

• 단순성 vs 성능: 저자들은 디코더(반복적 플로우 매칭)가 충분히 강력하다면 복잡한 도메인 특화 연산자가 반드시 필요한 것은 아니라고 주장합니다.
• 데이터 효율성: 결과는 공유된 원자 잠재 공간이 서로 다른 화학적 집단이 서로의 데이터를 통해 성능을 향상시킬 수 있음을 보여주는 직접적인 증거를 제공합니다.
• 기하학 및 사전 지식의 분리: 본 연구는 기하학적 복잡성이 디코더의 반복적 정교화를 통해 처리될 수 있음을 시사하며, 이를 통해 인코더와 사전 확률이 단순화되고 도메인에 무관한 잠재 공간에서 작동할 수 있게 합니다.

인정된 한계점:

• 단백질 평가는 현재 백본 수준의 생성으로 제한되어 있으며, 전원자(all-atom) 모델링은 아직 다루지 않았습니다.
• 현재 모델은 세 가지 도메인 모두에서 단일 체크포인트 내로 속성, 조성 또는 길이를 조건화하여 생성하는 기능을 지원하지 않습니다.
• 더 긴 단백질 체인(>256 잔기)의 경우 더 높은 길이 예산을 가진 재학습이 필요합니다.