# Few-step Cofolding with All-Atom Flow Maps

arXiv: https://arxiv.org/abs/2606.08375 · HTML: https://mlsanity.com/paper/2606.08375/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 53 references that are in the map.
- First posted on arXiv 7 Jun 2026.

## Builds on (most cited first, 30 of 31)

- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [Universal Guidance for Diffusion Models](https://mlsanity.com/paper/2302.07121/) (2023)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [A General Framework for Inference-time Scaling and Steering of Diffusion Models](https://mlsanity.com/paper/2501.06848/) (2025)
- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [Improved Mean Flows: On the Challenges of Fastforward Generative Models](https://mlsanity.com/paper/2512.02012/) (2025)
- [Flow map matching with stochastic interpolants: A mathematical framework for consistency models](https://mlsanity.com/paper/2406.07507/) (2024)
- [How to build a consistency model: Learning flow maps via self-distillation](https://mlsanity.com/paper/2505.18825/) (2025)
- [Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts](https://mlsanity.com/paper/2503.02819/) (2025)
- [Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review](https://mlsanity.com/paper/2501.09685/) (2025)
- [Align Your Flow: Scaling Continuous-Time Flow Map Distillation](https://mlsanity.com/paper/2506.14603/) (2025)
- [Flow Map Language Models: One-step Language Modeling via Continuous Denoising](https://mlsanity.com/paper/2602.16813/) (2026)
- [Proteina: Scaling Flow-based Protein Structure Generative Models](https://mlsanity.com/paper/2503.00710/) (2025)
- [One-step Latent-free Image Generation with Pixel Mean Flows](https://mlsanity.com/paper/2601.22158/) (2026)
- [Categorical Flow Maps](https://mlsanity.com/paper/2602.12233/) (2026)
- [Meta Flow Maps enable scalable reward alignment](https://mlsanity.com/paper/2601.14430/) (2026)
- [Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models](https://mlsanity.com/paper/2506.20701/) (2025)
- [La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching](https://mlsanity.com/paper/2507.09466/) (2025)
- [Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps](https://mlsanity.com/paper/2602.05993/) (2026)
- [Discrete Flow Maps](https://mlsanity.com/paper/2604.09784/) (2026)
- [Test-time scaling of diffusions with flow maps](https://mlsanity.com/paper/2511.22688/) (2025)
- [Scaling Atomistic Protein Binder Design with Generative Pretraining and Test-Time Compute](https://mlsanity.com/paper/2603.27950/) (2026)
- [How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance](https://mlsanity.com/paper/2604.27147/) (2026)
- [General Multimodal Protein Design Enables DNA-Encoding of Chemistry](https://mlsanity.com/paper/2604.05181/) (2026)
- [Pearl: A Foundation Model for Placing Every Atom in the Right Location](https://mlsanity.com/paper/2510.24670/) (2025)

## Built on by (most cited first, 1 of 1)

- [OMatG-flash: An All-Atom Flow Map with Reinforce Adjoint Matching for Scalable Materials Discovery](https://mlsanity.com/paper/2609.26402/) (2026)

## Abstract (arXiv)

All-atom generative modeling of 3D 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 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 σ-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-Boltz 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. Distilling the state-of-the-art Pearl cofolding model, DeCAF-Pearl outperforms diffusion-based cofolding models and matches its teacher on success rate while using 5x fewer NFEs. We release our code at https://github.com/genesistherapeutics/decaf.