Inference-time optimization for experiment-grounded protein ensemble generation
Abstract
Protein function relies on dynamic conformational ensembles, yet current generative models like AlphaFold3 (AF3) often fail to produce ensembles that match experimental data. Recent experiment-guided generators attempt to address this by steering the reverse diffusion process. However, these methods are limited by fixed sampling horizons and sensitivity to initialization, often yielding thermodynamically implausible results. We introduce a general inference-time optimization framework to solve these challenges. First, we optimize over latent representations to maximize ensemble log-likelihood, rather than perturbing structures post hoc. This approach eliminates dependence on diffusion length, removes initialization bias, and easily incorporates external constraints. Second, we present novel sampling schemes for drawing Boltzmann-weighted ensembles. By combining structural priors from AF3 with force-field–based priors, we sample from their product distribution while balancing experimental likelihoods. Our results show that this framework consistently outperforms state-of-the-art guidance, improving diversity, physical energy, and agreement with data in X-ray crystallography and NMR, sometimes fitting the experimental data better than deposited PDB structures. Finally, inference-time optimization experiments maximizing iPTM scores reveal that perturbing MSA embeddings can artificially inflate model confidence. This exposes a vulnerability in current design metrics, whose mitigation could offer a pathway to reduce false discovery rates in binder engineering.
Lay Summary
Proteins are not frozen sculptures - they constantly shift between several 3D shapes, and which shapes they favor is what lets them bind drugs or partner molecules. AlphaFold3 accurately predicts one shape per sequence, but it struggles to reproduce the full family of shapes and often disagrees with what lab experiments actually measure. Existing fixes "steer" the model by nudging atom positions as it generates, but this is fragile: it depends heavily on the random starting point, is tied to how many generation steps you run, and can yield physically impossible structures. Instead of pushing on the final atoms, we adjust AlphaFold3's internal representations — like tweaking the recipe rather than re-plating the dish. These tweaks carry over across many generation runs, shrug off the starting seed, and readily absorb experimental constraints. We pair this with a physics energy model so the predicted shapes are not only data-consistent but also thermodynamically realistic. On real NMR and X-ray data, our shapes fit the experiments better than steering - sometimes better than structures deposited in public databases - while staying physically sound and reproducible. We also show ipTM confidence scores can be inflated by tiny tweaks, warning against trusting them blindly in protein-binder design.