Autoregressive Boltzmann Generators

Abstract: Efficient sampling of molecular systems at thermodynamic equilibrium is a hallmark challenge in statistical physics. This challenge has driven the development of Boltzmann Generators (BGs), which allow rapid generation of uncorrelated equilibrium samples by combining a generative model with exact likelihoods and an importance sampling correction. However, modern BGs predominantly rely on normalizing flows (NFs), which either suffer from limited expressivity due to strict invertibility constraints (discrete time) or computationally expensive likelihoods (continuous time). In this paper, we propose Autoregressive Boltzmann Generators (ArBG) -- a novel autoregressive modelling framework -- that overcomes these limitations by departing from the flow-based BG paradigm. ArBG circumvents the topological constraints of flows and enables sequential inference-time interventions, while offering enhanced scalability by leveraging architectures effective in LLMs. We empirically demonstrate that ArBG leads to significant improvements over flow-based models across all benchmarks, but particularly in larger peptide systems such as the 10-residue Chignolin. Furthermore, we introduce Robin, a 132 million parameter transferable model trained with the ArBG framework which improves over the previous state-of-the-art, reducing the zero-shot energy error, E-W$_2$, on 8-residue systems by over 60$\%$. The code can be found at the following link: https://github.com/danyalrehman/autobg.

• The paper introduces ARBG that removes diffeomorphism constraints to efficiently model equilibrium distributions in high-dimensional molecular systems.
• It employs autoregressive factorization with uniform binning to achieve robust, direct likelihood evaluation and improved mode coverage.
• Empirical results demonstrate ARBG outperforms flow-based methods on energy and geometric metrics, enabling scalable zero-shot sampling via ROBIN.

Autoregressive Boltzmann Generators: A Diffeomorphism-Free Framework for Equilibrium Molecular Sampling

Introduction

The sampling of equilibrium distributions in high-dimensional molecular systems remains a critical challenge in statistical mechanics, primarily due to the sparsity and separation of metastable states in conformational landscapes. Existing approaches for equilibrium sampling, such as Molecular Dynamics (MD) and Monte Carlo methods, provide limited efficiency due to issues with mixing times and trapping within energy basins. The Boltzmann Generator (BG) framework—learning generative models with tractable likelihoods for importance resampling—has recently emerged as a scalable alternative. Notably, BGs have been dominated by normalizing flow (NF) architectures, each bearing significant theoretical and practical limitations: discrete NFs suffer from topological and expressivity barriers, while continuous NFs (CNFs) incur high computational costs during inference due to stiff ODE dynamics.

The paper "Autoregressive Boltzmann Generators" (2606.27361) introduces a principled departure from the flow-based paradigm by leveraging autoregressive (AR) modelling for equilibrium sampling. The proposed Autoregressive Boltzmann Generator (ARBG) framework circumvents diffeomorphic constraints, affords precise conditional likelihoods, and integrates well with advances in large-scale sequence modelling architectures. The authors provide a comprehensive study of proposal formulations, demonstrate the theoretical properties of ARBG, and empirically validate its strong performance and scalability across molecular benchmarks, including transferable models for zero-shot generalization.

Background and Limitations of Flow-Based Boltzmann Generators

Flow-based BGs require the generative map to be a diffeomorphism—an invertible, smooth, and topology-preserving mapping—between the latent and molecular configuration spaces. This necessity poses a severe expressivity bottleneck: the Gaussian latent prior must be warped through drastic, often ill-conditioned, transformations to capture disconnected or multiply-connected regions in the molecular Boltzmann distribution. Discrete NFs experience particularly poor conditioning with exploding Jacobians and limited mode coverage, whereas CNFs, though more expressive, are restrained by computationally intensive ODE solvers and high variance in numerical divergence estimation.

The unavoidable tradeoff between expressivity and computational tractability in flow-based BGs, especially for systems with complex multimodal manifolds (e.g., peptides with multiple stable conformations), motivates the exploration of fundamentally different generative mechanisms.

The ARBG Framework

Architecture and Density Factorization

ARBG models the molecular density $p_\theta(x)$ as a product of conditional densities over a specified atom/residue ordering:

$$p_\theta(x) = \prod_{j=1}^{d} p_\theta(x_j \mid x_{<j})$$

This autoregressive factorization supports efficient and unbiased evaluation of the likelihood without recourse to Jacobian determinants or ODE integration. It enables direct optimization via maximum likelihood estimation and compatibility with self-normalized importance sampling (SNIS) for uncorrelated equilibrium sampling.

Proposal Distributions and Parameterization

