PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Inference-Time Annealing with Surrogate Likelihood Estimators

Daniel Peñaherrera, Rishal Aggarwal, David Ryan Koes

arXiv:2605.31498v3cs.LGq-bio.BM

TL;DR

Efficient molecular Boltzmann sampling is difficult, and prior inference-time annealing methods require costly score-field divergence calculations. SITA retrains flow-based models along a temperature ladder using BoltzNCE surrogate likelihoods, achieving state-of-the-art results on alanine dipeptide and alanine tripeptide while avoiding divergence terms.

  • Problem

    Inference-time annealing methods require score-field divergence calculations for importance weights, making them difficult to scale to larger molecular systems.

  • Method

    SITA alternates flow-model retraining and temperature annealing while using a BoltzNCE energy-based model to provide surrogate likelihoods for importance weighting.

  • Results

    SITA achieves state-of-the-art results on alanine dipeptide and alanine tripeptide while avoiding costly divergence terms.

  • Takeaways & Limitations

    Surrogate likelihood estimators offer a tractable route to modeling molecular ensembles with many degrees of freedom, a regime where existing methods struggle.

  • Takeaways & Limitations

    Surrogate likelihood bias can break exact convergence guarantees, and the recovered distribution is exact only when the flow and surrogate densities match.

Abstract

from arXiv · show

A long standing challenge in computational chemistry and biophysics is efficiently sampling the Boltzmann distribution of molecules. Advances in generative modeling have been proposed to address the limitations of conventional sampling techniques by eliminating the computational cost of simulation. A promising direction is iteratively finetuning diffusion models along a temperature ladder whereby training data is generated via importance sampling during inference-time annealing. Unfortunately, these methods require computing a divergence over the score field to estimate importance weights, rendering them intractable for larger systems. Here we present scalable inference-time annealing (SITA), which retrains flow-based models to generate samples at progressively lower temperatures using an energy-based model to facilitate fast surrogate likelihoods. We demonstrate state-of-the-art performance on both Alanine Dipeptide and Alanine Tripeptide while avoiding costly divergence terms. Our code is available at https://github.com/countrsignal/sita.git

1 Introduction

Efficient Boltzmann sampling remains difficult because molecular energy landscapes trap conventional simulations and generative models need equilibrium data to train. SITA addresses this circularity with temperature-annealed flow models, surrogate likelihoods, and importance-weighted resampling.

  • Motivation: Boltzmann sampling provides access to thermodynamic observables but is difficult for molecular systems with rugged, non-convex energy functions.MCMC and molecular dynamics can become trapped in energy minima and require long simulations with highly correlated samples.
  • Motivation: Generative samplers promise fast amortized Boltzmann sampling, yet their training requires equilibrium ensembles that are themselves difficult to generate.Bootstrapping methods address this circularity by retraining models on outputs from earlier optimization steps.
  • Existing Direction: Temperature annealing trains models on high-temperature data and iteratively fine-tunes them at lower temperatures, improving exploration and data efficiency.The temperature sequence can be used to bootstrap recovery of the target distribution.
  • Existing Direction: PITA-style diffusion annealing requires divergence evaluation along the reverse-process integration path, creating serious computational overhead.This requirement limits inference-time importance weighting for larger systems.
  • SITA: SITA introduces flow-based inference-time annealing with BoltzNCE surrogate likelihoods and importance-weighted transport from high to lower temperatures.The training loop uses flow proposals, an energy-based model, and resampling to seed subsequent annealing steps without expensive Jacobian computations.
  • SITA: SITA achieves state-of-the-art results across several metrics despite bias from its surrogate likelihood estimator.

2 Background

SITA combines continuous-time flow modeling with BoltzNCE, an energy-based surrogate likelihood estimator. The surrogate is trained on flow-generated samples using score matching and noise contrastive estimation.

  • Continuous-Time Generative Models: Stochastic interpolants provide a continuous-time generative framework connecting a base distribution to a target distribution.The framework encompasses diffusion models, flow matching, and rectified flows under suitable interpolation functions.
  • Continuous-Time Generative Models: The probability-flow velocity field is learned by neural-network regression, after which generated samples are obtained by integrating the learned dynamics.Their density is represented through the change-of-variables formula.
  • Surrogate Likelihood Estimators: SITA uses an energy-based model as a surrogate likelihood estimator for efficient reweighting of flow-generated samples.BoltzNCE makes EBM training practical by combining score matching with noise contrastive estimation.
  • Surrogate Likelihood Estimators: BoltzNCE first draws samples from the flow and then fits an energy function to the flow’s output distribution.
  • Surrogate Likelihood Estimators: The full BoltzNCE objective combines score matching, which constrains energy gradients, with an NCE term that anchors energy values.The contrastive term discriminates between samples associated with different interpolation times.

3 Scalable Inference-Time Annealing

