Generate proteins like LLMs generate text — autoregressively, atom by atom.

Primer

Boltzmann Generators (BGs) combine an exact-likelihood generative model $p_\theta(x)$ with an energy function $\mathcal{E}(x)$ to align observables with the target Boltzmann distribution:

where $x \in \mathbb{R}^{n \times 3}$ is a conformation of $n$ atoms and $\mathcal{E}$ is the (dimensionless) potential energy that maps $\mathbb{R}^{n \times 3} \to \mathbb{R}$. The correction step — self-normalized importance sampling (SNIS) — assigns each sample $x^i \sim p_\theta$ a weight $w(x^i) = e^{-\mathcal{E}(x^i)}/p_\theta(x^i)$, which serves to de-bias samples in accordance with the Boltzmann density. For this to be consistent, the model must provide an exact likelihood $p_\theta(x)$ and a proposal that lies within the support of the target — two requirements that drive architecture choice.

Normalizing flows

Discrete-time (NFs)

Discrete-time NFs compose a stack of analytically invertible layers $f_1, \ldots, f_K$ that transport a prior sample $z_0 \sim p_0$ to $x = f_K \circ \cdots \circ f_1(z_0)$. The resulting transport map is $\mathbb{R}^d \to \mathbb{R}^d$, and conserves probability mass through the network via a discrete change of variables. Consequently, the likelihood follows directly from the change-of-variables formula:

Using this formulation, likelihoods are efficient, requiring one forward pass; however, every layer must admit a tractable Jacobian determinant, which sharply restricts the family of admissible architectures. A deeper, more structural problem is topological: a flow with a Gaussian prior is a homeomorphism, while the Boltzmann distribution of a peptide is multi-modal with approximately disconnected support. Mapping a single connected prior onto several well-separated basins forces the model to deform space across thin bridges, which can produce poorly-conditioned Jacobians, unstable training, and a practical ceiling on expressivity.

Continuous-time (CNFs)

CNFs replace the layer-by-layer composition with a learned vector field $v_\theta(x_t, t)$ that transports a prior sample at $t = 0$ to a model sample at $t = 1$. The likelihood is given by the instantaneous change-of-variables — an integral of the trace of the Jacobian:

This relaxes the architectural constraints, but the trace term is what makes likelihood evaluation expensive: computing it exactly costs $O(d)$ network evaluations per integration step. Stochastic estimators (e.g. Hutchinson’s trace estimator) bring the cost down but inject variance into the SNIS weights. Either way, the per-sample likelihood becomes the inference bottleneck — and since SNIS variance shrinks with the number of samples, this directly caps how well the corrected estimator can perform.

Autoregressive Boltzmann Generators

Both problems disappear if we drop the flow and factor the joint density over coordinates instead. Flattening $x \in \mathbb{R}^{n \times 3}$ to a $d$-dimensional vector $(x_1, \ldots, x_d)$ and applying the chain rule yields:

The joint log-likelihood is a sum of $d$ scalar log-probabilities computed in a single forward pass — no Jacobian determinant, no ODE integration, and no architectural constraint on invertibility.

There is, however, an immediate mismatch to resolve. Atomic coordinates live in continuous space ($x_j \in \mathbb{R}$), while autoregressive models — transformers in particular — are built to emit categorical distributions over a finite vocabulary of tokens. How do you parameterize a continuous density with an intrinsically discrete model?

Bin parameterization

Each conditional $p_\theta(x_j \vert x_{\lt j})$ is parameterized as a categorical distribution over a uniform partition of the coordinate’s support $[C_{\min}, C_{\max}]$ into $L$ bins. At sampling time, a continuous coordinate is recovered by drawing uniformly within the selected bin. Training uses maximum likelihood estimation (MLE) with teacher forcing: at every position $j$ the model sees the ground-truth prefix $x_{\lt j}$ and minimizes the cross-entropy of its predicted bin distribution against the bin index of the ground-truth $x_j$. Because each conditional is a categorical, the resulting density on $x$ is piecewise-constant; with sufficiently many bins it converges to the underlying continuous distribution and the SNIS weights remain exact in the continuous limit.

Architecture

We use a decoder-only transformer with causal self-attention — the same family of architectures used to power modern large language models. Per-coordinate embeddings carry atom type, residue type, and positional information. For Robin, our 132 million parameter transferable Boltzmann Generator, the conditioning is augmented with residue position and overall sequence length, which turns the model into a sequence-to-structure generator: given a peptide sequence, it samples an equilibrium ensemble.

Inference-time techniques

