Scalable Inference-Time Annealing with Surrogate Likelihood Estimators

This paper introduces Scalable Inference-Time Annealing (SITA), a method that overcomes the intractability of existing inference-time annealing techniques for large molecular systems by utilizing energy-based models to provide fast surrogate likelihoods, thereby achieving state-of-the-art sampling performance on peptides without costly divergence calculations.

The Gist

Imagine you are trying to find the most comfortable spot to sleep in a giant, dark, mountainous cave. This cave represents a molecule, and the "comfortable spots" (the valleys) represent the most stable, natural shapes the molecule can take. The goal of scientists is to map out all these valleys to understand how the molecule behaves.

The problem is that the cave is huge, and the paths are tricky. If you just start walking randomly, you might get stuck in one small valley and never find the others. If you try to map the whole cave by walking every single inch, it would take longer than the age of the universe.

This paper introduces a new, clever way to map the cave called SITA (Scalable Inference-Time Annealing). Here is how it works, broken down into simple steps:

1. The Old Way: The "Perfect Map" Trap

Traditionally, scientists tried to use computer models (like a GPS) to learn the map of the cave. But to train the GPS, you needed a perfect map of the cave to begin with. This is a catch-22: you need the map to make the map.

Another method involved "simulating" the walk step-by-step (like a very slow, careful hiker). While accurate, it's incredibly slow and expensive, like trying to map a continent by walking every single step.

2. The New Idea: The "Hot Start" Strategy

The authors realized they could cheat the system by starting in a different place. Imagine heating up the cave until the walls melt and the ground becomes flat and smooth (a high temperature). In this "hot" state, it's easy to run around and explore the whole cave quickly because there are no deep valleys to get stuck in.

They trained their AI model on this "hot, flat" version of the cave. Now, they have a model that knows how to run around freely.

3. The Problem with "Cooling Down"

The goal is to find the comfortable spots in the cold cave (room temperature), where the valleys are deep and sharp. If you just tell the "hot" model to slow down and look for valleys, it often gets confused. It might miss some valleys or get stuck in the wrong ones.

Previous attempts to fix this involved a very expensive calculation (like checking a massive, complex ledger for every single step) to ensure the model didn't make mistakes. This was too slow for big molecules.

4. The SITA Solution: The "Surrogate Guide"

This is where SITA comes in. Instead of doing the expensive ledger check, the authors use a Surrogate Likelihood Estimator. Think of this as a smart, cheap guide dog.

Here is the process:

1. The Hot Run: The AI model (the runner) generates a bunch of random paths in the "hot" cave.
2. The Guide Learns: A second, smaller AI (the guide dog) looks at these paths and learns to guess which ones are good and which are bad. It doesn't need to be perfect; it just needs to be a good "surrogate" (a stand-in) for the expensive calculation.
3. The Filter: The guide dog helps sort the runner's paths. It says, "Hey, this path looks like a good valley, keep it. That one looks like a dead end, throw it away."
4. The Cool Down: The runner then tries again, but this time it uses the guide dog's advice to focus on the "cooler," more stable parts of the cave.
5. Repeat: They do this over and over, slowly lowering the temperature (making the cave colder and the valleys deeper), with the guide dog getting better at its job each time.

5. Why It's a Big Deal

• Speed: By using the "guide dog" (the surrogate) instead of the "expensive ledger," they can handle much larger and more complex molecules without the computer crashing or taking forever.
• Accuracy: They tested this on two small protein molecules (Alanine Dipeptide and Tripeptide). The results showed that SITA found all the important "valleys" (stable shapes) better than previous methods, even though it used a shortcut.
• No "Mode Collapse": Sometimes, AI models get lazy and only find one valley, ignoring the others. SITA managed to find all the major valleys, not just the easiest one.

Summary

In short, the authors built a system that learns to explore a complex molecular landscape by starting in a "hot," easy-to-navigate version of the world. They use a smart, lightweight "guide" to help the system slowly cool down and find the precise, stable shapes of molecules, avoiding the need for slow, expensive calculations that used to make this impossible for large systems.

What the paper does NOT claim:

• It does not claim to cure diseases or be used in hospitals yet.
• It does not claim to work on any molecule instantly; it was tested specifically on small protein chains (Alanine Dipeptide and Tripeptide).
• It does not claim the method is perfect; it admits there is a small "bias" (a slight guesswork element) introduced by the guide dog, but the results show this bias is acceptable for getting high-quality answers.

Technical Summary

Technical Summary: Scalable Inference-Time Annealing with Surrogate Likelihood Estimators (SITA)

Problem Statement

Sampling the equilibrium ensemble of molecular configurations, defined by the Boltzmann distribution $\pi(x) \propto \exp(-E(x)/k_B T)$, is a foundational yet computationally intractable task in statistical physics and computational chemistry. Traditional methods like Molecular Dynamics (MD) and Markov Chain Monte Carlo (MCMC) suffer from high computational costs due to femtosecond timesteps and the tendency to become trapped in local energy minima.

