Few-Step Cofolding with All-Atom Flow Maps
Gianluca Scarpellini ⋅ Ron Shprints ⋅ Peter Holderrieth ⋅ Juno Nam ⋅ Pranav Murugan ⋅ Rafael Gomez-Bombarelli ⋅ Tommi Jaakkola ⋅ Maruan Al-Shedivat ⋅ Nicholas Boffi ⋅ Joey Bose
Abstract
All-atom generative modeling of $3\mathrm{D}$ biomolecular complexes has emerged as the dominant paradigm for predicting the structure of proteins and protein-ligand systems. Generating structures at the atomic level of fidelity, however, typically requires expensive iterative diffusion rollouts, making both conventional deployment and inference-time search techniques computationally costly. In this paper, we introduce the Denoiser Cofolding All-atom Flowmap (DeCAF) framework for distilling state-of-the-art all-atom cofolding models into all-atom flow maps that produce high-quality samples in only a few inference steps. We build DeCAF on a denoiser-based formulation of flow maps with endpoint losses that naturally support $\mathrm{SE(3)}$ rigid alignment, which we show is critical for training accurate models. We further derive a simple change of variables that lets DeCAF operate in the $\sigma$-space noise schedule of EDM-style architectures, enabling direct distillation from pretrained cofolding diffusion models. Equipped with DeCAF's flowmap lookahead, we introduce a purpose-built inference-time framework that improves sampling through reward-guided search. Empirically, DeCAF statistically improves over Boltz-1x in both accuracy (RMSD) and physical validity scores of protein-ligand poses at strict NFE budgets on the challenging Runs N Poses, while also showing a more optimal Pareto frontier across all inference compute budgets on PoseBusters.
Video
Chat is not available.
Successful Page Load