A key technical contribution is the examination of proposal parameterizations for the conditionals $p(x_j \mid x_{<j})$ in high-dimensional continuous spaces. The paper analyzes mixture density networks (MDN) with both Mixture of Logistics (MoL-PixelCNN++) and Gaussian mixture (GMM-PixelCNN++) outputs. However, these approaches exhibit training instability and mode collapse phenomena.

The paper then proposes a simple uniform binning discretization, converting the conditional prediction task to categoricals over discretized bins, paralleling next-token prediction in LLMs. This enables robust training dynamics, predictable scaling, and practical dequantization for continuous-valued inference. The minimum achievable error, quantified via a KL-entropy-based bound, is shown to be determined solely by bin width.

Inference-Time Interventions: Autoregressive Twisted SMC

ARBG uniquely admits sequential inference-time interventions, leveraging the AR factorization for early rejection and correction via Twisted Sequential Monte Carlo (SMC). By conditioning resampling steps on physically plausible partial structures (e.g., using residue-level energy-based filters), it enables computational efficiency not available to standard BGs, which must wait for full candidate completion before importance resampling.

Empirical Evaluation

Benchmarks and Metrics

The paper evaluates ARBG on equilibrium sampling for peptides of increasing sizes, including Chignolin (10 residues), using energy-based and geometric metrics:

• 2-Wasserstein Energy Distance (E-W2): Alignment of sample and reference energy distributions
• Torus 2-Wasserstein Distance (T-W2): Coverage in torsional (dihedral) space
• TICA 2-Wasserstein (TICA-W2): Alignment in collective slow modes from time-lagged independent component analysis

Flow-based baselines include discrete NFs, ECNF++, SBG, GIVT, as well as MDN-based AR models.

Single-System Performance and Scaling

ARBG demonstrates consistent dominance over all flow-based models across E-W2 and T-W2 for all tested systems. Notably, for Chignolin, ARBG achieves lower energy and geometric discrepancies—a regime where discrete NFs collapse due to expressivity limits and CNFs become intractably expensive. ARBG's training and inference time scale favorably with model size, paralleling LLM-like scaling laws, and model performance improves monotonically with bin resolution.

Transferable Autoregressive Boltzmann Generation: ROBIN

The authors introduce ROBIN, a 132M parameter transformer trained via ARBG, conditioned to enable zero-shot equilibrium sampling for unseen peptide sequences. On 8-residue peptide benchmarks, ROBIN reduces E-W2 error by over 60% compared to the best previous transferable BG, Prose, and achieves favorable scaling in inference budget (order-of-magnitude fewer samples for the same quality).

The twisted SMC refinement offers marginal gains given the high base quality of ROBIN but is demonstrated to be critical for larger/less well-learned systems.

Ablations

Performance is systematically analyzed over bin resolutions and temperature sweeps. Mode coverage and energy fidelity improve with higher bin counts (finer discretization), and an optimal generation temperature near unity (with slight system-dependent shifts) is established for best tradeoff between diversity and energy accuracy.

Theoretical and Practical Implications

ARBG breaks the topological bottleneck imposed by invertible, topology-preserving flows, thus capturing disconnected and highly multimodal Boltzmann densities with greater flexibility. The tractable, exact likelihood construction overcomes the computational infeasibility of continuous flows at inference. ARBG's sequential autoregressive generation unlocks incremental physical validation and resampling, essential for scalable and safe generation in larger, more complex molecular systems.

Practically, ARBG sets new state-of-the-art results for scalable peptide equilibrium sampling, both on single-molecule and zero-shot transfer settings. The framework aligns with high-throughput, distributed generative workflows, supporting token-level steering, diversity-rigidity balancing (via temperature), and compatibility with downstream molecular simulation protocols.

Limitations and Future Directions

A primary limitation of ARBG is the necessity to impose an explicit ordering over molecular coordinates, which may not align with intrinsic graph symmetries; performance may be sensitive to this choice, particularly for small or undirected graphs. The piecewise-uniform binning bounds model precision, which may become restrictive for systems with acutely peaked energy profiles or when ultra-precise local sampling is required. Unlike flows, ARBG does not yet exploit data-informed priors or explicit physical symmetries (e.g., SE(3)-equivariance), an area open for architectural innovation. Investigations into hybrid approaches, informative priors, and the exploitation of symmetry-aware AR architectures may further ameliorate these gaps.

Conclusion

Autoregressive Boltzmann Generators represent a significant advance in the generative modelling of equilibrium molecular systems. By decoupling BG performance from flow-based architectural restrictions, ARBG combines the computational stability and efficiency of autoregressive likelihoods with the expressivity required for challenging multimodal distributions. The empirical dominance of ARBG and the transferable successes of ROBIN foreshadow broader applicability in molecular design, computational chemistry, and other scientific domains requiring rigorous, scalable sampling from high-dimensional, structurally complex physical distributions (2606.27361).