기술 요약: 전원자 플로우 맵을 이용한 소수 단계 공폴딩 (DECAF)

문제 정의

3D 생체 분자 복합체, 특히 단백질-리간드 공폴딩(cofolding)의 정확한 계산 모델링은 구조 기반 약물 설계(SBDD)의 지배적인 패러다임이 되었습니다. 현대적인 전원자(all-atom) 생성 모델(예: AlphaFold 3, Boltz-1)은 원자 해상도로 구조를 예측할 수 있지만, 비용이 많이 드는 반복적인 확산 롤아웃(diffusion rollouts)에 의존합니다. 고충실도 구조를 생성하려면 일반적으로 수백 번의 신경 함수 평가(NFE)가 필요하며, 이는 배포 및 추론 시 탐색 기법 모두에 상당한 계산 병목 현상을 초래합니다.

나아가, 원자 수준에서 물리적으로 유효한 구조를 생성하는 것은 매우 중요합니다. 부정확한 국소 모델링은 입체 충돌(steric clashes), 잘못된 결합 길이/각도, 변형된 측쇄(side chains)와 같은 치명적인 실패로 이어질 수 있습니다. 엄격한 물리적 제약 조건은 종종 추론 시점의 방법론으로 계산 부담을 전가하며, 여기서 보상 가이드 탐색(예: 순차적 몬테카를로, MCTS)이 샘플을 정교화하기 위해 사용됩니다. 그러나 이러한 탐색 방법은 궤적 동안 여러 입자에 대해 여러 번의 보상 쿼리를 수행해야 하므로 추론 비용을 가중시킵니다. 핵심 연구 질문은 단 몇 번의 NFE만으로도 효율적으로 보상 최적화된 샘플을 생성할 수 있는 전원자 공폴딩 모델을 학습시키는 것이 가능한가 하는 점입니다.

방법론: DECAF 프레임워크

저자들은 최첨단 전원자 공폴딩 확산 모델을 소수 단계의 플로우 맵(flow maps)으로 증류(distill)하도록 설계된 프레임워크인 **DENOISER COFOLDING ALL-ATOM FLOWMAP (DECAF)**를 소개합니다. 이 방법론은 표준 확산 모델을 전원자 단백질 모델링을 위한 플로우 맵으로 적응시키는 과정에서의 특정 기술적 과제들을 해결합니다.

1. $\sigma$-공간 플로우 맵 재매개변수화

표준 확산 모델은 종종 단일 노이즈 레벨 $\sigma$에 의존하는 EDM(Elucidating the Design Space) 매개변수화를 사용합니다. 이러한 형식은 듀얼 타임(dual-time) 플로우 맵으로 깔끔하게 확장되지 않으며, EDM에서 사용되는 비선형 시간 재매개변수화 $\sigma(t)$는 플로우 맵 목적 함수에 필요한 시간 편미분($\partial \sigma / \partial t$)을 계산할 때 수치적 불안정성을 유발합니다.

이를 우회하기 위해, DECAF는 플로우 맵을 직접 $\sigma$-노이즈 공간에서 정의합니다. 시간 $t$를 매핑하는 대신, 노이즈 레벨 $\rho \to \sigma$ ($\rho > \sigma$) 사이를 매핑합니다.

• $\sigma$-속도: 프레임워크는 순시 $\sigma$-속도 $v_\sigma(x_\sigma) = dx_\sigma/d\sigma$를 정의합니다.
• 평균 속도: 플로우 맵은 두 노이즈 레벨 사이의 평균 $\sigma$-속도를 이용한 변위로 매개변수화되며, 이를 통해 모델이 노이즈 공간 상에서 확률 흐름 ODE 궤적을 따라 직접 "점프"할 수 있게 합니다.
• 안정성: 이 재매개변수화는 문제가 되는 $\partial \sigma / \partial t$ 인자를 제거하여, 사전 학습된 EDM 스타일의 공폴딩 모델로부터 안정적인 학습과 직접적인 증류를 가능하게 합니다.

2. SE(3) 정렬을 이용한 디노이저 매개변수화

핵심적인 혁신은 직접적인 속도 예측자가 아닌 디노이저 기반 형식을 사용하는 것입니다.

• 엔드포인트 예측: 모델은 순시 속도가 아니라 목표 엔드포인트 $\hat{x}_{tgt}$ (깨끗한 구조)를 예측합니다.
• SE(3) 강체 정렬: 손실 함수는 예측된 구조 $\hat{x}_{tgt}$를 실제 값(또는 교사 모델의 예측값)과 정렬하기 위해 SE(3) 군에 대한 최소화(Kabsch 알고리즘 사용)를 포함합니다. 이는 SE(3) 대칭성을 부드럽게 강제하여 모델의 안정성을 높이고 분산을 줄이는 데 결정적입니다.
• 손실 함수: 학습 목적 함수는 예측된 엔드포인트를 강체 정렬 후 교사 모델의 예측값과 일치시키며, 이는 대각(score matching) 및 오프-대각(jump) 학습을 모두 지원합니다.

3. DECAF-SEARCH: 추론 시 프레임워크

플로우 맵을 보상 최적화에 활용하기 위해 저자들은 DECAF-SEARCH를 도입했습니다.

