Technical Summary: Few-step Cofolding with All-Atom Flow Maps (DECAF)

Problem Statement

The accurate computational modeling of 3D biomolecular complexes, particularly for protein-ligand cofolding, has become a dominant paradigm in structure-based drug design (SBDD). While modern all-atom generative models (e.g., AlphaFold 3, Boltz-1) can predict structures at atomic resolution, they rely on expensive iterative diffusion rollouts. Generating high-fidelity structures typically requires hundreds of neural function evaluations (NFEs), creating a significant computational bottleneck for both deployment and inference-time search techniques.

Furthermore, generating physically valid structures at the atomic level is critical; inaccurate local modeling leads to catastrophic failures such as steric clashes, incorrect bond lengths/angles, and strained side chains. Strict physical constraints often shift the computational burden to inference-time methods, where reward-guided search (e.g., Sequential Monte Carlo, MCTS) is used to refine samples. However, these search methods compound the inference cost, as they require multiple reward queries over multiple particles during the trajectory. The central research question is whether it is possible to train an all-atom cofolding model that generates reward-optimized samples efficiently using only a small number of NFEs at inference time.

Methodology: The DECAF Framework

The authors introduce DENOISER COFOLDING ALL-ATOM FLOWMAP (DECAF), a framework designed to distill state-of-the-art all-atom cofolding diffusion models into few-step flow maps. The methodology addresses specific technical challenges in adapting standard diffusion models to flow maps for all-atom protein modeling.

1. $\sigma$-Space Flow Map Reparametrization

Standard diffusion models often use the EDM (Elucidating the Design Space) parameterization, which conditions on a single noise level $\sigma$. This formulation does not extend cleanly to dual-time flow maps, and the non-linear time reparameterization $\sigma(t)$ used in EDM leads to numerical instability when computing time partials ($\partial \sigma / \partial t$) required for flow map objectives.

To circumvent this, DECAF defines the flow map directly in the $\sigma$-noise space. Instead of mapping time $t$, it maps between noise levels $\rho \to \sigma$ (where $\rho > \sigma$).

• $\sigma$-Velocity: The framework defines an instantaneous $\sigma$-velocity $v_\sigma(x_\sigma) = dx_\sigma/d\sigma$.
• Average Velocity: The flow map is parametrized as a displacement using the average $\sigma$-velocity between two noise levels, allowing the model to take "jumps" along the probability flow ODE trajectory directly in noise space.
• Stability: This reparameterization eliminates the problematic $\partial \sigma / \partial t$ factor, enabling stable training and direct distillation from pretrained EDM-style cofolding models.

2. Denoiser Parametrization with SE(3) Alignment

A critical innovation is the use of a denoiser-based formulation rather than a direct velocity predictor.