While deep generative models offer a path toward amortized, fast sampling, they face a "circular" training limitation: they require equilibrium ensembles for training data, which are precisely the difficult-to-generate target. Recent "bootstrapping" approaches attempt to resolve this by iteratively retraining models on their own outputs via temperature annealing. However, existing methods, particularly those based on diffusion models (e.g., PITA), rely on self-normalized importance sampling (SNIS) that requires computing the divergence of the score field along the full reverse-time integration path. This divergence computation scales prohibitively with system dimensionality, rendering these methods intractable for larger molecular systems.

Methodology: SITA

The authors propose SITA (Scalable Inference-Time Annealing), a framework that combines continuous flow-matching models with a surrogate likelihood estimator to bypass the computational bottleneck of divergence calculations.

Core Components

1. Continuous Flow-Matching Models: SITA utilizes stochastic interpolants to define a continuous-time process transporting samples from a base Gaussian distribution to a target distribution. Unlike diffusion models, flow-matching learns a velocity field $v_t(x)$ via regression, avoiding the need for explicit noise schedules.
2. Surrogate Likelihood Estimators (BoltzNCE): To avoid calculating the intractable divergence term ($\nabla \cdot v_t$) required for exact likelihood evaluation in flow-based reweighting, SITA employs an Energy-Based Model (EBM) parameterized as a surrogate likelihood. This EBM is trained using BoltzNCE, a method combining score matching with Noise Contrastive Estimation (NCE). The EBM learns an energy function $U_\phi$ over the flow's output distribution, providing a tractable density estimate $q_\phi(x)$.
3. Temperature Steering: SITA leverages the implicit temperature encoding in flow models. By rescaling the variance of the base Gaussian distribution at inference time ($x_0 \sim \mathcal{N}(0, \kappa^{-1}I)$ where $\kappa = T_{high}/T_{low}$), the flow can generate samples at lower temperatures without architectural changes or explicit temperature conditioning.

The SITA Algorithm

The method operates via an iterative bootstrap loop over a decreasing temperature ladder $\{T_k\}$:

1. Anneal the Flow: Generate samples at temperature $T_{k+1}$ by integrating the flow ODE with a rescaled base distribution.
2. Finetune the EBM: Use the generated flow samples to update the EBM parameters $\phi$ via the BoltzNCE objective, refining the surrogate likelihood $q_\phi$.
3. Importance Sampling: Compute importance weights $w(x) \propto \exp(-E(x)/k_B T_{k+1}) / q_\phi(x)$. Resample the flow outputs using these weights to create a dataset targeting the lower temperature $T_{k+1}$.
4. Finetune the Flow: Retrain the flow model $\theta$ on the reweighted dataset to better approximate the target distribution at $T_{k+1}$.

This process repeats until the target temperature (e.g., 300K) is reached. The method avoids expensive Jacobian computations by relying on the EBM's surrogate likelihood for the reweighting step.

Key Contributions

• First Inference-Time Annealing for Flows: SITA is the first method to apply inference-time annealing to continuous flow-matching models, enabling large temperature jumps across a pre-defined ladder.
• Surrogate-Driven Importance Sampling: The integration of BoltzNCE-style surrogates allows for importance-weighted transport from high-temperature distributions to room-temperature targets, circumventing the prohibitive cost of path-integral-based SNIS estimators used in diffusion-based methods.
• Scalability: By replacing divergence calculations with a tractable EBM, the method scales to systems with many degrees of freedom where previous annealing approaches fail.

Experimental Results

The authors evaluated SITA on two benchmark molecular systems: Alanine Dipeptide (ADP) and Alanine Tripeptide (ATP), comparing performance against PITA, Temperature Annealed Boltzmann Generators (TA-BG), and models trained directly on 300K MD data (MD-Diff, MD-NF).

• Alanine Dipeptide: SITA achieved state-of-the-art performance in Ramachandran KL divergence (Rama-KL) and Energy Wasserstein-2 metrics, outperforming PITA. While MD-NF showed lower Energy-W1, it exhibited signs of mode collapse (poor Rama-KL), whereas SITA successfully captured all major conformational basins.
• Alanine Tripeptide: SITA outperformed all baselines on nearly all metrics (Rama-KL, Energy-W1, Energy-W2) without requiring post-hoc MD refinement. In contrast, PITA and TA-BG required short MD refinement steps to remain competitive.
• Efficiency: After the initial shared cost of generating training data, the bootstrapping phase incurred one to two orders of magnitude fewer energy evaluations compared to baseline methods.
• Refinement: The authors demonstrated that applying Independent Metropolis-Hastings (IMH) using the learned surrogate further improved metrics (particularly Rama-KL and T-W2) while preserving sample diversity, though at the cost of additional energy evaluations.

Significance and Claims

The paper claims that SITA establishes a new state-of-the-art for sampling Boltzmann distributions in molecular systems by resolving the computational intractability of divergence terms in flow-based annealing. The authors emphasize that surrogate likelihood estimators offer a practical route to modeling molecular ensembles with many degrees of freedom, a regime where existing diffusion-based bootstrapping methods struggle with mode collapse or computational cost.

The work highlights that while the use of a surrogate likelihood introduces a theoretical bias (converging to a tilted distribution rather than the exact target), empirical results on benchmark systems show that this bias does not preclude superior performance in capturing equilibrium ensembles. The authors position SITA as a scalable alternative that maintains the architectural advantages of continuous-time flows while enabling efficient, iterative refinement toward low-temperature targets.