Few-Step Cofolding with All-Atom Flow Maps

Motivation

• Known: AlphaFold3와 Boltz-1 등 EDM 스타일의 diffusion 기반 all-atom cofolding model이 단백질-리간드 복합체의 3D 구조를 원자 수준으로 정확하게 예측하는 표준 패러다임으로 자리잡았으며, SE(3) rigid alignment와 smooth-LDDT 등의 보조 손실이 물리적으로 타당한 구조 생성에 중요하다는 것이 알려져 있다.
• Gap: 기존 all-atom diffusion cofolding model은 정확한 국소 구조(local structure) 예측을 위해 many-step의 반복적 numerical ODE/SDE 시뮬레이션을 필요로 하며, 이는 배치 추론뿐 아니라 SMC, FK-steering, MCTS 등 reward-guided inference-time search 기법의 계산 비용을 더욱 증폭시킨다. 특히 flow map을 통한 few-step distillation은 이미지·비디오 생성 등에서 연구되었지만, all-atom 생물분자 구조 예측에는 아직 적용되지 않았다.
• Why: Structure-based drug design(SBDD)에서 in-silico 후보 구조 생성과 reward 기반 필터링의 실용성은 inference 비용에 직접적으로 좌우되며, few-step으로 물리적으로 타당하고 정확한 구조를 생성할 수 있다면 대규모 후보 탐색과 실시간 reward-guided search가 가능해져 실제 신약 개발 파이프라인에 큰 영향을 줄 수 있다.
• Approach: 사전학습된 Boltz-1 diffusion cofolding model을 teacher로 사용하여, denoiser 기반 flow map 파라미터화에 endpoint loss와 SE(3) weighted rigid alignment를 결합하고, EDM의 σ-space noise schedule에서 동작하도록 변수 변환을 도출하여 all-atom flow map(DECAF)을 distillation한다.

Achievement

Figure 4. DECAF is cost-effective as we increase compute. We compare DECAF-SEARCH and Boltz-1x as we increase NFEs budge

1. 최초의 all-atom protein flow map(DECAF): Boltz-1 teacher로부터 효율적으로 distillation되는 최초의 all-atom cofolding flow map을 설계함.
2. σ-space reparameterization: EDM 스타일 아키텍처의 σ-space noise schedule에서 동작하도록 flow map objective를 변환하는 간단한 change of variables를 도출하여 사전학습 diffusion cofolding model로부터 직접 distillation을 가능하게 함.
3. denoiser 기반 SE(3) rigid alignment: denoiser 기반 flow map parametrization이 ground truth 구조에 대한 SE(3) weighted rigid alignment를 자연스럽게 지원함을 보이고, 이것이 학습 안정성에 critical함을 실증함.
4. DECAF-SEARCH inference 프레임워크: DECAF의 flow map lookahead를 활용해 diffusion model보다 높은 fidelity의 reward 추정을 제공하는 reward-guided search 방법을 제안함.
5. 경험적 성능 향상: 어려운 Runs N' Poses 벤치마크에서 낮은 NFE budget에서 Boltz-1x보다 RMSD 정확도 및 물리적 타당성 점수에서 통계적으로 유의한 개선을 보였으며, PoseBusters에서는 전체 NFE budget의 Boltz-1x(600 NFE) 성능을 20배 적은 compute로 매칭함.

How

Figure 3. Mean RMSD per structure for DECAF-SEARCH (150

• flow map 프레임워크(Boffi et al., 2025b;a)를 기반으로 average velocity $u_{s,t}(x_s)$를 이용한 displacement 파라미터화 $X_{s,t}(x_s) = x_s + (t-s)u_{s,t}(x_s)$를 채택
• Lagrangian, Eulerian, progressive consistency 조건에서 유도되는 PINN 스타일 손실을 이용해 pretrained diffusion teacher(Boltz-1)를 distillation
• denoiser 기반 flow map 파라미터화를 도입하여 endpoint loss가 SE(3) weighted rigid alignment를 자연스럽게 지원하도록 구성
• EDM의 cskip, cin, cout, cnoise 등 σ-dependent preconditioning을 유지하면서 flow map이 σ-space에서 동작할 수 있도록 하는 change of variables 도출
• DECAF의 flow map lookahead를 이용해 reward 함수를 완전히 denoise된 구조가 아닌 중간 예측에서도 고fidelity로 평가 가능하게 하여, 적은 reward query로 reward-guided search(DECAF-SEARCH) 수행
• Runs N' Poses와 PoseBusters 벤치마크에서 RMSD 정확도, 물리적 타당성 점수, NFE-성능 Pareto frontier를 비교 평가

Originality

• 기존 flow map distillation 연구(주로 이미지/비디오 도메인)를 all-atom 생물분자 cofolding 도메인에 최초로 확장함
• denoiser 기반 flow map parametrization과 SE(3) rigid alignment를 결합하여 biomolecular 구조의 대칭성을 명시적으로 다루는 새로운 학습 목적함수를 설계함
• EDM σ-space에서 동작하는 flow map으로의 변수 변환을 고안하여 기존 pretrained EDM 스타일 diffusion cofolding model(Boltz-1)로부터 직접 distillation이 가능하도록 함
• flow map의 lookahead 특성을 SMC/FK-steering/MCTS 등 기존 multi-particle inference-time search 기법과 차별화되는 low-NFE reward-guided search(DECAF-SEARCH)로 재구성함

Limitation & Further Study

• Boltz-1이라는 단일 teacher 모델에 대한 distillation 결과만 제시되어, 다른 최신 cofolding model(AlphaFold3, Chai 등)에 대한 일반화 가능성이 명확히 검증되지 않음
• 적은 NFE에서의 성능 향상이 보고되지만 극단적으로 1-step 생성과 같은 최저 한계 성능 및 실패 모드에 대한 세부 분석이 제한적일 수 있음
• reward-guided search(DECAF-SEARCH)의 reward 함수 설계와 다양한 downstream task(virtual screening, affinity ranking 등)로의 확장성에 대한 추가 검증이 필요함
• workshop paper 형태로 발췌된 내용이라 전체 실험 설정, ablation, 통계적 검증의 세부사항이 충분히 드러나지 않을 가능성이 있음

Evaluation

총평: All-atom cofolding diffusion model을 few-step flow map으로 distillation하고 이를 reward-guided search에 결합한 실용적이고 시의성 있는 연구로, 명확한 방법론적 기여와 실험적 개선을 보여주지만 단일 teacher 모델 의존성과 workshop 규모의 검증이라는 점에서 후속 확장 연구가 필요하다.