Protein structure prediction models following the AlphaFold 3 architecture involve three stages: input preparation (tokenization, retrieval e.g. of MSAs, embeddings), then a transformer-based representation learning “trunk”, and finally a diffusion-based structure module to predict structure. During inference, each denoising step runs one forward pass through the diffusion module, and the AlphaFold 3,¹ Boltz-1,² and Boltz-2³ papers all used 200 sampling steps for prediction.

Recently there have been several works aiming to reduce the computational work of this sort of diffusion-based structure prediction process using flow maps - either distilled from a pretrained model⁴ or trained from scratch⁵.

Below are the results of some simple experiments which take the pretrained Boltz-1 model and look at:

• What the diffusion trajectories look like: the predicted denoised structure often arrives at the broad shape of the final prediction very early on, even after a single step.
• How much we can reduce the number of sampling steps without sacrificing prediction quality too much. This is a much simpler way to achieve more efficient inference than approaches that involve training new models.
• A comparison between outputs of the distogram head and the predicted structures, which suggests that the coarse geometry is often already present in the trunk representations before the diffusion module.

I chose to use Boltz-1 for these experiments because the evaluation code for Boltz-2 is not yet released. For experimentation I sampled a 50-target validation split from the released Boltz-1 test set, plus another 100 targets as a held-out test split. For details of the experimental setup, see appendix.

After running these experiments I came across Protenix-Mini⁶, which made a very similar observation for Protenix models: it is possible to dramatically reduce the number of sampling steps with very little impact on performance. This requires changes to the default AF3-style sampler; see their Section 3.1 and my appendix for more details.

Visualising 200-step trajectories

→ Open the animated trajectory viewer

This viewer shows denoising trajectories for two types of atom coordinates: x0_pred, the predicted denoised coordinates (atom_coords_denoised in Boltz code) and x_t, the post-update coordinates (atom_coords_next in Boltz code).

Looking at x0_pred, the rough shape of the structure often appears very early, even after the first denoising step. (Some counterexamples: 8iha, 8tp8.)

Here are some plots which look at the x0_pred structure throughout the diffusion trajectories and compare it to A) the final prediction, B) the ground truth structure:

We see that predictions gradually get more similar to the final prediction (no surprise), but predictions after a single step are often already quite close to the final prediction.

There is not much improvement in RMSD against ground truth (sometimes it even gets worse over the trajectory). This is consistent with the broad shape of the structure being decided early on. However, the early x0_preds are not plausible structure predictions. Looking at the first frame of an x0_pred trajectory shows implausible atom placements (e.g. 8agr sample 4, even though after just 1 step the RMSD to ground truth is ~1Å). The local atom geometry takes some time to be resolved⁷.

Sweep over the number of sampling steps

Given that we’ve seen x0_pred reaches the rough shape of the final prediction pretty fast, can we get the final prediction more efficiently by changing the sampling approach? The simplest change is just to use fewer steps, though there are various other tricks one could try.

I ran a sweep over the number of denoising steps, trying {100, 50, 25, 20, 15, 10, 5} steps in comparison with the 200-step baseline. I used Boltz’s default noise schedule, which is a Karras-style schedule⁸ with sigma_min=0.0004, sigma_max=160.0, sigma_data=16.0, and rho=8.

Here’s what the noise schedules look like for these:

Here are the results, looking at whether there is a degradation in the prediction quality metrics:

We see that 100 steps shows no degradation at all. 50, 25, and 20 steps are a bit ambiguous with some metrics showing positive changes and some negative, but all very small in magnitude and not statistically significant. 15 steps shows a small but statistically significant degradation in lDDT, but no degradations in RMSD, TM-score, or DockQ.

Decreasing to 10 and 5 steps results in a catastrophic degradation. I later found out that this is largely due to step_scale=1.638 being set higher than 1 and setting step_scale=1 helps. Scarpellini et al.⁴ apply this same fix for their Boltz-1x baselines when using $\le 15$ diffusion steps. See appendix for further details on this issue.

Further reducing the number of steps

The 15-step experiment uses the same Karras noise schedule parameters as the 200-step baseline, just with many fewer steps. However, there is no reason we need to stick to this functional form, or to expect that it is well-suited to the task of structure prediction.

