Scallop shell SCALLOP Few-Step Boltzmann Generators via
Scalable Likelihood Flow Maps

1 University of Cambridge 2 University of Helsinki 3 Peking University 4 Carnegie Mellon University 5 Amazon FAR 6 FAIR at Meta

* Correspondence: ro352@cam.ac.uk † Equal contributions

NeurIPS 2026

Few-step molecular generation, with likelihoods along the way.
From Gaussian noise to peptide conformations: SCALLOP jointly predicts samples and their log densities, enabling importance resampling.

Abstract

Recent progress in flow-based generative modeling has led to models that output high-quality samples while using only a small number of function evaluations. However, at present, there is a lack of similar advances in estimating the model likelihood. In particular, most existing methods either rely on restrictive architectures that enable exact calculations, or use stochastic approximations such as Hutchinson’s trace estimator that introduce substantial variance.

In this work, we introduce SCAlable LikeLihood distillation of flOw maPs (SCALLOP). SCALLOP builds on the recently proposed F2D2 (Ai et al.), a likelihood flow map model that can generate samples and their densities in a small number of function evaluations. While F2D2 uses Hutchinson’s estimator during training, we introduce an alternative and more scalable likelihood distillation objective that is Hutchinson-free and admits a vectorized formulation.

Empirically, we demonstrate the effectiveness of SCALLOP as a Boltzmann generator in molecular science, and further validate its benefit on image datasets. SCALLOP significantly reduces both training variance and training time while consistently improving performance compared to F2D2, and is competitive with the state-of-the-art while achieving up to 10× inference speedup over the fastest baseline.

Method

Samples and likelihoods, learned together.

Gaussian noise

x0 ∼ 𝒩(0, I)

Likelihood flow map

Ψθ(xt, t, s)

Samples + log density

Ball-and-stick structure of alanine dipeptide

xs, log ps(xs)

Pretrain the velocity

Learn a velocity field that transports Gaussian noise toward the data distribution. Use this pretrained model as the teacher for a likelihood flow map.

Conditional Divergence Matching:
Distill without Hutchinson

Replace stochastic divergence estimation with conditional divergence matching: a vectorized objective that reduces training variance and computation.

Generate and reweight

Update the sample and its estimated log density together in a few network evaluations. Use these likelihoods for self-normalized importance sampling.

The key change is the likelihood training objective. At inference, a shared network predicts the sample update and log-density change without computing a full Jacobian.

Vectorized conditional divergence matching

View formulation

For the linear interpolation xt = (1 − t)x0 + tx1, with Gaussian noise x0 ∼ 𝒩(0, I) and a data sample x1 ∼ pdata, SCALLOP trains a vector-valued predictor Dθ ∈ ℝd using:

ℒCDM-v(θ;t)= 𝔼x0,x1 [ ‖ Dθ(xt,t,t)+ x01−t⊙ (vt(xt)−(x1−x0))− 11−t𝟏d ‖ 2 ]. \mathcal{L}_{\mathrm{CDM\text{-}v}}(\theta;t)=\mathbb{E}_{x_0,x_1}\left[\left\|D_\theta(x_t,t,t)+\frac{x_0}{1-t}\odot\bigl(v_t(x_t)-(x_1-x_0)\bigr)-\frac{1}{1-t}\mathbf{1}_d\right\|^2\right].

Here, ⊙ is elementwise multiplication, 1d is the all-ones vector, and vt is the pretrained velocity teacher. This objective replaces Hutchinson’s stochastic trace estimator with a coordinatewise training signal. The vector output targets the negative diagonal of the velocity Jacobian; summing its entries gives the scalar log-density-change prediction.

Δlogp^=(s−t) ∑j=1d Dθ,j(xt,t,s). \widehat{\Delta\log p}=(s-t)\sum_{j=1}^{d}D_{\theta,j}(x_t,t,s).

In practice, the CDM-v loss is weighted by (1 − t) and combined with flow matching and self-distillation. Training samples t from [0, 1 − ε], with ε = 10−5, to avoid the endpoint singularity.

Results

Faster likelihood evaluation. More stable training.

On alanine systems, SCALLOP improves on F2D2 across all reported metrics and achieves competitive performance with state-of-the-art Boltzmann generators.

Molecular sampling

Full results from the paper’s molecular benchmark, evaluated using self-normalized importance sampling. Values are mean ± standard deviation over three random seeds with 200,000 samples. Bold marks the best result; underlining marks the second best.