• Endpoint Prediction: The model predicts the target endpoint $\hat{x}_{tgt}$ (the clean structure) rather than the instantaneous velocity.
• SE(3) Rigid Alignment: The loss function incorporates a minimization over the SE(3) group (using the Kabsch algorithm) to align the predicted structure $\hat{x}_{tgt}$ with the ground truth (or the teacher's prediction). This softly enforces SE(3) symmetry, which is critical for training stable models and reducing variance.
• Loss Function: The training objective matches the predicted endpoint against the teacher's prediction after rigid alignment, supporting both diagonal (score matching) and off-diagonal (jump) training.

3. DECAF-SEARCH: Inference-Time Framework

To leverage the flow map for reward optimization, the authors introduce DECAF-SEARCH.

• Lookahead Mechanism: Unlike standard diffusion models where reward estimation requires full denoising, DECAF's flow map provides a "lookahead" prediction of the clean endpoint $\hat{x}_0$ from any intermediate noisy state $x_\sigma$. This allows for higher-fidelity reward estimation ($R(\hat{x}_0)$) without completing the full trajectory.
• Search Algorithm: The framework adapts inference-time steering methods (Feynman-Kac, MCTS) to the flow map setting. It uses a stochastic $\gamma$-sampler to generate particles, evaluates their lookahead rewards, and performs clean-space gradient ascent or resampling to steer generation toward physically valid, high-reward configurations.

Key Contributions

1. First All-Atom Protein Flow Map: The design of DECAF-Boltz, the first all-atom protein flow map that successfully distills a pretrained Boltz-1 teacher.
2. Methodological Innovations:
   • A novel $\sigma$-space reparameterization that enables stable training of flow maps for all-atom modeling.
   • A denoiser-based parametrization that enables SE(3) weighted rigid alignment, crucial for stability and accuracy.
3. Inference-Time Search: The introduction of DECAF-SEARCH, a framework that utilizes flow map lookahead to improve reward estimation and alignment, enabling efficient search with fewer NFEs.
4. Empirical Validation: Demonstration that DECAF outperforms diffusion-based baselines in low-NFE regimes and matches full-budget performance with significantly fewer computations.

Experimental Results

The authors evaluated DECAF on the Runs N' Poses (RnP) and PoseBusters benchmarks, comparing against Boltz-1x (a physical potential variant of Boltz-1) and other frontier models like AlphaFold 3 and Pearl.

Performance at Low NFE Budgets

• RnP Benchmark: At strict NFE budgets (e.g., 20–50 NFEs), DECAF-SEARCH significantly outperforms Boltz-1x (tuned) across all metrics (PB-Valid, Success Rate, lDDT-PLI). For instance, at 20 NFEs, DECAF achieves an 83.5% PB-Valid score compared to 75.9% for Boltz-1x.
• Generalization: The performance gain holds across all quartiles of pocket-similarity, including out-of-distribution targets, confirming the model's robustness.
• Efficiency: DECAF matches the performance of Boltz-1x at 600 NFEs using only 30 NFEs (a 20x reduction in compute) on PoseBusters.

Distillation of State-of-the-Art Models

• DECAF-Pearl: The framework was applied to distill the state-of-the-art Pearl cofolding model. DECAF-Pearl matches the teacher's success rate while using 5x fewer NFEs (40 NFEs vs. full simulation).
• Comparison: DECAF-Pearl outperforms other potent baselines (including full simulation of Pearl) in physical validity and success rate metrics at low NFE budgets.

Pareto Frontier

Analysis of the compute-optimal frontier shows that DECAF-SEARCH dominates Boltz-1x across all inference compute budgets. The optimal search strategy varies by budget:

• Low NFE ($\le$ 30): Particle-based SMC (FK-steering) is optimal.
• Moderate NFE (100–250): Monte Carlo gradient estimation (MC-GRAD) is most effective.
• High NFE ($\ge$ 142): MCTS-based search yields the best results.

Significance and Claims

The paper claims that DECAF represents a significant step forward in making all-atom biomolecular generation computationally feasible for high-throughput applications.

• Acceleration: By compressing diffusion trajectories into few-step flow maps, DECAF reduces inference latency by orders of magnitude (e.g., 5x to 20x reductions in NFEs) without sacrificing structural accuracy or physical validity.
• Reward Optimization: The flow map lookahead enables more efficient and accurate reward-guided search, addressing the "inference tax" associated with multiple reward queries in standard diffusion models.
• Practical Impact: The speedup enables new workflows, such as cofolding entire ligand libraries against a target for virtual screening and generating large-scale synthetic datasets for training downstream models, which were previously impractical due to the cost of full diffusion simulations.
• Modesty: The authors acknowledge that failure modes (e.g., non-planar sp2 bonds) are inherited from the reward signals of the teacher models (Boltz-1x) rather than the search procedure itself. They also note that the relationship between NFE budget and success rate is non-monotone, suggesting future work in adaptive search and joint training with task-specific rewards.

The code is released at https://github.com/genesistherapeutics/decaf.