SITA anneals a flow model through a decreasing temperature ladder by alternating EBM refinement, surrogate-based importance resampling, and flow fine-tuning. It enables large temperature jumps without divergence computations, but surrogate mismatch biases the recovered distribution.

  • Annealing Procedure: SITA starts with a flow trained on high-temperature simulation data and an EBM trained on flow-generated outputs.The process uses a decreasing temperature sequence and an exactly evaluable target energy function.
  • Annealing Procedure: At each annealing step, the flow generates samples, the EBM is fine-tuned, and surrogate-based importance weights create data for the next lower temperature.The flow is then fine-tuned on the resampled data until the final temperature is reached.
  • Annealing Procedure: The same flow and EBM are maintained throughout annealing, with their optimizers re-initialized at each fine-tuning step.
  • Temperature Steering: Temperature steering raises the flow density to κ = T_high/T_low and rescales the base-distribution variance at inference time.This enables large temperature jumps without changing the velocity-field architecture or enforcing volume preservation.
  • Surrogate-Likelihood Bias: The surrogate EBM introduces bias into importance weights, so the recovered distribution is generally a tilted target and is exact only when ρ_θ = q_ϕ.Despite this bias, the EBM-derived weights deliver superior benchmark performance empirically.

4 Related Work

Related work spans flow-based, diffusion-based, and importance-sampling-based molecular samplers. SITA builds on continuous-time generative models while targeting temperature-annealed flow bootstrapping.

  • Flow-Based Boltzmann Samplers: Flow-based Boltzmann samplers use bijections from tractable base distributions, but mode coverage remains a major empirical constraint.
  • Diffusion-Based Sampling: Diffusion-based samplers differ in how they estimate optimal drifts and balance fast mode mixing against simulation costs during training.
  • Importance-Sampling Bootstrapping: Importance-sampling bootstrapping corrects flow outputs along annealing paths inspired by Annealed Importance Sampling and Sequential Monte Carlo.Prior approaches use geometric paths, constrained mass transport, or temperature schedules.
  • Continuous-Time Generative Models: SITA uses continuous-time generative models and stochastic interpolants, which include flow matching and rectified flows as special cases.

5 Experiments

SITA is evaluated against established baselines on alanine dipeptide and tripeptide, using distributional, energy, torsion, efficiency, and TICA-based measures. It achieves strong benchmark performance, captures major conformational basins, and supports post-hoc refinement with distinct diversity–accuracy trade-offs.

  • Experimental setup: SITA is evaluated on alanine dipeptide and alanine tripeptide against PITA, TA-BG, MD-Diff, and MD-NF.Evaluation uses 300K MD test data, 10,000 generated samples, three seeds, Wasserstein distances, Rama-KL, torsion metrics, and energy-evaluation costs.
  • Efficiency: Bootstrapping requires one to two orders of magnitude fewer energy evaluations after the shared training-set cost.The shared upfront cost is 5 × 10^7 evaluations for each system.
  • Alanine Dipeptide: SITA captures all major alanine-dipeptide conformational basins and closely matches the MD reference in energy distribution and Ramachandran free-energy landscape.The comparison uses 30,000 samples and an annealing schedule ending at 300K after pretraining at 1200K.
  • Alanine Dipeptide: On alanine dipeptide, SITA achieves the best Rama-KL and energy Wasserstein-2 performance, while outperforming PITA on energy Wasserstein-1.SITA remains competitive with PITA on the torsion metric, where recovering mode populations is challenging.
  • Alanine Tripeptide: On alanine tripeptide, SITA outperforms all methods on nearly all metrics except T-W2, without requiring relaxation of generated samples.PITA and TA-BG require short MD refinement at the target temperature to remain competitive.
  • Metropolis-Hastings Refinement: SITA-IMH improves every reported metric while preserving sample diversity, whereas SITA-IS lowers energy distances but degrades Rama-KL with effective sample size 0.191.SITA-IMH attains the best Rama-KL and T-W2 scores at the cost of additional energy evaluations.
  • TICA Evaluations: TICA evaluation depends materially on lag time and MD-frame down-sampling, with uniform sampling improving the reported TICA metrics.The adopted protocol fits TICA on the full trajectory, uniformly subsamples projection frames, and uses lag times of 100 for ADP and 10 for ATP.

6 Conclusion

SITA uses temperature annealing with continuous flow models and surrogate likelihood estimators to recover Boltzmann distributions efficiently. It avoids vector-field divergence computations and reports state-of-the-art results on alanine dipeptide and alanine tripeptide.

  • SITA trains continuous flow models to recover target Boltzmann distributions through temperature annealing.
  • The method avoids vector-field divergence computations, reducing the computational cost of likelihood evaluation.
  • SITA establishes state-of-the-art performance on alanine dipeptide and alanine tripeptide while offering a tractable route to molecular ensembles with many degrees of freedom.
  • Continuous normalizing flows transport a base distribution through a time-indexed velocity field whose density satisfies the continuity equation.
  • SITA uses temperature-scaled flow and base-distribution relations to model samples across progressively lower temperatures.

