Autoregressive Boltzmann Generators

Noosaga Placements

• Deep Learningsubfield95%
  The paper proposes an autoregressive model leveraging transformer architectures and scaling properties similar to large language models, which is a core deep learning contribution.
  ArBG circumvents the topological constraints of flows and enables sequential inference-time interventions, while offering enhanced scalability by leveraging architectures effective in Large Language Models.ArBG benefits from the investment and advances in discrete generative modelling that power modern LLMs, and exhibits similar scaling properties in both model size and inference samples.
• Deep Learningframework95%
  The paper proposes an autoregressive model using transformer architectures, which is a deep learning framework, and leverages scaling properties similar to LLMs.
  ArBG circumvents the topological constraints of flows and enables sequential inference-time interventions, while offering enhanced scalability by leveraging architectures effective in Large Language Models.ArBG benefits from the investment and advances in discrete generative modelling that power modern LLMs, and exhibits similar scaling properties in both model size and inference samples.
• Statistical Learningframework90%
  The paper uses importance sampling and self-normalized importance sampling, which are core techniques in statistical learning, to correct the proposal distribution.
  A commonly-used approach to obtain consistent samples from the Boltzmann distribution leverages self-normalized importance sampling (SNIS) with an easy-to-sample proposal distribution p(x).A crucial requirement of any proposal in a BG is to facilitate SNIS-based correction of samples—necessitating access to fast, unbiased, and preferably exact likelihood evaluation pθ(x).
• Machine Learningsubfield80%
  The work involves generative modeling, likelihood-based training, and importance sampling, which are core machine learning techniques, but the primary focus is on deep learning architectures.
  Boltzmann Generators (BGs) have emerged as a powerful framework to circumvent this. BGs learn a generative model pθ(x) to propose samples for importance sampling, leveraging the exact model likelihood pθ(x) and target energy E(x).We train our models using maximum likelihood estimation (MLE) on the Boltzmann distribution.
• Transformer Architectureframework90%
  The autoregressive model is based on the transformer architecture, which is explicitly mentioned and used for the conditional density modeling.
  The architecture choice employed follows a standard but performant recipe of Transformer-based building blocks.We employ FlashAttention for all models.
• Probabilistic Aisubfield70%
  The paper deals with sampling from a target Boltzmann distribution using importance sampling and sequential Monte Carlo, which are probabilistic methods, but the main innovation is in the autoregressive architecture.
  Efficient sampling of molecular systems at thermodynamic equilibrium is a hallmark challenge in statistical physics.We implement this via an Autoregressive Twisted SMC.
• Deep Generative Modelsframework85%
  The paper proposes a generative model for sampling from the Boltzmann distribution, which is a deep generative model, and uses autoregressive modeling as the generative approach.
  Boltzmann Generators (BGs) have emerged as a powerful framework to circumvent this. BGs learn a generative model pθ(x) to propose samples for importance sampling.We introduce AUTOREGRESSIVE BOLTZMANN GENERATORS (ArBG), the first scalable, autoregressive and diffeomorphism-free method for Boltzmann Generation.

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 Large Language Models. 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.

Paper Context