# Robust Inference-Time Steering of Protein Diffusion Models via Embedding Optimization

arXiv: https://arxiv.org/abs/2602.05285 · HTML: https://mlsanity.com/paper/2602.05285/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 14 of its 39 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 14 of 14)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Practical and Asymptotically Exact Conditional Sampling in Diffusion Models](https://mlsanity.com/paper/2306.17775/) (2023)
- [Optimizing Prompts for Text-to-Image Generation](https://mlsanity.com/paper/2212.09611/) (2022)
- [Manifold Preserving Guided Diffusion](https://mlsanity.com/paper/2311.16424/) (2023)
- [Prompt-tuning latent diffusion models for inverse problems](https://mlsanity.com/paper/2310.01110/) (2023)
- [DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models](https://mlsanity.com/paper/2509.21655/) (2025)
- [Multiscale guidance of protein structure prediction with heterogeneous cryo-EM data](https://mlsanity.com/paper/2506.04490/) (2025)
- [Unlocking hidden biomolecular conformational landscapes in diffusion models at inference time](https://mlsanity.com/paper/2512.03312/) (2025)
- [A Statistical Benchmark for Diffusion Posterior Sampling Algorithms](https://mlsanity.com/paper/2509.12821/) (2025)

## Built on by (most cited first, 6 of 6)

- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [ConforNets: Latents-Based Conformational Control in OpenFold3](https://mlsanity.com/paper/2604.18559/) (2026)
- [CrystalBoltz: End-to-End Protein Structure Determination via Experiment-Guided Diffusion for X-Ray Crystallography](https://mlsanity.com/paper/2605.15564/) (2026)
- [A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models](https://mlsanity.com/paper/2608.08770/) (2026)
- [Fold'EM: Direct atomic structure inference from Cryo-EM particles](https://mlsanity.com/paper/2610.01358/) (2026)
- [Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models](https://mlsanity.com/paper/2610.01933/) (2026)

## Abstract (arXiv)

A core challenge in structural biophysics is generating biomolecular conformations that are both physically plausible and consistent with experimental measurements. While sequence-to-structure diffusion models provide powerful priors, posterior sampling methods steer generation by perturbing atomic coordinates with gradients from experimental likelihoods. However, when the target lies in a low-density region of the prior, these methods require aggressive upweighting of the likelihood that can destabilize sampling and be sensitive to hyperparameters. We propose EmbedOpt, an inference-time steering framework that introduces an orthogonal optimization axis: rather than performing posterior sampling under a fixed prior, EmbedOpt directly optimizes the prior by updating the model's conditional embedding. This embedding space encodes rich coevolutionary signals, so optimizing it shifts the structural prior to align with experimental constraints. Empirically, EmbedOpt matches coordinate-based posterior sampling baselines on sparse distance constraints and outperforms them on cryo-electron microscopy map fitting, including real, noisy experimental ones. Furthermore, EmbedOpt's smooth optimization behavior yields robustness to hyperparameters spanning two orders of magnitude and enables comparable performance with fewer diffusion steps. Code is available at https://github.com/rs-station/embedopt.