How to Spend Your Oracle Budget: Practical Guidance for Protein Structure Prediction Models

AI summary

Overview

Research area: Guidance and optimisation of generative protein structure prediction models under limited evaluation budgets (AI for Science / computational biology).

Technical level: Intermediate. Readers will benefit from familiarity with diffusion/generative models, Bayesian optimisation, and preference-based fine-tuning, although the paper's practical recommendations are stated plainly.

Scope in one sentence: A systematic benchmark of four oracle-budget-aware guidance strategies (O3, FK-steering, DPO, and Best K-of-N) applied to the Boltz-2 protein structure prediction model on two protein targets, distilled into practical guidance for choosing a method at a given oracle budget.

What This Paper Is About

Protein structure foundation models such as AlphaFold3, Chai-1, and Boltz-2 can predict structures from sequence alone, but they remain unreliable on some targets — collapsing onto one conformation, producing physically implausible geometries, or mis-assembling complexes. External oracles (for example molecular dynamics simulations or wet-lab assays) can flag and correct these failures, but they are expensive, making the number of oracle calls a hard constraint. This paper asks which guidance method spends a limited oracle budget most effectively, and provides the first practical reference for that choice.

Key Contributions

1. First systematic comparison of oracle-budget-aware guidance methods. The paper benchmarks FK-steering, DPO, Best K-of-N sampling, and Optimisation Over Outputs (O3) under a shared constraint of exactly N oracle calls and a returned batch of K candidates, across six (N, K) configurations.

2. First application of O3 to a protein structure prediction model. The authors extend the O3 framework of Willis et al. (2025) to Boltz-2, which required converting Boltz-2's stochastic generation into its equivalent probability-flow formulation (Karras et al., 2022) and disabling stochastic SE(3) augmentation, since O3 needs a deterministic generator.

3. A transferable decision recipe. The paper converts its findings into actionable advice: use O3 at constrained oracle budgets, consider FK-steering at moderate budgets, and reach for DPO once budgets are large enough to fine-tune.

4. Evaluation on two targets and two oracle types. Results are reported for calmodulin (PDB: 1CLL) with a TM-score oracle, and for E. coli aspartate transcarbamoylase (PDB: 9EEH) with a reference-free MolProbity oracle.

Main Findings

• No universal winner, but O3 leads in the tested range: The abstract states that no single method consistently dominates across all budgets and oracles; the conclusions state that across the budgets evaluated, O3 dominates the performance.

• O3 is strongest at low-to-mid budgets on 1CLL: Using mean-of-K TM-score, O3 outperforms all baselines at N ≤ 1000 with the mean value plateauing at approximately 0.81. It is the only method that meaningfully improves on the Best K-of-N baseline at low oracle budgets (N ≤ 100).

• Best K-of-N is flat: Best K-of-N stays roughly constant at approximately 0.60 across all budgets, consistent with a fixed K/N ratio.

• FK-steering and DPO improve with budget but start weak: Neither is competitive with Best K-of-N at low budgets (N ≤ 100); FK-steering achieves only approximately 0.55 at N = 20, then reaches approximately 0.73 at N = 1000. DPO improves with N but reaches only approximately 0.71 at N = 1000, below FK-steering and O3 within the evaluated range.

• Both the subspace and the optimiser matter for O3: At N = 100, uniform random sampling within the O3 subspace already beats Best K-of-N, and the full O3 result comes from combining subspace construction with Bayesian optimisation.

• The best O3 subspace dimension depends on budget: When N = 20, the best-performing dimension is d = 6; when N = 200, d = 10 yields better predictions. There is a trade-off between the richness of the subspace (more seed latents d) and the difficulty of optimising within it.

• Higher FK-steering reward scaling helps, at a possible diversity cost: With K and N fixed, increasing λ improves performance across all budgets; a high λ = 50 amplifies the oracle signal and drives mean-of-K and max-of-K closer together, which likely hinders output diversity. Increasing the number of resampling steps N/K at the cost of particle population K generally yields better performance at higher oracle budgets.

• Online DPO beats offline DPO: At N = 1000, online DPO reaches a mean TM-score of 0.708 and a max TM-score of 0.783, while offline DPO plateaus at 0.55–0.57 across all budgets. The gap is negligible at N = 100 and becomes substantial at N = 500, suggesting gains come primarily from adaptive on-policy resampling rather than extra optimisation steps. A large oracle budget is required to support effective fine-tuning at foundation-model scale.

• On 9EEH with the MolProbity oracle: O3 scores highest at all but the smallest budget, where it matches Best K-of-N. Best K-of-N is relatively constant across budgets, DPO improves with N but performs significantly worse than Best K-of-N at low budgets, and FK-steering performs poorly throughout without improving with budget — unlike 1CLL, where it beats Best K-of-N above N = 100. Under max-of-K, Best K-of-N is strongest at large budgets and both it and FK-steering improve with N, while O3 is unchanged. The MolProbity score has a narrow effective range, with most structures scoring between 0.5 and 0.7 under this implementation.

Methodology in Plain English

The setup: a frozen generative model g (Boltz-2) maps random noise z to a structure x, and an oracle r scores each structure. Every method gets exactly N oracle calls and must return K candidate structures.

O3 (Bayesian optimisation in a learned subspace). Sample M structures and score them; keep the d highest-scoring latents as "seeds" forming a matrix Z. These seeds define a low-dimensional subspace U of dimension d − 1, where a point u is mapped to a latent via a Knothe–Rosenblatt transform (producing weights on the simplex) followed by an LOL projection (ℓ(w, Z) = wᵀZ). Once the d seeds are chosen the subspace is fixed and never rebuilt. A Gaussian process is fitted on d + 2 initial points (the d projected seeds plus 2 random points in U), and the remaining budget is spent through Bayesian optimisation rounds using Log Expected Improvement, with an RBF kernel, constant mean, single-task GP, and Monte Carlo acquisition from BOTorch.