Because the model is autoregressive, decoding tricks developed for large language models apply almost without modification. Below, $\tau$ denotes the sampling temperature (distinct from the physical temperature, which we have already absorbed into $\mathcal{E}$):

• Temperature scaling. Logits are divided by $\tau > 0$ before the softmax, which controls sample diversity. For transferable generation, we tune $\tau$ to optimize performance, and use $\tau = 0.95$ for all studies conducted.
• KV-caching. The attention keys and values computed for the prefix are stored and reused, so generating each new coordinate is a single-token append rather than a full re-encode.
• Twisted SMC. SMC maintains a population of partial generations and resamples at residue boundaries using twist functions: scalar approximations of the future likelihood under the target, computed from the partial energy of the residues generated so far. Implausible substructures (e.g., steric clashes) are filtered out before they propagate. Autoregressive sampling is well-suited to this, since intermediate prefixes can be scored before the rest of the conformation is generated.

Robin

Single-system results

We start by training a separate model per system on five peptides, ranging from the 2-residue alanine dipeptide to the 10-residue Chignolin. Following prior work, the ground-truth reference is taken from long unbiased molecular dynamics (MD) trajectories. The first thing to check is whether the proposal $p_\theta$ overlaps the target well enough that the SNIS-corrected distribution recovers the Boltzmann statistics. Figure 5 overlays the proposal, the SNIS-resampled distribution, and the MD reference for each system.

The quantitative picture across four systems is summarized in Table 1; the energy and structural distributions behind these numbers follow.

Agreement on the energy is a local consistency check; a more demanding question is whether the sampler actually covers the relevant regions of conformational space. Chignolin is the standard benchmark for this: its $\beta$-hairpin secondary structure produces several well-separated torsional modes, and prior flow-based methods tend to miss some of them. Figure 6 places the per-residue Ramachandran plots from the MD reference and from our model side by side.

Transferable generation

The next question is whether a single network can learn to sample many peptides at once. We train Robin, a 132 million parameter conditional causal transformer, on ManyPeptidesMD — a corpus of biased MD trajectories spanning many distinct peptide sequences — and evaluate it on held-out unbiased 8-residue peptides that never appear during training. As an aggregate metric we report the squared 2-Wasserstein distance between predicted and reference energy distributions, denoted $\mathcal{E}\text{-}\mathcal{W}_2$. Averaged over 30 unseen octapeptides, Robin reduces $\mathcal{E}\text{-}\mathcal{W}_2$ by roughly 60% relative to Prose, the prior flow-based state-of-the-art. Table 2 reports the aggregate zero-shot numbers; Figure 7 then zooms in on which conformational modes the model is actually capturing.

References

1. Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu (2019). Boltzmann Generators: Sampling equilibrium states of many-body systems with deep learning. Science.
2. Danilo Jimenez Rezende and Shakir Mohamed (2015). Variational Inference with Normalizing Flows. ICML.
3. Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud (2018). Neural Ordinary Differential Equations. NeurIPS.
4. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin (2017). Attention Is All You Need. NeurIPS.
5. Danyal Rehman, Charlie B. Tan, Yoshua Bengio, Joey Bose, and Alexander Tong (2026). Autoregressive Boltzmann Generators. ICML.

Cite

If you found this useful, the BibTeX is below:

@inproceedings{rehman2026arbg,
  title        = {{A}utoregressive {B}oltzmann {G}enerators},
  author       = {Rehman, Danyal and Tan, Charlie B and Bengio, Yoshua and Bose, Joey and Tong, Alexander},
  booktitle    = {International Conference on Machine Learning (ICML)},
  year         = {2026},
  url          = {https://arxiv.org/abs/2606.27361},
  doi          = {10.48550/arXiv.2606.27361},
  eprint       = {2606.27361},
  archivePrefix= {arXiv},
  primaryClass = {cs.LG},
  note         = {Spotlight at ICML 2026}
}

@inproceedings{rehman2026arbg,
  title        = {{A}utoregressive {B}oltzmann {G}enerators},
  author       = {Rehman, Danyal and Tan, Charlie B and Bengio, Yoshua and Bose, Joey and Tong, Alexander},
  booktitle    = {International Conference on Machine Learning (ICML)},
  year         = {2026},
  url          = {https://arxiv.org/abs/2606.27361},
  doi          = {10.48550/arXiv.2606.27361},
  eprint       = {2606.27361},
  archivePrefix= {arXiv},
  primaryClass = {cs.LG},
  note         = {Spotlight at ICML 2026}
}