• 룩어헤드(Lookahead) 메커니즘: 전체 디노이징이 보상 추정을 위해 필요한 표준 확산 모델과 달리, DECAF의 플로우 맵은 임의의 중간 노이즈 상태 $x_\sigma$로부터 깨끗한 엔드포인트 $\hat{x}_0$에 대한 "룩어헤드" 예측을 제공합니다. 이를 통해 전체 궤적을 완료하지 않고도 고충실도 보상 추정($R(\hat{x}_0)$)이 가능합니다.
• 탐색 알고리즘: 이 프레임워크는 추론 시 스티어링(steering) 방법(Feynman-Kac, MCTS)을 플로우 맵 설정에 맞게 조정합니다. $\gamma$-샘플러를 사용하여 입자를 생성하고, 룩어헤드 보상을 평가하며, 물리적으로 유효하고 보상이 높은 구성으로 생성을 유도하기 위해 클린 공간(clean-space) 그래디언트 상승 또는 리샘플링을 수행합니다.

주요 기여

1. 최초의 전원자 단백질 플로우 맵: 사전 학습된 Boltz-1 교사 모델을 성공적으로 증류한 최초의 전원자 단백자 플로우 맵인 DECAF-Boltz를 설계했습니다.
2. 방법론적 혁신:
   • 전원자 모델링의 안정적인 학습을 가능하게 하는 새로운 $\sigma$-공간 재매개변수화.
   • 안정성과 정확도에 필수적인 SE(3) 가중 강체 정렬을 가능하게 하는 디노이저 기반 매개변수화.
3. 추론 시 탐색: 플로우 맵 룩어헤드를 활용하여 보상 추정과 정렬을 개선함으로써, 더 적은 NFE로 효율적인 탐색을 가능하게 하는 DECAF-SEARCH 프레임워크를 도입했습니다.
4. 실증적 검증: DECAF가 낮은 NFE 환경에서 확산 기반 베이스라인보다 우수한 성능을 보이며, 훨씬 적은 계산량으로 전체 예산 성능에 도달함을 입증했습니다.

실험 결과

저자들은 Runs N' Poses (RnP) 및 PoseBusters 벤치마크에서 DECAF를 Boltz-1x(Boltz-1의 물리적 포텐셜 변형) 및 AlphaFold 3, Pearl과 같은 최첨단 모델들과 비교 평가했습니다.

낮은 NFE 예산에서의 성능

• RnP 벤치마크: 엄격한 NFE 예산(예: 20–50 NFE)에서 DECAF-SEARCH는 모든 지표(PB-Valid, Success Rate, lDDT-PLI)에서 Boltz-1x(튜닝됨)를 크게 능가합니다. 예를 들어, 20 NFE에서 DECAF는 Boltz-1x의 75.9%에 비해 83.5%의 PB-Valid 점수를 달성합니다.
• 일반화: 성능 향상은 포켓 유사도의 모든 사분위수에서 유지되며, 이는 OOD(out-of-distribution) 타겟에 대해서도 모델의 견고함을 확인시켜 줍니다.
• 효율성: DECAF는 PoseBusters에서 단 30 NFE를 사용하여 600 NFE에서의 Boltz-1x 성능에 도달합니다 (20배의 계산량 감소).

최첨단 모델의 증류

• DECAF-Pearl: 이 프레임워크는 최첨단 Pearl 공폴딩 모델을 증류하는 데 적용되었습니다. DECAF-Pearl은 교사 모델의 성공률을 유지하면서도 5배 적은 NFE(전체 시뮬레이션 대비 40 NFE)를 사용합니다.
• 비교: DECAF-Pearl은 낮은 NFE 예산에서 물리적 유효성 및 성공률 지표에서 다른 강력한 베이스라인(전체 시뮬레이션된 Pearl 포함)을 능가합니다.

파레토 프런티어 (Pareto Frontier)

계산 최적 프런티어 분석 결과, DECAF-SEARCH는 모든 추론 계산 예산에서 Boltz-1x를 압도합니다. 최적의 탐색 전략은 예산에 따라 다릅니다:

• 낮은 NFE ($\le$ 30): 입자 기반 SMC (FK-steering)가 최적입니다.
• 중간 NFE (100–250): 몬테카를로 그래디언트 추정 (MC-GRAD)이 가장 효과적입니다.
• 높은 NFE ($\ge$ 142): MCTS 기반 탐색이 가장 좋은 결과를 냅니다.

의의 및 주장

본 논문은 DECAF가 고처리량 애플리케이션을 위해 전원자 생체 분자 생성을 계산적으로 실현 가능하게 만드는 중요한 진전임을 주장합니다.

• 가속화: 확산 궤적을 소수 단계 플로우 맵으로 압축함으로써, DECAF는 구조적 정확도나 물리적 유효성을 희생하지 않고 추론 지연 시간을 몇 자릿수 감소시킵니다 (예: NFE 5~20배 감소).
• 보상 최적화: 플로우 맵 룩어헤드는 표준 확산 모델에서 발생하는 "추론 세금(inference tax)" 문제를 해결하며, 더 효율적이고 정확한 보상 가이드 탐색을 가능하게 합니다.
• 실질적 영향: 이러한 속도 향상은 가상 스크리닝을 위해 타겟에 대한 전체 리간드 라이브러리를 공폴딩하거나, 기존에는 전체 확산 시뮬레이션 비용 때문에 불가능했던 대규모 합성 데이터셋을 생성하는 등의 새로운 워크플로우를 가능하게 합니다.
• 겸허함: 저자들은 실패 모드(예: 비평면 sp2 결합)가 탐색 절차 자체가 아닌 교사 모델(Boltz-1x)의 보상 신호에서 유래했음을 인정합니다. 또한 NFE 예산과 성공률 사이의 관계가 비단조적(non-monotone)임을 언급하며, 적응형 탐색 및 특정 작업 보상과의 공동 학습에 대한 향후 연구 과제를 제시합니다.

코드는 https://github.com/genesistherapeutics/decaf에 공개되어 있습니다.