FK-steering (inference-time). M interacting particles are generated and resampled proportionally to weights based on the oracle reward, using the reward difference across the trajectory. Because most biological oracles are non-differentiable, the original Boltz-2 transition kernel is used. With a budget of N and a batch of K distinct samples, M ≥ K particles each get N/M oracle calls. The oracle is called once at the start and once at the end of denoising, with the remainder spread over the first 3/4 of the trajectory (Boltz-2 stops injecting noise in the last quarter to avoid atomic jitter).

DPO (fine-tuning). Preference pairs are formed from the N scored structures: a preferred structure above the median oracle score and a rejected one below it. Log-likelihoods in the DPO objective are approximated by differences in denoising-score-matching losses, following Diffusion-DPO. The offline variant draws all N structures in one batch and trains on a fixed dataset; the online variant interleaves generation, scoring, and one epoch of updates across E epochs, resetting the reference model to the current policy each epoch, which softens the KL regularisation.

Best K-of-N. Draw N samples and return the K with the highest oracle scores.

Oracles and metrics. The main oracle is TM-score against the ground-truth 1CLL backbone, with values above 0.5 indicating the same fold and below 0.2 indicating random pairs. The 9EEH experiments use a reference-free MolProbity-based validity score in [0, 1], combining Ramachandran and Cβ deviation checks, clashscore, RMS Z-scores of backbone bond lengths and angles, and peptide planarity terms. Two metrics are reported: max-of-K (best single structure) and mean-of-K (average over the returned batch).

Budget configurations: (20,2), (50,5), (100,10), (200,20), (500,50), and (1000,100). Experiments are repeated over 5 random seeds unless stated otherwise (some ablations use 3 or 10 seeds). Compute details are reported for H100 GPUs, with per-run wall-clock times ranging from minutes at low budgets up to 153.5 (± 1.6) minutes for the n1000_k100 O3 configuration in Table A.1.

Why This Matters

Impact on research. The paper fills a stated gap: guidance methods for protein structure prediction differ in how they spend oracle budget, but no systematic comparison existed. It also demonstrates that a latent-subspace optimisation framework previously used elsewhere transfers successfully to a protein structure prediction model, and it makes the trade-offs between inference-time steering, fine-tuning, and simple sampling explicit rather than anecdotal.

Real-world applications.

• Drug discovery pipelines that need reliable predicted structures before committing to expensive experimental validation.
• Antibody and therapeutic design, where downstream wet-lab screening tests batches of candidates rather than a single proposal — motivating the paper's focus on K > 1 settings.
• Selecting and prioritising candidates for molecular dynamics simulations, which can take days of GPU compute per run.
• Batched wet-lab assay screening, where the mean quality across a returned batch matters as much as the single best candidate.

Industry relevance. Oracle budget is a direct cost line for computational and experimental biology teams. The paper's practical advice — O3 at constrained budgets, FK-steering possibly competitive at moderate budgets, DPO once budgets are large enough to fine-tune — gives teams a rule of thumb for allocating that spend. It also flags a deployment consideration: DPO amortises its training budget into model weights, making large subsequent batches cheap, whereas inference-time and search methods usually discard oracle values afterwards.

Future Directions

• More protein targets and more oracle types. The authors state that future work will include additional protein targets paired with other, potentially expensive and non-differentiable, oracles.
• Controlled DPO hyperparameter studies. Training hyperparameters such as preference batch size and the number of samples per epoch can meaningfully affect performance at large N, but the paper leaves a controlled study isolating each factor for future work.
• Scaling to budgets beyond N = 1000. DPO is identified as the natural candidate at substantially larger budgets than those evaluated, an untested regime here.
• Reconciling the budget-dependent subspace dimension. The best-performing O3 dimension shifts with budget (d = 6 at N = 20, d = 10 at N = 200), raising the question of how to choose d automatically for a given budget.

Target Audience

Practitioners and researchers who need to apply guidance methods to protein structure prediction models under real cost constraints: ML engineers working on generative biology, computational structural biologists, and teams in drug discovery or protein design who must decide how to spend a limited oracle budget. It is also useful for method developers interested in latent-space optimisation and inference-time steering, and for reviewers or readers wanting a concrete, empirically grounded comparison of these families of techniques rather than a single-method claim.

Authors’ abstract

Foundation models for protein structure prediction remain unreliable on certain targets. External oracles can flag and correct these failures, but biological oracles are expensive, making oracle budget a critical constraint. Existing guidance methods, such as FK-steering, DPO, and Best K-of-N sampling, differ in how they spend this budget, yet no systematic comparison exists to guide method selection. To bridge this gap, we benchmark these methods alongside the recently proposed Optimisation Over Outputs (O3), which applies off-the-shelf optimisers within a generative model's latent subspace. We extend the usage of O3 to protein structure prediction models. Overall, our work provides the first practical reference for oracle budget-aware guidance. Our evaluation on two protein targets, calmodulin (1CLL) and E. coli aspartate transcarbamoylase (9EEH), reveals that no single method consistently dominates across all budgets and oracles. Specifically, O3 proves most effective at low oracle budgets, while FK-steering and DPO demonstrate improved performance as the budget increases. We distil these findings into actionable recommendations for practitioners operating under real-world oracle-budget constraints.