Appendix B. Independent Metropolis-Hastings with Surrogate Likelihoods

The surrogate-likelihood IMH construction replaces an unavailable exact proposal density with a learned surrogate. This changes the stationary distribution unless the surrogate matches the proposal, so efficiency and correctness depend on that discrepancy.

  • Independent Metropolis-Hastings proposes candidates independently of the chain’s current state and accepts them using the target density and proposal density.
  • Replacing the exact proposal density with a surrogate q yields a modified acceptance ratio while candidates remain drawn from the true proposal p.
  • The surrogate chain has tilted stationary distribution ˜π(x) ∝ π(x)p(x)/q(x), rather than generally targeting π.
  • The deviation from the target is governed by the discrepancy between p and q, and standard IMH is recovered when q ≡ p.
  • The acceptance-probability derivation uses detailed balance and compares forward and reverse transition kernels through the surrogate importance ratio.

Appendix C. Architecture Details

SITA represents molecular conformers as fully connected graphs and processes Cartesian coordinates with an equivariant flow and a structure-aware energy-based model. The architecture injects geometric information through distances and graph topology.

  • Molecular conformers are encoded as fully connected graphs, with the flow and EBM operating directly on Cartesian atomic coordinates.
  • The flow uses an E(3)-equivariant graph neural network with geometric vector perceptrons.
  • Message passing combines scalar node features, vector messages, pairwise distances, and atomic positions.
  • The EBM uses a Graphormer whose structure-aware attention bias is derived from the Euclidean distance matrix through an MLP.
  • All training experiments ran on a single NVIDIA L40 GPU with 48 GB GDDR6 memory.

D.2 Pre-training

The pre-training setup trains the flow and energy-based model separately before annealing bootstrap modifications are applied. The appendix identifies the bootstrap procedure as a later fine-tuning stage.

  • Pre-training: The flow and EBM were each pre-trained for 500 epochs with batch size 512.
  • Pre-training: Separate Adam optimizers used learning rates of 1e−3 for the flow and 5e−4 for the EBM.
  • Pre-training: Both models used reduce-on-plateau schedulers with patience 30, reduction factor 0.5, and exponential moving averages with decay 0.999.
  • Bootstrap modifications: Annealing-bootstrap training was nearly identical to pre-training, with its modifications summarized in Table 7.

Appendix E. Metrics

The appendix evaluates generated samples using optimal-transport distances, energy and torsional costs, Ramachandran KL divergence, and TICA-based comparisons. These metrics compare model and reference distributions in physically relevant representations, including periodic dihedral geometry and slow dynamical modes.

  • Effective sample size is reported using normalized importance weights.
  • Optimal-transport evaluation: The 2-Wasserstein distance measures minimum-cost couplings between generated and reference sample point clouds.Optimal couplings are obtained with the POT library.
  • Optimal-transport evaluation: Energy-space agreement is assessed with an energetic cost sensitive to local bonded geometry and long-range interactions.This produces the E-W2 metric.
  • Optimal-transport evaluation: For alanine dipeptide, a periodic torsional cost compares Ramachandran angles ϕ and ψ on their torus.The angles fully characterize the molecule’s backbone geometry in this setting.
  • Distributional metrics: Ramachandran KL divergence approximates the discrepancy between reference and generated distributions using 2D histograms on [−π, π]2.
  • TICA metrics: TICA identifies projections maximizing autocorrelation at a chosen lag, with eigenvalues near one indicating slowly persistent dynamics.The resulting basis is applied to both MD and generated features before computing Wasserstein distances.
  • TICA metrics: Wasserstein-p distances use identical molecular featurization and discrete optimal transport with Euclidean ground costs.For W1 the transport cost is unsquared; for W2 squared-Euclidean cost is used and the optimum’s square root is returned.

F.1 Sources of TICA Evaluation Error

The appendix identifies lag-time selection and trajectory down-sampling as sources of error in TICA evaluation. These choices can change the projections and bias reference distributions, affecting comparisons between generated and MD samples.

  • PITA’s TICA evaluations contain errors from lag-time choice and down-sampling trajectory frames.
  • Choice of lag time: Different lag times define different TICA optimization problems because each lag targets autocorrelation at a single timescale.Small lags can mix fast molecular motions with thermal noise and integration artifacts, whereas larger lags filter these effects but discard some fast processes.
  • Down-sampling: Using the first 10,000 temporally correlated MD frames as a reference introduces systematic bias and can drop modes.This can make mode-collapsed models appear better when their missing modes match the biased reference.
  • Down-sampling: Figure 4 compares TICA down-sampling at different lag times using projections of MD-derived samples.
  • Down-sampling: Table 8 reports alanine-tripeptide TICA-W1 and TICA-W2 at τ = 100, contrasting uniform frame sampling with first-frame down-sampling across 10,000 samples and 3 seeds.
Loading 2605.31498v3…