Full molecular benchmark from the paper: all five methods and all four metrics for ALA-2, ALA-3, ALA-4, and ALA-6. Best values are bold; second-best values are underlined.
SystemMetricSBG-ISFALCONFALCON-AF2D2SCALLOP
ALA-2d = 66ESS ↑0.030 ± 0.0120.067 ± 0.0130.097 ± 0.0070.066 ± 0.0010.088 ± 0.001
ℰ-W2 ↓0.873 ± 0.3380.225 ± 0.1040.512 ± 0.0380.958 ± 0.1140.555 ± 0.116
𝕋-W2 ↓0.439 ± 0.1290.402 ± 0.0210.180 ± 0.0050.383 ± 0.0220.310 ± 0.027
NFE14444
ALA-3d = 99ESS ↑0.052 ± 0.0130.077 ± 0.0040.104 ± 0.0040.079 ± 0.0010.088 ± 0.001
ℰ-W2 ↓0.758 ± 0.5060.544 ± 0.0131.385 ± 0.1822.155 ± 0.0331.596 ± 0.028
𝕋-W2 ↓0.502 ± 0.0160.452 ± 0.0110.343 ± 0.0040.219 ± 0.0010.212 ± 0.001
NFE18888
ALA-4d = 126ESS ↑0.046 ± 0.0140.055 ± 0.0030.094 ± 0.0070.053 ± 0.0010.067 ± 0.001
ℰ-W2 ↓1.068 ± 0.4950.686 ± 0.0472.929 ± 0.0684.209 ± 0.0661.545 ± 0.033
𝕋-W2 ↓0.969 ± 0.0670.858 ± 0.0771.094 ± 0.0341.435 ± 0.0231.399 ± 0.039
NFE18888
ALA-6d = 189ESS ↑0.034 ± 0.0150.060 ± 0.0170.077 ± 0.0070.033 ± 0.0010.042 ± 0.001
ℰ-W2 ↓1.021 ± 0.2390.892 ± 0.3111.211 ± 0.10510.140 ± 0.0655.845 ± 0.101
𝕋-W2 ↓1.431 ± 0.0851.256 ± 0.1321.163 ± 0.1121.143 ± 0.0391.120 ± 0.014
NFE116161616

ESS is the effective sample size (higher is better); ℰ-W2 and 𝕋-W2 are the 2-Wasserstein distances for energy and torsional-angle distributions (lower is better). NFE is the number of function evaluations. The SBG-IS normalizing-flow baseline comes from SBG (Tan et al.). SBG-IS, FALCON, and FALCON-A values are taken from FALCON (Rehman et al.), as in the paper.

Training stability and efficiency

ALA-2 training comparison: SCALLOP has lower loss-gradient variance, higher gradient signal-to-noise ratio, and converges to a higher effective sample size than F2D2.
More stable training on ALA-2. Lower gradient variance and a stronger signal-to-noise ratio lead to better effective sample size as training progresses. Select the figure to view full size.

Sampling and density evaluation are up to 10× faster than the normalizing-flow baseline from SBG and 100× faster than FALCON.

Training and inference GPU-hours across ALA-2, ALA-3, ALA-4, and ALA-6. SCALLOP and F2D2 have nearly identical inference costs, below TarFlow and FALCON.
Training and inference time. Benchmarked on one NVIDIA RTX PRO 6000. Inference includes generating 200,000 samples and their densities. Training is roughly 10% faster than F2D2 for the same number of updates. Select the figure to view full size.

Beyond molecules: CelebA-64

The same objective improves per-sample likelihood accuracy on image data. SCALLOP consistently reduces the mean absolute error in bits per dimension (BPD) relative to a 1,024-step teacher, while retaining competitive sample quality.

Full CelebA-64 benchmark from the paper: NLL, absolute deviation from reference BPD, per-sample BPD error, and FID for Flow map, F2D2, and SCALLOP. Best values are bold.
StepsFlow mapF2D2SCALLOP
NLL|Δ| ↓Err ↓FID ↓NLL|Δ| ↓Err ↓FID ↓NLL|Δ| ↓Err ↓FID ↓
1-69.8371.5871.6412.961.640.110.536.941.600.150.176.91
2-32.7234.4734.466.321.730.020.333.861.690.060.123.81
4-14.8916.6416.594.041.750.000.242.751.750.000.122.70
8-6.788.538.453.331.640.110.182.411.680.070.102.37

All models use Lagrangian self-distillation (LSD). NLL is measured in bits per dimension (BPD); |Δ| is its absolute deviation from the reference BPD; Err is the mean absolute per-sample BPD error against the teacher. The reference flow-matching model achieves BPD 1.75 with 1,024 Euler steps and FID 2.48 with 200 Euler steps. Best results are bold. On CelebA-64, SCALLOP’s loss computation takes 83.1% of F2D2’s wall-clock time.

References

BibTeX

If you find our work useful, please cite:

@inproceedings{ouyang2026fewstep,
  title = {Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps},
  author = {OuYang, RuiKang and Yu, Hanlin and Ai, Xinyue and He, Yutong and Boffi, Nicholas Matthew and Ravikumar, Pradeep and Hern{\'a}ndez-Lobato, Jos{\'e} Miguel and Simchowitz, Max and Miller, Benjamin Kurt and Chehab, Omar},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year = {2026},
  url = {https://openreview.net/forum?id=BW8H7viD5K},
}