Inference-time optimization for experiment-grounded protein ensemble generation

Sai Advaith Maddipatla, Anar Rzayev, Marco Pegoraro, Martin Pacesa, Paul Schanda, Ailie Marx, Sanketh Vedula, Alexander Bronstein
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:85483-85523, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-maddipatla26a, title = {Inference-time optimization for experiment-grounded protein ensemble generation}, author = {Maddipatla, Sai Advaith and Rzayev, Anar and Pegoraro, Marco and Pacesa, Martin and Schanda, Paul and Marx, Ailie and Vedula, Sanketh and Bronstein, Alexander}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {85483--85523}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/maddipatla26a/maddipatla26a.pdf}, url = {https://proceedings.mlr.press/v306/maddipatla26a.html}, 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.} }
Endnote
%0 Conference Paper %T Inference-time optimization for experiment-grounded protein ensemble generation %A Sai Advaith Maddipatla %A Anar Rzayev %A Marco Pegoraro %A Martin Pacesa %A Paul Schanda %A Ailie Marx %A Sanketh Vedula %A Alexander Bronstein %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-maddipatla26a %I PMLR %P 85483--85523 %U https://proceedings.mlr.press/v306/maddipatla26a.html %V 306 %X 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.
APA
Maddipatla, S.A., Rzayev, A., Pegoraro, M., Pacesa, M., Schanda, P., Marx, A., Vedula, S. & Bronstein, A.. (2026). Inference-time optimization for experiment-grounded protein ensemble generation. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:85483-85523 Available from https://proceedings.mlr.press/v306/maddipatla26a.html.

Related Material