When sampling from a trained diffusion model, […] we need to choose how to space things out as we traverse the different noise levels from high to low. In a range of noise levels that is more important, we’ll want to spend more time evaluating the model, and therefore space the noise levels closer together.

Sander Dieleman, “Noise schedules considered harmful”

To try to bring the number of steps below 15 without further degrading performance, I tested five modifications of the 15-step schedule, each compressing a different block of three steps into one (each of these schedules therefore takes 13 steps). Here are the results compared to the 15-step schedule as a baseline:

There is almost no impact in compressing the first three steps. There is very little impact in compressing the last three steps on RMSD, TM-score, or DockQ, but there is a large and statistically significant drop in lDDT - perhaps the last steps are performing local refinements that lDDT is more sensitive to.

Since compressing the first three steps had almost no effect, and since looking at a trajectory of $x_t$, the initial noise levels seem extremely high, I tried simply starting from a lower initial noise. I used the level reached after 3 steps in the 15-step schedule, $\sigma \approx 568$, resulting in a 12-step schedule. This showed almost no degradation compared to the 15-step schedule.

I then tried two further tweaks to this 12-step schedule starting from lower initial noise:

1. Turn off churn (noise injection)
2. In addition to turning off churn, set step_scale to 1 (default is 1.638)

I tried these after seeing something strange in the trajectory of x_t in the 15-step schedules: after a single step, x_t is already fairly close to the rough shape of the final prediction, then the second and third steps increase the noise a lot. At first I thought this was due to churn (hence variant 1) but it is actually caused by step_scale=1.638, the same issue as mentioned above. Briefly, at each step, the update is of the form

atom_coords_next = (1 - b) * (atom_coords + eps) + b * atom_coords_denoised

for some scalar coefficient b, which can become greater than 1 when step_scale > 1. More details in the appendix. Setting step_scale=1 fixes this “overshoot” issue so I tried applying this to the 12-step schedule to see if performance increases - unfortunately it degraded performance substantially in all metrics.

Here’s a summary of the accuracy of these candidate short schedules compared to the 200-step baseline:

Focussing on the 12-step schedule starting from lower initial noise, here are some more detailed comparisons between the 12-step schedule and the 200-step baseline:

12 steps vs 200 steps on held-out test set

I picked the 12-step schedule from a lower initial noise to evaluate on a test set consisting of 100 targets, a stratified sample of the Boltz-1 test set (see appendix for details) which was held out during the above iteration on schedules.

Open the held-out per-target plot.

The coarse structure is often already visible before the structure module

A distogram is a histogram, or binned probability distribution, of pairwise distances.

Boltz-1 has a distogram head which sits after the trunk and before the diffusion-based structure module. The distogram head simply symmetrises the pair representation $z$ across token pairs, then applies a linear projection from pair-channel dimension to distance-bin logits (Boltz-1 code). The outputs are logits for a distogram of pairwise distances between tokens, not atoms.

The distogram is used in Boltz-1 in two ways: it feeds directly into the loss function for training, and its outputs are also passed into the confidence module. It isn’t used to produce the structure prediction during inference.

I wanted to see if distogram predictions match up with the predicted structures output by the structure module. To compare a distogram with a predicted structure, I compute pairwise distances between the representative atoms in the predicted structure for each token. The distogram is a distribution over distances so I take a crude approach and for each token pair, pick the modal bin, and ask whether the corresponding distance in the predicted structure is in the modal bin, or close to it.

This was done using the default schedule with 200 sampling steps. Here are the results:

The four targets with the worst ±1 Å match rate are all quite complicated, and the Boltz-1 predictions are not close to the reference structures:

The match-rate numbers are a little hard to interpret on their own, so here are three examples in more detail showing the targets with the lowest, median, and highest ±1 Å match rate:

8iks

trunk distogram mode

sample 0

sample 1

sample 2

sample 3

sample 4

scale

sample 0

sample 1

sample 2

sample 3

sample 4

reference

8ec3

trunk distogram mode

sample 0

sample 1

sample 2

sample 3

sample 4

scale

sample 0

sample 1

sample 2

sample 3

sample 4

reference

7y4h

trunk distogram mode

sample 0

sample 1

sample 2

sample 3

sample 4

scale

sample 0

sample 1

sample 2

sample 3

sample 4

reference

The comparisons here are between token-pair distograms and distances between corresponding representative atoms. The structure module does useful work in resolving the structure for all atoms and embedding this in 3D space. However, for the majority of these 50 structures, the distogram head outputs closely match the pairwise distances between representative atoms in the predicted structures. This raises the question of whether the diffusion module plays much of a role at all in determining the broad shape of the structure, or if this is generally already “decided” by the end of the trunk.

Discussion

Bearing in mind that these experiments use a small number of targets, I think these results point towards the following conclusions:

• The number of diffusion steps for Boltz-1 can be decreased substantially from 200 without harming performance. The 12-step schedule we evaluated on the test set did show a statistically significant degradation in average lDDT, although the degradation is not that large in magnitude. Looking at the detailed results, it’s a small number of targets which suffer from large degradations in quality.
• Denoising through the high-noise part of the schedule in particular doesn’t add value.
• The coarse geometry is often already present before diffusion (see the similarity between trunk distogram outputs and final predicted structures).

The final point makes me wonder how much the structure module is genuinely deciding the structure, versus just realising it in 3D.

I think changes to sampling are an attractive way to improve efficiency because they are simple to try, requiring no changes to the pretrained model.

Related work

The closest related work I’m aware of is Protenix-Mini⁶, which speeds up Protenix models in three ways, including using a 2-step sampler for inference. To get this to work they change the sampler to set step_scale=1 and gamma_0=0, i.e. no churn. See Figure 4 of their paper for some results on a small-scale Protenix model. They hypothesised that these trends would generalise to other AF3-style models. I didn’t try 2-step inference but I did try 5- and 10-step schedules with step_scale=1 and the Boltz-1 default gamma_0=0.605 (results in appendix) which perform reasonably well.

Another work which tried few-step sampling for Protenix is DCFold⁹. As a baseline for their distilled single-step model they run the pretrained model with a single sampling step (and also only one recycling step), also setting the step size to 1 and turning off churn - see their Section 4.1 / “AF3 ODE”.

The observation that the highest noise parts of the noise schedule don’t add much value at inference time was also made by Candido et al. for the ESMFold2 model¹⁰. ESMFold2 is a protein complex prediction model which uses representations from hidden layers of a pretrained protein language model, ESM-C, processed using recurrent folding layers, and finally an all-atom diffusion module. At inference, they use a schedule with 68 diffusion steps which is “truncated” from a 100 step schedule by starting from a lower initial noise ($\sigma=256$). See Section A.2.6 of the paper for their discussion.

Scarpellini et al.⁴ train flow map models distilled from Boltz-1 and Genesis’s proprietary Pearl model, to achieve few-step inference, supporting steering. Their Appendix C.1, Fig. 6 reports a comparison between their sampler and standard Boltz-1 sampling without steering, which is therefore comparable with my experiments. Their DECAF flow map method slightly beats standard Boltz-1 sampling at 5, 10, and 20 steps, with DECAF at 20 steps matching Boltz-1 at 200 steps, on the PoseBusters benchmark.

Finally, Kim¹¹ examines the denoising trajectories of AlphaFold 3 and Boltz-2 using sparse autoencoders.

Thanks to Rishabh Anand for feedback on a draft.

ost compare-structures \
  -m <prediction.cif> \
  -r <reference.cif[.gz]> \
  --fault-tolerant \
  --min-pep-length 4 \
  --min-nuc-length 4 \
  -o <structure_out.json> \
  --lddt \
  --bb-lddt \
  --qs-score \
  --dockq \
  --ics \
  --ips \
  --rigid-scores \
  --tm-score

atom_coords_next = 
  atom_coords_noisy
+ step_scale * (sigma_t - t_hat) * (atom_coords_noisy - atom_coords_denoised) / t_hat

atom_coords_next = (1 - b) * atom_coords_noisy + b * atom_coords_denoised

sigma_t / sigma_tm < (step_scale - 1) * (1 + gamma) / step_scale