# Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules

Joseph C. Kim1,2, David Bloore1, Karan Kapoor1, Jun Feng1
Ming-Hong Hao1, Mengdi Wang2
1Ensem Therapeutics, Waltham, MA
2Department of Electrical and Computer Engineering, Princeton University

###### Abstract

The Boltzmann distribution of a protein provides a roadmap to all of its
functional states. Normalizing flows are a promising tool for modeling this
distribution, but current methods are intractable for typical pharmacological
targets; they become computationally intractable due to the size of the
system, heterogeneity of intra-molecular potential energy, and long-range
interactions. To remedy these issues, we present a novel flow architecture
that utilizes split channels and gated attention to efficiently learn the
conformational distribution of proteins defined by internal coordinates. We
show that by utilizing a 2-Wasserstein loss, one can smooth the transition
from maximum likelihood training to energy-based training, enabling the
training of Boltzmann Generators for macromolecules. We evaluate our model and
training strategy on villin headpiece HP35(nle-nle), a 35-residue subdomain,
and protein G, a 56-residue protein. We demonstrate that standard
architectures and training strategies, such as maximum likelihood alone, fail
while our novel architecture and multi-stage training strategy are able to
model the conformational distributions of protein G and HP35.

##  1 Introduction

The structural ensemble of a protein determines its functions. The
probabilities of the ground and metastable states of a protein at equilibrium
for a given temperature determine the interactions of the protein with other
proteins, effectors, and drugs, which are keys for pharmaceutical development.
However, enumeration of the equilibrium conformations and their probabilities
is infeasible. Since complete knowledge is inaccessible, we must adopt a
sampling approach. Conventional approaches toward sampling the equilibrium
ensemble rely on Markov-chain Monte Carlo or molecular dynamics (MD). These
approaches explore the local energy landscape adjacent a starting point;
however, they are limited by their inability to penetrate high energy
barriers. In addition, MD simulations are expensive and scale poorly with
system size. This results in incomplete exploration of the equilibrium
conformational ensemble.

In their pioneering work, Noé et al. (2019) proposed a normalizing flow model
(Rezende & Mohamed, 2015), that is trained on the energy function of a many-
body system, termed Boltzmann generators (BGs). The model learns an invertible
transformation from a system’s configurations to a latent space
representation, in which the low-energy configurations of different states can
be easily sampled. As the model is invertible, every latent space sample can
be back-transformed to a system configuration with high Boltzmann probability,
i.e., $p(\mathbf{x})\propto e^{-u(\mathbf{x})/(kT)}$.

A normalizing flow-based generative model is constructed by a sequence of
invertible transformations (Rezende & Mohamed, 2015). BGs typically employ
flow models because they can be sampled from efficiently and they describe a
tractable probability density function. This allows us to employ reverse KL
divergence training since we can compute an unnormalized density for the
target Boltzmann distribution (Noé et al., 2019; Wirnsberger et al., 2022;
Köhler et al., 2021).

BGs in the literature have often struggled with even moderate-sized proteins,
due to the complexity of conformation dynamics and scarcity of available data.
Most works have focused on small systems like alanine dipeptide (22 atoms)
(Köhler et al., 2021; Midgley et al., 2022; Invernizzi et al., 2022). To date,
only two small proteins, BPTI and bromodomain, have been modeled by BGs. Noé
et al. (2019) trained a BG for BPTI, a 58 amino acid structure, at all-atom
resolution. Unfortunately, the training dataset used is licensed by DESRES
(Shaw et al., 2010) and not open-source. No works since have shown success on
proteins of similar size at all-atom resolution or reported results for BPTI.
Mahmoud et al. (2022) trained a BG for bromodomain, a 100 residue protein,
with a SIRAH coarse-grained representation. However, drug design applications
require much finer resolution than resolvable by SIRAH. A thorough review of
related works is detailed in Appendix A.

The limited scope of flow model BG applications is due to the high
computational expense of their training process. Their invertibility
requirement limits expressivity when modeling targets whose supports have
complicated topologies (Cornish et al., 2019), necessitating the use of many
transformation layers. Another hurdle in scaling BGs is that proteins often
involve long-range interactions; atoms far apart in sequence can interact with
each other. In this work, we present a new BG method for general proteins with
the following contributions:

  * •

We use a global internal coordinate representation with fixed bond-lengths and
side-chain angles. From a global structure and energetics point-of-view,
little information is lost by allowing side-chain bonds to only rotate. Such a
representation not only reduces the number of variables but also samples
conformations more efficiently than Cartesian coordinates (Noé et al., 2019;
Mahmoud et al., 2022).

  * •

The global internal coordinate representation is initially split into a
backbone channel and a side-chain channel. This allows the model to
efficiently capture the distribution of backbone internal coordinates, which
most controls the overall global conformation.

  * •

A new NN architecture for learning the transformation parameters of the
coupling layers of the flow model which makes use of gated attention units
(GAUs) (Hua et al., 2022) and a combination of rotary positional embeddings
(Su et al., 2021) with global, absolute positional embeddings for learning
long range interactions.

  * •

To handle global conformational changes, a new loss-function, similar in
spirit to the Fréchet Inception Distance (FID) (Heusel et al., 2017), is
introduced to constrain the global backbone structures to the space of native
conformational ensemble.

We show in this work that our new method can efficiently generate Boltzmann
distributions and important experimental structures in two different protein
systems. We demonstrate that the traditional maximum likelihood training for
training flow models is insufficient for proteins, but our multi-stage
training strategy can generate samples with high Boltzmann probability.

##  2 Background

###  2.1 Normalizing Flows

Normalizing flow models learn an invertible map
$f:\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$ to transform a random variable
$\bm{z}\sim q_{Z}$ to the random variable $\bm{x}=f(\bm{z})$ with distribution

$q_{X}(\bm{x})=q_{Z}(\bm{z})|\text{det}(J_{f}(\bm{z}))|^{-1},$ (1)

where $J_{f}(\bm{z})=\partial f/\partial\bm{z}$ is the Jacobian of $f$. We can
parameterize $f$ to approximate a target distribution $p(\bm{x})$. To simplify
notation, we refer to the flow distribution as $q_{\theta}$, where $\theta$
are the parameters of the flow. If samples from the target distribution are
available, the flow can be trained via maximum likelihood. If the unnormalized
target density $p(\bm{x})$ is known, the flow can be trained by minimizing the
KL divergence between $q_{\theta}$ and $p$, i.e.,
$\text{KL}(q_{\theta}||p)=\int_{X}q_{\theta}(\bm{x})\log(q_{\theta}(\bm{x})/p(\bm{x}))d\bm{x}$.

###  2.2 Distance Matrix

A protein distance matrix is a square matrix of Euclidean distances from each
atom to all other atoms. Practitioners typically use $C\alpha$ atoms or
backbone atoms only. Protein distance matrices have many applications
including structural alignment, protein classification, and finding homologous
proteins (Holm & Sander, 1993; Holm, 2020; Zhu et al., 2023). They have also
been used as representations for protein structure prediction algorithms,
including the first iteration of AlphaFold (Senior et al., 2019; Xu & Wang,
2019; Hou et al., 2019).

###  2.3 2-Wasserstein Distance

The 2-Wasserstein Distance is a measure of the distance between two
probability distributions. Let $P=\mathcal{N}(\bm{\mu}_{P},\bm{\Sigma}_{P})$
and $Q=\mathcal{N}(\bm{\mu}_{Q},\bm{\Sigma}_{Q})$ be two normal distributions
in $\mathbb{R}^{d}$. Then, with respect to the Euclidean norm on
$\mathbb{R}^{d}$, the squared 2-Wasserstein distance between $P$ and $Q$ is
defined as

$W_{2}(P,Q)^{2}=\mathopen{\|}\bm{\mu}_{P}-\bm{\mu}_{Q}\mathclose{\|}_{2}^{2}+\text{trace}(\bm{\Sigma}_{P}+\bm{\Sigma}_{Q}-2(\bm{\Sigma}_{P}\bm{\Sigma}_{Q})^{1/2}).$
(2)

In computer vision, the Fréchet Inception Distance (FID) (Heusel et al., 2017)
computes the 2-Wasserstein distance and is often used as an evaluation metric
to measure generated image quality.

Figure 1: (a) Our split flow architecture. (b) Each transformation block
consists of a gated attention rational quadratic spline (RQS) coupling layer.
(c) Example structures of protein G from the flow $q_{\theta}$ (left) and from
molecular dynamics simulation $p$ (right). We also show sample distance
matrices $\mathbf{D}(\mathbf{x}_{q_{\theta}})$ and
$\mathbf{D}(\mathbf{x}_{p})$.

##  3 Scalable Boltzmann Generators

###  3.1 Problem Setup

BGs are generative models that are trained to sample from the Boltzmann
distribution for physical systems, i.e., $p(\mathbf{x})\propto
e^{-u(\mathbf{x})/(kT)}$, where $u(\mathbf{x})$ is the potential energy of the
conformation $\mathbf{x}$, $k$ is the Boltzmann constant, and $T$ is the
temperature. A protein conformation is defined as the arrangement in space of
its constituent atoms (Fig. 2), specifically, by the set of 3D Cartesian
coordinates of its atoms. Enumeration of metastable conformations for a
protein at equilibrium is quite challenging with standard sampling techniques.
We tackle this problem with generative modeling. Throughout this work, we
refer to $p$ as the ground truth conformation distribution and $q_{\theta}$ as
the distribution parameterized by the normalizing flow model $f_{\theta}$.

###  3.2 Reduced internal coordinates

Energetically-favored conformational changes take place via rotations around
single chemical bonds while bond vibrations and angle bending at physiologic
temperature result in relatively small spatial perturbations (Vaidehi & Jain,
2015). Our focus on near ground and meta-stable states therefore motivates the
use of internal coordinates: $N-1$ bond lengths $d$, $N-2$ bond angles
$\theta$, and $N-3$ torsion angles $\phi$, where $N$ is the number of atoms of
the system (see Fig. 2). In addition, internal coordinate representation is
translation and rotation invariant.

Figure 2: A two residue chain. Hydrogens on carbon atoms are omitted. Backbone
atoms are highlighted green. Shown is an example of a bond length $d$, a bond
angle $\theta$, and a dihedral/torsion angle $\phi$.

A protein can be described as a branching graph structure with a set of
backbone atoms and non-backbone atoms (we will colloquially refer to these as
side-chain atoms). Previous works have noted the difficulty in working with
internal coordinate representations for the backbone atoms (Noé et al., 2019;
Köhler et al., 2022; Mahmoud et al., 2022). This is due to the fact that
protein conformations are sensitive to small changes in the backbone torsion
angles. Noé et al. (2019) introduced a coordinate transformation whereby the
side-chain atom coordinates are mapped to internal coordinates while the
backbone atom coordinates are linearly transformed via principal component
analysis (PCA) and the six coordinates with the lowest variance are
eliminated. However, as mentioned by Midgley et al. (2022), the mapping of
vectors onto a fixed set of principal components is generally not invariant to
translations and rotations. In addition, PCA suffers from distribution shift.

A full internal coordinate system requires $3N-6$ dimensions where $N$ is the
number of atoms. Bond lengths hardly vary in equilibrium distributions while
torsion angles can vary immensely. We treat non-backbone bond angles as
constant, again replaced by their mean. Heterocycles in the sidechains of Trp,
Phe, Tyr and His residues are treated as rigid bodies. Our final
representation is

$\mathbf{x}=[\bm{\theta}_{bb},\bm{\phi}_{bb},\bm{\phi}_{sc}],$

where the subscripts $bb$ and $sc$ indicate backbone and sidechain,
respectively. This dramatically reduces the input dimension and keeps the most
important features for learning global conformation changes in the equilibrium
distribution.

Recent works have adopted similar approaches to reduce the input number of
dimensions. Wu et al. (2022) utilize only the backbone torsion and bond angles
to represent proteins for novel protein design and generation, while Wang et
al. (2022) simply use the backbone torsion angles to represent the polypeptide
AiB9 in a similar task of modeling the polypeptide conformational landscape.

###  3.3 Training and evaluation

We train BGs with MD simulation data at equilibrium, i.e., the distribution of
conformations is constant and not changing as with, for example, folding. We
seed the simulation with energetically stable native conformations. BG
training aims to learn to sample from the Boltzmann distribution of protein
conformations. We compute the energy of samples generated by our model under
the AMBER14 forcefield (Case et al., 2014) and report their mean. In addition,
in order to evaluate how well the flow model generates the proper backbone
distribution, we define a new measure:

######  Definition 3.1 (Distance Distortion).

Let $\mathcal{A}_{bb}$ denote the indices of backbone atoms. Define
$\mathbf{D}(\mathbf{x})$ as the pairwise distance matrix for the backbone
atoms of $\mathbf{x}$. Define
$\mathcal{P}=\\{(i,j)|i,j\in\mathcal{A}_{bb}\text{ and }i<j\\}$. The distance
distortion is defined as

$\Delta
D\coloneqq\mathop{\mathbb{E}}_{\begin{subarray}{c}\mathbf{x}_{q_{\theta}}\sim{q_{\theta}}\\\
\mathbf{x}_{p}\sim
p\end{subarray}}\left[\frac{1}{|\mathcal{P}|}\sum_{(i,j)\in\mathcal{P}}|\mathbf{D}(\mathbf{x}_{q_{\theta}})_{ij}-\mathbf{D}(\mathbf{x}_{p})_{ij}|\right],$
(3)

###  3.4 Split Flow Architecture

We use Neural Spline Flows (NSF) with rational quadratic splines (Durkan et
al., 2019) having 8 bins each. The conditioning is done via coupling. Torsion
angles $\bm{\phi}$ can freely rotate and are therefore treated as periodic
coordinates (Rezende et al., 2020).

The full architectural details are highlighted in Fig. 1. We first split the
input into backbone and sidechain channels:

$\displaystyle\mathbf{x}_{bb}$
$\displaystyle=[\bm{\theta}_{bb},\bm{\phi}_{bb}],\quad\quad\mathbf{x}_{sc}=[\bm{\phi}_{sc}].$

We then pass the backbone inputs through $L_{bb}=48$ gated attention rational
quadratic spline (GA-RQS) coupling blocks. As all the features are angles in
$[-\pi,\pi]$, we augment the features with their mapping on the unit circle.
In order to utilize an efficient attention mechanism, we employ gated
attention units (GAUs) (Hua et al., 2022). In addition, we implement relative
positional embeddings (Shaw et al., 2018) on a global level so as to allow
each coupling block to utilize the correct embeddings. The backbone latent
embeddings are then concatenated with the side chain features and passed
through $L=10$ more GA-RQS coupling blocks.

###  3.5 Multi-stage training strategy

Normalizing flows are most often trained by maximum likelihood, i.e.,
minimizing the negative log likelihood (NLL)

$\mathcal{L}_{\text{NLL}}(\theta)\coloneqq-\mathbb{E}_{\mathbf{x}\sim p}[\log
q_{\theta}(\mathbf{x})],$ (4)

or by minimizing the reverse KL divergence 111We refer to the reverse KL
divergence as just “KL divergence” or “KL loss”, as is often done in the
literature.:

$\mathcal{L}_{\text{KL}}(\theta)\coloneqq\mathbb{E}_{\mathbf{x}\sim
q_{\theta}}[\log(q_{\theta}(\mathbf{x})/p(\mathbf{x}))].$ (5)

In the BG literature, minimizing the KL divergence is often referred to as
“training-by-energy”, as the expression can be rewritten in terms of the
energy of the system. The reverse KL divergence suffers from mode-seeking
behavior, which is problematic when learning multimodal target distributions.
While minimizing the NLL is mass-covering, samples generated from flows
trained in this manner suffer from high variance. In addition, for larger
systems, maximum likelihood training often results in high-energy generated
samples.

In their seminal paper, Noé et al. (2019) used a convex combination of the two
loss terms, in the context of BGs, in order to both avoid mode-collapse and
generate low-energy samples. However, for larger molecules, target evaluation
is computationally expensive and dramatically slows iterative training with
the reverse KL divergence objective. In addition, during the early stages of
training, the KL divergence explodes and leads to unstable training. One way
to circumvent these issues is to train with the NLL loss, followed by a
combination of both loss terms. Unfortunately, for larger systems, the KL term
tends to dominate and training often get stuck at non-optimal local minima. In
order to remedy these issues, we consider a sequential training scheme,
whereby we smooth the transition from maximum likelihood training to the
combination of maximum likelihood and reverse KL divergence minimization.

(1) As mentioned previously, we first train with maximum likelihood to
convergence.

(2) Afterward, we train with a combination of the NLL and the 2-Wasserstein
loss with respect to distance matrices of the backbone atoms:

$\mathcal{L}_{\text{W}}(\theta)\coloneqq\mathopen{\|}\bm{\mu}_{q_{\theta}}-\bm{\mu}_{p}\mathclose{\|}_{2}^{2}+\text{trace}(\bm{\Sigma}_{q_{\theta}}+\bm{\Sigma}_{p}-2(\bm{\Sigma}_{q_{\theta}}\bm{\Sigma}_{p})^{1/2}),$
(6)

where

$\displaystyle\bm{\mu}_{p}$ $\displaystyle\coloneqq\mathbb{E}_{\mathbf{x}\sim
p}[\mathbf{x}_{bb}],\quad$
$\displaystyle\bm{\Sigma}_{p}\coloneqq\mathbb{E}_{\mathbf{x}\sim
p}[(\mathbf{x}_{bb}-\bm{\mu}_{p})(\mathbf{x}_{bb}-\bm{\mu}_{p})^{\top}]$ (7)
$\displaystyle\bm{\mu}_{q_{\theta}}$
$\displaystyle\coloneqq\mathbb{E}_{\mathbf{x}\sim
q_{\theta}}[\mathbf{x}_{bb}],\quad$
$\displaystyle\bm{\Sigma}_{q_{\theta}}\coloneqq\mathbb{E}_{\mathbf{x}\sim
q_{\theta}}[(\mathbf{x}_{bb}-\bm{\mu}_{q_{\theta}})(\mathbf{x}_{bb}-\bm{\mu}_{q_{\theta}})^{\top}]$
(8)

are mean and covariance, respectively, of the vectorized backbone atom
distance matrices.

(3) As a third stage of training, we train with a combination of the NLL, the
2-Wasserstein loss, and the KL divergence. In our final stage of training, we
drop the 2-Wasserstein loss term and train to minimize a combination of the
NLL and the KL divergence.

##  4 Results

###  4.1 Protein Systems

Alanine dipeptide (ADP) is a two residue (22-atoms) common benchmark system
for evaluating BGs (Noé et al., 2019; Köhler et al., 2021; Midgley et al.,
2022). We use the MD simulation datasets provided by Midgley et al. (2022) for
training and validation.

HP35(nle-nle), a 35-residue double-mutant of the villin headpiece subdomain,
is a well-studied structure whose folding dynamics have been observed and
documented (Beauchamp et al., 2012). For training, we use the MD simulation
dataset made publicly available by Beauchamp et al. (2012) and remove faulty
trajectories and unfolded structures as done by Ichinomiya (2022).

Protein G is a 56-residue cell surface-associated protein from Streptococcus
that binds to IgG with high affinity (Derrick & Wigley, 1994). In order to
train our model, we generated samples by running a MD simulation. The crystal
structure of protein G, 1PGA, was used as the seed structure. The
conformational space of Protein G was first explored by simulations with
ClustENMD (Kaynak et al., 2021). From 3 rounds of ClustENMD iteration and
approximately 300 generated conformations, 5 distinctly diverse structures
were selected as the starting point for equilibrium MD simulation by Amber. On
each starting structure, 5 replica simulations were carried out in parallel
with different random seeds for 400 ns at 300 K. The total simulation time of
all the replicas was accumulated to 1 microsecond. Thus, $10^{6}$ structures
of protein G were saved over all the MD trajectories.

As a baseline model for comparison, we use Neural Spline Flows (NSF) with 58
rational quadratic spline coupling layers (Durkan et al., 2019). NSFs have
been used in many recent works on BGs (Köhler et al., 2022; Midgley et al.,
2022; Mahmoud et al., 2022). In particular, Midgley et al. (2022) utilized the
NSF model (with fewer coupling layers) in their experiments with alanine
dipeptide, a two residue system. In our experiments with ADP and HP35, we
utilize 48 GA-RQS coupling layers for the backbone followed by 10 GA-RQS
coupling layers for the full latent size. We also ensure that all models have
a similar number of trainable parameters. We use a Gaussian base distribution
for non-dihedral coordinates. For ADP and HP35, we use a uniform distribution
for dihedral coordinates. For protein G, we use a von Mises base distribution
for dihedral coordinates; we noticed that using a von Mises base distribution
improved training for the protein G system as compared to a uniform or
truncated normal distribution.

###  4.2 Main Results

Table 1: Training BGs with different strategies. We compute $\Delta D$, energy $u(\cdot)$, and mean NLL of $10^{6}$ generated structures after training with different training strategies with ADP, protein G, and Villin HP35. $\Delta D$ is computed for batches of $10^{3}$ samples. Means and standard deviations are reported. Statistics for $u(\cdot)$ are reported for structures with energy below the median sample energy. Best results are bold-faced. For reference, the energy for training data structures is $-317.5\pm 125.5$ kcal/mol for protein G and $-1215.5\pm 222.2$ kcal/mol for villin HP35. We compare our results against a Neural Spline Flows (NSF) baseline model. |  | Training strategy |  |  |
---|---|---|---|---|---
System | Arch. | NLL | KL | W2 |  $\Delta D$ (Å) | Energy $u(\mathbf{x})$ (kcal/mol) | $-\mathbb{E}_{p(\mathbf{x})}[\log q_{\theta}(\mathbf{x})]$
| NSF | ✓ |  |  | $0.09\pm 0.01$ | $\mathbf{(-1.19\pm 0.61)\times 10^{1}}$ | $38.29\pm 0.19$
|  | ✓ |  |  | $0.08\pm 0.01$ | $\mathbf{(-1.18\pm 0.65)\times 10^{1}}$ | $\mathbf{36.15\pm 0.15}$
ADP | Ours | ✓ | ✓ |  | $0.05\pm 0.01$ | $\mathbf{(-1.20\pm 0.59)\times 10^{1}}$ | $38.66\pm 0.19$
|  | ✓ |  | ✓ | $\mathbf{0.04\pm 0.00}$ | $\mathbf{(-1.06\pm 0.74)\times 10^{1}}$ | $38.12\pm 0.03$
| Ours | ✓ | ✓ | ✓ | $\mathbf{0.03\pm 0.01}$ | $\mathbf{(-1.31\pm 0.52)\times 10^{1}}$ | $37.67\pm 0.09$
| NSF | ✓ |  |  | $2.92\pm 0.80$ | $(2.15\pm 3.31)\times 10^{10}$ | $-263.46\pm 0.13$
|  | ✓ |  |  | $1.81\pm 0.14$ | $(9.47\pm 15.4)\times 10^{8}$ | $\mathbf{-310.11\pm 0.08}$
Protein G | Ours | ✓ | ✓ |  | $16.09\pm 1.14$ | $(2.86\pm 0.62)\times 10^{2}$ | $-308.68\pm 0.08$
|  | ✓ |  | ✓ | $\mathbf{0.18\pm 0.01}$ | $(2.68\pm 4.31)\times 10^{6}$ | $-307.17\pm 0.01$
| Ours | ✓ | ✓ | ✓ | $\mathbf{0.19\pm 0.01}$ | $\mathbf{(-3.04\pm 1.24)\times 10^{2}}$ | $\mathbf{-309.10\pm 0.91}$
| NSF | ✓ |  |  | $0.81\pm 0.06$ | $(7.78\pm 17.4)\times 10^{7}$ | $687.95\pm 1.92$
|  | ✓ |  |  | $0.65\pm 0.04$ | $(5.29\pm 11.7)\times 10^{6}$ | $\mathbf{651.90\pm 2.88}$
HP35 | Ours | ✓ | ✓ |  | $0.61\pm 0.04$ | $(6.46\pm 14.3)\times 10^{2}$ | $678.38\pm 0.87$
|  | ✓ |  | ✓ | $\mathbf{0.38\pm 0.03}$ | $(1.15\pm 1.76)\times 10^{7}$ | $678.31\pm 1.55$
| Ours | ✓ | ✓ | ✓ | $\mathbf{0.39\pm 0.03}$ | $\mathbf{(-4.66\pm 3.52)\times 10^{2}}$ | $667.45\pm 2.04$
Figure 3: Sample conformations generated by BG via different training
strategies. (a) Root mean square fluctuation (RMSF) computed for each residue
(C$\alpha$ atoms) in HP35 and protein G. Matching the training dataset’s plot
is desirable. (b) Examples of HP35 from ground truth training data, generated
samples from our model, and generated samples from the baseline model. (c)
Example of two metastable states from protein G training data. (d) Low-energy
conformations of protein G generated by our model superimposed on each other.
We also show some examples of pathological structures generated after training
with different training paradigms: NLL (maximum likelihood), both NLL and KL
divergence, and NLL and the 2-Wasserstein loss. Atom clashes are highlighted
with red circles.

From Table 1, we see that our model has marginal improvements over the
baseline model for ADP. This is not surprising as the system is extremely
small, and maximum likelihood training sufficiently models the conformational
landscape.

For both proteins, our model closely captures the individual residue
flexibility as analyzed by the root mean square fluctuations (RMSF) of the
generated samples from the various training schemes in Fig. 3(a). This is a
common metric for MD analysis, where larger per-residue values indicate larger
movements of that residue relative to the rest of the protein. Fig. 3(a)
indicates that our model generates samples that present with similar levels of
per-residue flexibility as the training data.

Table 1 displays $\Delta D$, the mean energy, and the mean NLL of structures
generated from flow models trained with different strategies. For each model
(architecture and training strategy), we generate $3\times 10^{6}$
conformations ($10^{6}$ structures over 3 random seeds) after training with
either protein G or Villin HP35. Due to the cost of computing $\Delta D$, we
compute it for batches of $10^{3}$ samples and report statistics (mean and
standard deviation) over the $3\times 10^{3}$ batches. Before we computed
sample statistics for the energy $u$, we first filtered out the samples with
energy higher than the median value. This was done to remove high energy
outliers that are not of interest and would noise the reported mean and
standard deviations. We also report the mean and standard deviation (across 3
seeds) for the average NLL. We see that our model is capable of generating
low-energy, stable conformations for these two systems while the baseline
method and ablated training strategies produce samples with energies that are
positive and five or more orders of magnitude greater.

Table 1 highlights a key difference in the results for protein G and villin
HP35. For villin, models trained by reverse KL and without the 2-Wasserstein
loss do not result in completely unraveled structures. This is consistent with
the notion that long-range interactions become much more important in larger
structures. From Fig. 3(b), we see that Villin HP35 is not densely packed, and
local interactions, e.g., as seen in $\alpha$-helices, are more prevalent than
long-range interactions/contacts. In addition, we see that our model generates
diverse alternative modes of the folded villin HP35 protein that are
energetically stable compared to the structures obtained from the baseline
model.

Fig. 3(d) visualizes pathological structures of protein G generated via
different training schemes. In Fig. 3(d, left), we see that minimizing the NLL
generally captures local structural motifs, such as the $\alpha$-helix.
However, structures generated by training only with the NLL loss tend to have
clashes in the backbone, as highlighted with red circles in Fig 3(d), and/or
long range distortions. This results in large van der waals repulsion as
evidenced by the high average energy values in Table 1.

In Fig. 3(d, middle), we see that structures generated by minimizing a
combination of the NLL loss and the reverse KL divergence unravel and present
with large global distortions. This results from large, perpetually unstable
gradients during training. In Fig. 3(d, right), we see that training with a
combination of the NLL loss and the 2-Wasserstein loss properly captures the
backbone structural distribution, but tends to have clashes in the sidechains.
Table 1 demonstrates that only our model with our multistage training strategy
is able to achieve both low energy samples and proper global structures. The
2-Wasserstein loss prevents large backbone distortions, and thus,
simultaneously minimizing the reverse KL divergence accelerates learning for
the side chain marginals with respect to the backbone and other side chain
atoms.

###  4.3 BGs can generate novel samples

Figure 4: BGs can generate novel sample conformations. (a) Protein G 2D UMAP
embeddings for the training data, test data, and $2\times 10^{5}$ generated
samples. (b) A representative example of generated structures by the BG model
which was not found in training data (cyan) and the closest structure in the
training dataset (magenta) by RMSD. Both structures are depicted as stars with
their respective structural colors in (a). (c) Protein G energy distribution
of training dataset (orange) and samples (blue) generated by our model. The
second energy peak of the sampled conformations covers the novel structure
shown in (b). (d) An overlay of high-resolution, lowest-energy all-atom
structures of protein G generated by the BG model. This demonstrates that our
model is capable of sampling low-energy conformations at atomic resolution.

One of the primary goals for development of BG models is to sample important
metastable states that are unseen or difficult to sample by conventional MD
simulations. Protein G is a medium-size protein with diverse metastable states
that provide a system for us to evaluate the capability of our BG model.
First, we visualize 2D UMAP embeddings (McInnes et al., 2018) for the training
data set, test dataset, and for $2\times 10^{5}$ generated samples of protein
G in Fig. 4(a). We see that the test dataset 4(a, middle), an independent MD
dataset, covers far less conformational space than an equivalent number of BG
samples as shown in Fig. 4(a, right).

Secondly, we computed, respectively, the energy distributions of the training
set from MD simulations and sample set from the BG model as shown by Fig.
4(c). Unlike the training set, the BG sample energy distribution is bimodal.
Analysis of structures in the second peak revealed a set of conformations not
present in the training set. These new structures are characterized by a large
bent conformation in the hair-pin loop which links beta-strand 3 and 4 of
protein G. Fig. 4(b) compares representative structures of the discovered new
structure (magenta) with the closest structure (by RMSD) in the training
dataset (cyan). We also see vastly different sidechain conformations along the
bent loops between two structures. Energy minimization on the discovered new
structures demonstrated that these new structures are local-minimum metastable
conformations. Thirdly, we carefully examined the lowest-energy conformations
generated by the BG model. Fig. 4(d) showcases a group of lowest-energy
structures generated by the BG model, overlaid by backbone and all-atom side
chains shown explicitly. All of these structures are very similar to the
crystal structure of protein G, demonstrating that the trained BG model is
capable of generating protein structures with high quality at the atomic
level.

##  5 Discussion

The scalability of the Boltzmann generator for large macromolecules is a major
challenge. This study developed a new flow model architecture to address this
challenge. We represented protein structures using internal coordinates and
conducted conformational sampling primarily in dihedral angle space. This
approach reduces the number of variables to be trained in the flow model, and
conformational sampling follows the energetics of proteins. We made several
innovations to the neural network architecture, such as using gated attention
units for transformation maps and rotary positional embedding to capture non-
local interactions. We also introduced split channels to allocate more
transformation layers for backbone representations and employed a
2-Wasserstein loss with respect to distance matrices of the backbone atoms to
consider long-distance interactions. We demonstrated the feasibility of this
new flow model by successfully training it on medium-sized proteins. The model
generated interesting results by sampling new metastable states that are
difficult to obtain through conventional MD simulations.

We envision further improvement to scalability may be possible via various
means. The success of the 2-Wasserstein loss motivates further exploration of
spatial loss functions such as using distance matrices or contact maps. Li et
al. (2023) showed that the native folding conformations of proteins can be
more accurately sampled when backbone bond angles are conditioned on backbone
dihedral angles, which could further simplify our representation. One primary
limitation of our work is the lack of transferability between molecular
systems. Another limitation is that, as a normalizing flow model, the model
has a large number of parameters. Conditional diffusion-based approaches are a
promising direction to address these limitations.

## References

  * Beauchamp et al. (2012) Kyle A. Beauchamp, Robert McGibbon, Yu-Shan Lin, and Vijay S. Pande.  Simple few-state models reveal hidden complexity in protein folding.  _The Proceedings of the National Academy of Sciences_ , 109(44):17807–17813, 2012.
  * Case et al. (2014) D. A. Case, V. Babin, J. T. Berryman, R. M. Betz, Q. Cai, D.S. Cerutti, T.E. Cheatham III, T.A. Darden, R. E. Duke, H. Gohlke, A. W. Goetz, S. Gusarov, N. Homeyer, P. Janowski, J. Kaus, I. Kolossváry, A. Kovalenko, T. S. Lee, S. LeGrand, T. Luchko, R. Luo, B. Madej, K. M. Merz, F. Paesani, D. R. Roe, A. Roitberg, C. Sagui, R. Salomon-Ferrer, G. Seabra, C. L. Simmerling, W. Smith, J. Swails, R. C. Walker, J. Wang, R. M. Wolf, X. Wu, and P. A. Kollman.  Amber 14, 2014.
  * Cornish et al. (2019) Rob Cornish, Anthony L. Caterini, George Deligiannidis, and Arnaud Doucet.  Relaxing bijectivity constraints with continuously indexed normalising flows, 2019.
  * Derrick & Wigley (1994) Jeremy P. Derrick and Dale B. Wigley.  The third igg-binding domain from streptococcal protein g. an analysis by x-ray crystallography of the structure alone and in a complex with fab.  _Journal of Molecular Biology_ , 243(5):906–918, 1994.
  * Dibak et al. (2022) Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé.  Temperature steerable flows and boltzmann generators.  _Phys. Rev. Research_ , 4(L042005), 2022.
  * Dinh et al. (2017) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  In _International Conference on Learning Representations_ , 2017.
  * Durkan et al. (2019) Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  Neural spline flows.  In _NeurIPS_ , 2019.
  * Fu et al. (2023) Xiang Fu, Tian Xie, Nathan J. Rebello, Bradley D. Olsen, and Tommi Jaakkola.  Simulate time-integrated coarse-grained molecular dynamics with multi-scale graph networks.  _Transactions on Machine Learning Research_ , 2023.
  * Grcić et al. (2021) Matej Grcić, Ivan Grubišić, and Siniša Šegvić.  Densely connected normalizing flows, 2021.
  * Heusel et al. (2017) Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/8a1d694707eb0fefe65871369074926d-Paper.pdf.
  * Holm (2020) Liisa Holm.  Using dali for protein structure comparison.  _Methods in Molecular Biology_ , pp. 29–42, 2020.
  * Holm & Sander (1993) Liisa Holm and Chris Sander.  Protein structure comparison by alignment of distance matrices.  _Journal of Molecular Biology_ , 233(1):123–138, 1993.
  * Hou et al. (2019) Jie Hou, Tianqi Wu, Renzhi Cao, and Jianlin Cheng.  Protein tertiary structure modeling driven by deep learning and contact distance prediction in casp13.  _Proteins Structure Function Bioinformatics_ , 87(12):1165–1178, 2019.
  * Hua et al. (2022) Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le.  Transformer quality in linear time, 2022.
  * Ichinomiya (2022) Takashi Ichinomiya.  Topological data analysis gives two folding paths in hp35(nle-nle), double mutant of villin headpiece subdomain.  _Scientific Reports_ , 12(1), 2022.
  * Invernizzi et al. (2022) Michele Invernizzi, Andreas Krämer, Cecilia Clementi, and Frank Noé.  Skipping the replica exchange ladder with normalizing flows.  _J. Phys. Chem. Lett._ , 13(50):11643––11649, 2022.
  * Jing et al. (2022) Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola.  Torsional diffusion for molecular conformer generation.  _arXiv preprint arXiv:2206.01729_ , 2022.
  * Kaynak et al. (2021) Burak T. Kaynak, She Zhang, Ivet Bahar, and Pemra Doruker.  Clustenmd: efficient sampling of biomolecular conformational space at atomic resolution.  _Bioinformatics_ , 37(21):3956–3958, 2021.
  * Klein et al. (2023a) Leon Klein, Andrew Y. K. Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Marc Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka.  Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics, 2023a.
  * Klein et al. (2023b) Leon Klein, Andreas Krämer, and Frank Noé.  Equivariant flow matching, 2023b.
  * Köhler et al. (2021) Jonas Köhler, Andreas Krämer, and Frank Noé.  Smooth normalizing flows.  In _Advances in Neural Information Processing Systems 34_ , 2021.
  * Köhler et al. (2022) Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé.  Flow-matching – efficient coarse-graining molecular dynamics without forces.  _arXiv preprint arXiv:2203.11167_ , 2022.
  * Kondor & Trivedi (2018a) Risi Kondor and Shubhendu Trivedi.  On the generalization of equivariance and convolution in neural networks to the action of compact groups.  In _ICML_ , 2018a.
  * Kondor & Trivedi (2018b) Risi Kondor and Shubhendu Trivedi.  3d steerable cnns: learning rotationally equivariant features in volumetric data.  In _NeurIPS_ , 2018b.
  * Li et al. (2023) Jie Li, Oufan Zhang, Seokyoung Lee, Ashley Namini, Zi Hao Liu, Joao Miguel Correia Teixeira, Julie D Forman-Kay, and Teresa Head-Gordon.  Learning correlations between internal coordinates to improve 3d cartesian coordinates for proteins.  _J. Chem. Theory Comput._ , 19(14):4689––4700, 2023.
  * Ma et al. (2023) Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer.  Mega: Moving average equipped gated attention, 2023.
  * Mahmoud et al. (2022) Amr H. Mahmoud, Matthew Masters, Soo Jung Lee, , and Markus A. Lill.  Accurate sampling of macromolecular conformations using adaptive deep learning and coarse-grained representation.  _Journal of Chemical Information and Modeling_ , 62(7), 2022.
  * Martino et al. (2017) Luca Martino, Víctor Elvira, and Francisco Louzada.  Effective sample size for importance sampling based on discrepancy measures.  _Signal Processing_ , 131:386–401, February 2017.  ISSN 0165-1684.  doi: 10.1016/j.sigpro.2016.08.025.  URL http://dx.doi.org/10.1016/j.sigpro.2016.08.025.
  * McInnes et al. (2018) Leland McInnes, John Healy, Nathaniel Saul, and Lukas Grossberger.  Umap: Uniform manifold approximation and projection.  _The Journal of Open Source Software_ , 3(29):861, 2018.
  * Midgley et al. (2023) Laurence I. Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato.  Se(3) equivariant augmented coupling flows, 2023.
  * Midgley et al. (2022) Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato.  Flow annealed importance sampling bootstrap.  _arXiv preprint arXiv:2208.01893_ , 2022.
  * Nguyen & Salazar (2019) Toan Q. Nguyen and Julian Salazar.  Transformers without tears: Improving the normalization of self-attention.  In _IWSLT_ , 2019.
  * Noé et al. (2019) Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu.  Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning.  _Science_ , 365(6457), 2019.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer, 2020.
  * Ramachandran et al. (2017) Prajit Ramachandran, Barret Zoph, and Quoc V. Le.  Searching for activation functions, 2017.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In Francis Bach and David Blei (eds.), _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 1530–1538, Lille, France, 07–09 Jul 2015. PMLR.  URL https://proceedings.mlr.press/v37/rezende15.html.
  * Rezende et al. (2020) Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer.  Normalizing flows on tori and spheres.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 8083–8092. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/rezende20a.html.
  * Senior et al. (2019) Andrew W. Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Žídek, Alexander W. R. Nelson, Alex Bridgland, Hugo Penedones, Stig Petersen, Karen Simonyan, Steve Crossan, Pushmeet Kohli, David T. Jones, David Silver, Koray Kavukcuoglu, and Demis Hassabis.  Protein structure prediction using multiple deep neural networks in the 13th critical assessment of protein structure prediction (casp13).  _Proteins Structure Function Bioinformatics_ , 87:1141–1148, 2019.
  * Shaw et al. (2010) David E. Shaw, Paul Maragakis, Kresten Lindorff-Larsen, Stafano Piana, Ron O. Dror, Michael P. Eastwood, Joseph A. Bank, John M. Jumper, John K. Salmon, and Willy Wriggers.  Atomic-level characterization of the structural dynamics of proteins.  _Science_ , 330(6002):341–346, 2010.
  * Shaw et al. (2018) Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani.  Self-attention with relative position representations, 2018.
  * Sittel et al. (2017) Florian Sittel, Thomas Filk, and Gerhard Stock.  Principal component analysis on a torus: Theory and application to protein dynamics.  _J. Chem. Phys._ , 147(244101), 2017.
  * Su et al. (2021) Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  _arXiv preprint arXiv:2104.09864_ , 2021.
  * Thomas et al. (2018) Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley.  Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds, 2018.
  * Vaidehi & Jain (2015) Nagarajan Vaidehi and Abhinandan Jain.  Internal coordinate molecular dynamics: A foundation for multiscale dynamics.  _J. Phys. Chem. B_ , 119(4):1233–1242, 2015.
  * Wang et al. (2022) Yihang Wang, Lukas Herron, and Pratyush Tiwary.  From data to noise to data for mixing physics across temperatures with generative artificial intelligence.  _The Proceedings of the National Academy of Sciences_ , 119(32), 2022.
  * Wirnsberger et al. (2022) Peter Wirnsberger, George Papamakarios, Borja Ibarz, Sébastien Racanière, Andrew J. Ballard, Alexander Pritzel, and Charles Blundell.  Normalizing flows for atomic solids.  _Machine Learning: Science and Technology_ , 3(2), 2022.
  * Wu et al. (2020) Hao Wu, Jonas Köhler, and Frank Noé.  Stochastic normalizing flows.  In _NeurIPS_ , 2020.
  * Wu et al. (2022) Kevin E. Wu, Kevin K. Yang, Rianne van den Berg, James Y. Zou, Alex X. Lu, and Ava P. Amini.  Protein structure generation via folding diffusion.  _arXiv preprint arXiv:2209.15611_ , 2022.
  * Xu & Wang (2019) Jinbo Xu and Sheng Wang.  Analysis of distance-based protein structure prediction by deep learning in casp13.  _Proteins Structure Function Bioinformatics_ , 87(12):1069–1081, 2019.
  * Zhu et al. (2023) Kun Zhu, Hong Su, Zhenling Peng, and Jianyi Yang.  A unified approach to protein domain parsing with inter-residue distance matrix.  _Bioinformatics_ , 39(2), 2023.

##  Appendix A Related works

###  A.1 Boltzmann Generators

Bolzmann generators (Noé et al., 2019) are normalizing flows that approximate
Boltzmann distributions. Noé et al. (2019) utilized the fact that normalizing
flows are tractable density models and introduced a notion of training by
energy via reverse KL-divergence minimization. Recently, there has been a
growing interest in Boltzmann generators. Dibak et al. (2022) proposed
temperature steerable flows that generalized to families of ensembles across
multiple temperatures, thereby increasing the range of thermodynamic states
accessible for sampling. Unfortunately, this model tends to undersample
significant local minimas for systems as small as alanine dipeptide. The
authors believed that this was due to the limited expressivity of the flow
model. Wu et al. (2020) proposed stochastic normalizing flows, which combine
flows with MCMC methods by introducing sampling layers between flow layers to
improve model expressivity. Unfortunately, this method is computationally
expensive as it involves many more target evaluations. In addition, stochastic
normalizing flows tend to miss modes (Midgley et al., 2022). Köhler et al.
(2021) introduced smooth normalizing flows, which are $C^{\infty}$-smooth,
thus making them more physically amenable. They also introduce force-matching
as an added loss term. While they have impressive results and modal coverage
for alanine dipeptide, they utilize a root-finding algorithm to approximate
the inverse for their smooth flows, which becomes computationally prohibitive
for higher-dimensional systems.

This work has focused on normalizing flows. However, diffusion models have
also shown great promise as an alternative generative model for learning
Boltzmann generators. Jing et al. (2022) train a diffusion model to learn the
Boltzmann distribution over the torsion angles of multiple drug-like
molecules, while using cheminformatics methods for the bond lengths and
angles. They perform energy-based training (similar in spirit to the reverse
KL divergence in flow model training) via estimation of a score matching loss
using samples generated by the model. However, this method does not scale well
to larger molecules and inherits the same problem of unstable training at
initialization.

###  A.2 Loss functions

Wirnsberger et al. (2022) trained a flow model without MD samples by
minimizing the KL divergence to approximate the Boltzmann distribution of
atomic solids with up to 512 atoms. However, the KL-divergence suffers from
mode-seeking behavior, which severely impairs training for multimodal target
distributions. While the forward KL-divergence, i.e. maximum likelihood, is
mass covering, the Monte Carlo approximations of such an objective have a very
high variance in loss. To circumvent this, Midgley et al. (2022) trains a flow
to approximate a target $p$ by minimizing the alpha-divergence with
$\alpha=2$, which is estimated with annealed importance sampling (AIS) using
the flow $q$ as the base distribution and $p^{2}/q$ as target. This method is
notable in that it does not require any MD samples but still achieves
impressive results for alanine dipeptide. Nonetheless, the AIS component is
computationally expensive and scales poorly for larger systems.

###  A.3 Coarse-graining

Several works have attempted to scale flow-based Boltzmann generators to
larger systems. Mahmoud et al. (2022) trained a flow model on coarse-grained
protein representations which they then mapped back to full-atom
representations using a language model. On a similar note, Köhler et al.
(2022) trained a normalizing flow to represent the probability density for
coarse-grained (CG) MD samples in order to learn the parameters of a CG model.
Unfortunately, coarse-grain approaches tend to lose significant information
compared to full-atom resolution for downstream applications. Importantly,
both works note that using internal-coordinate representations do not scale
well as small changes in torsion angles can lead to large global distortions.
Our results indicate that this is not necessarily true, as we use a reduced
internal-coordinate representation.

###  A.4 Normalizing flow architectures

Our flow model, while novel, shares some similarities to previous works.
DenseFlow (Grcić et al., 2021) fuses a densely connected convolutional block
with Nyström self-attention in modules with both cross-unit and intra-module
couplings. This architecture is specifically designed for image data and
utilizes a linear approximation for the self-attention mechanism. In contrast,
we use gated attention and rotary positional embeddings in order to handle the
sequential nature of proteins.

Multiscale flow architectures were first introduced by Dinh et al. (2017) In
the protein domain, previous works also split the inputs into different
channels (Noé et al., 2019; Köhler et al., 2021; 2022). However, they split
the input dimensions into torsion, angle, and bond channels. In contrast, our
model splits the input into separate backbone and sidechain channels to better
capture the global distribution.

###  A.5 Transferable models

As mentioned in the discussion section of the main text, one of the primary
limitations of this work is the inability to transfer across molecular
systems. Several works have attempted to overcome this limitation. Klein et
al. (2023a) developed Timewarp: an enhanced sampling method which uses a
normalising flow as a proposal distribution in a Markov Chain Monte Carlo
(MCMC) method targeting the Boltzmann distribution. However, the
transferability of Timewarp is demonstrated only for small peptides (2-4 amino
acids), and its capabilities are yet to be validated on larger systems. One
promising direction for developing transferable models targeting the Boltzmann
distribution is diffusion modeling. Jing et al. (2022) develop a torsion score
model that allows for transferability across systems. However, their model is
only trained and validated on small, drug-like molecules that are around the
same size as alanine dipeptide or smaller. Fu et al. (2023) trained a multi-
scale graph neural network that directly simulates coarse-grained MD with a
very large time step and used a diffusion model as a refinement module to
mitigate simulation instability. The degree of coarse-graining as presented in
the paper diminishes the resolution, thereby making downstream drug-design
applications infeasible. In addition, coarse-graining dynamics often do not
mimic real transitions that occur in nature for proteins.

###  A.6 Equivariant flow models

Several recent works have attempted to bring the benefits of equivariance
(particularly SE(3) equivariance) (Thomas et al., 2018; Kondor & Trivedi,
2018a; b) to normalizing flow models. Two recent works, in particular, were
able to model the Boltzmann distribution for alanine dipeptide in Cartesian
coordinates. Midgley et al. (2023) develop an augmented coupling flow that
preserve SE(3) and permutation equivariance that can sample from the Boltzmann
distribution of alanine dipeptide via importance weighting. Klein et al.
(2023b) utilize a different generative modeling method called flow matching.
Specifically, they utilize equivariant flow matching to exploit the physical
symmetries of the Boltzmann distribution and achieve significant sampling
efficiency for alanine dipeptide. Unfortunately, both works still fall short
of internal coordinate-based methods for alanine dipeptide. However, as a
Cartesian coordinate representation is more generalizable and will often
present with smoother gradients and more stable training, they are a promising
direction for developing scalable BGs.

##  Appendix B Training by energy

Below, we show the connection between minimizing the reverse KL-divergence and
minimizing the energy of generated samples.

$\displaystyle KL(q_{\theta}||p)$ $\displaystyle=\mathbb{E}_{\mathbf{x}\sim
q_{\theta}}\left[\log q_{\theta}(\mathbf{x})-\log p(\mathbf{x})\right]$
$\displaystyle=\mathbb{E}_{\mathbf{z}\sim q}\left[\log
q(\mathbf{z})-\log|\text{det}(J_{f_{\theta}}(\mathbf{z}))|-\log
p(f_{\theta}(\mathbf{z}))\right]$ $\displaystyle=-H_{\mathbf{z}}+\log
C+\mathbb{E}_{\mathbf{z}\sim
q}\left[u(f_{\theta}(\mathbf{z}))-\log|\text{det}(J_{f_{\theta}}(\mathbf{z}))|\right],$

where $H_{\mathbf{z}}$ is the entropy of the random variable $\mathbf{z}$ and
$C=\int e^{-u(\mathbf{x})/(kT)}d\mathbf{x}$ is the normalization constant for
the Boltzmann distribution $p(\mathbf{x})\propto e^{-u(\mathbf{x})/(kT)}$.
When minimizing the KL-divergence with respect to the parameters $\theta$, the
entropy term and the log normalization constant disappear as they are not
dependent on $\theta$:

$\displaystyle\theta^{*}$
$\displaystyle=\operatorname*{arg\,min}_{\theta}-H_{\mathbf{z}}+\log
C+\mathbb{E}_{\mathbf{z}\sim
q}\left[u(f_{\theta}(\mathbf{z}))-\log|\text{det}(J_{f_{\theta}}(\mathbf{z}))|\right]$
$\displaystyle=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{\mathbf{z}\sim
q}\left[u(f_{\theta}(\mathbf{z}))-\log|\text{det}(J_{f_{\theta}}(\mathbf{z}))|\right].$

The expectation here is usually approximated with a Monte Carlo estimate, but
a variety of different sampling procedures can be utilized. The log
determinant Jacobian (ldj) term can be seen as promoting entropy, or
exploration, of the sample space.

##  Appendix C Coordinate Transformation

###  C.1 Protein Structure

Protein structure refers to the three-dimensional arrangement of atoms in an
amino acid-chain molecule. There are four distinct levels by which we can
describe protein structure. The primary structure of a protein refers to the
sequence of amino acids in the polypeptide chain. The secondary structure
refers to regularly patterned local sub-structures on the actual polypeptide
backbone chain. The two most common secondary structure motifs are
$\alpha$-helices and $\beta$-sheets. Tertiary structure refers to the overall
three-dimensional structure created by a single polypeptide. Tertiary
structure is primarily driven by non-specific hydrophobic interactions as well
as long-range intramolecular forces. Quaternary structure refers to the three-
dimensional structure consisting of two or more polypeptide chains that
operate as a single functional unit.

###  C.2 Coordinate Representations

Figure S.1: An illustration of the definition of bond length, bond angle, and
dihedral angle by four atoms. Subscripts indicate the atoms that define the
value, where order is given by the bond graph connectivity. In internal
coordinate system, the position or Cartesian coordinate of atom 4 is
determined by atom 1,2 and 3 based on bond length, bond angle and dihedral
angle.

Boltzmann generators usually do not operate directly with Cartesian
coordinates. The primary global conformational changes of a protein do not
described efficiently by the atomic Cartesian coordinates. This is driven by
the fact that chemical bonds are very stiff, and energetically-favored
conformational changes take place via rotations around single chemical bonds
(Vaidehi & Jain, 2015). A more commonly used alternative is internal
coordinates. Internal coordinates are defined by bond lengths $d$, bond angles
$\theta$, and dihedral angles $\phi$ (Fig. S.1).

In their seminal work, Noé et al. (2019) introduced a coordinate
transformation whereby the protein backbone atoms (primarily defined as the
$N$, $C_{\alpha}$, and $C$ atoms) are mapped PCA coordinates while the rest of
the atoms are mapped to internal coordinates. The motivation behind this mixed
coordinate transformation is that protein conformations are highly sensitive
to changes in backbone internal coordinates. This often results in unstable
training and difficulty in generating natural, i.e., high Boltzmann
probability, structures. Most works since have used full internal coordinate
representations but experimented only with small systems, the most common of
which is alanine dipeptide (22 atoms). Köhler et al. (2022) note that scaling
Boltzmann generators to larger systems is difficult with internal coordinate
representations.

In our work, rather than using a full internal coordinate representation,
which would be $3N-6$ dimensional (where $N$ is the number of atoms in the
system), we utilize a reduced internal coordinate representation. For training
features, we use the dihedral angles and the bond angles for the 3 backbone
atoms ($N$, $C_{\alpha}$, $C$). For side-chain atoms, we use all rotatable
dihedral angles around single bond. All bond lengths and bond angles other
than the 3 defining backbone atoms and improper torsion angles are kept at
their mean values calculated from input protein structures. By examining all
protein structures generated, we confirmed that such a reduced internal
coordinate system can represent all protein structures to very high accuracy
and quality. Recent works have adopted similar approaches; Wu et al. (2022)
utilize only the backbone torsion and bond angles to represent various
proteins, while Wang et al. (2022) simply use the backbone torsion angles to
represent the polypeptide AiB9.

##  Appendix D Training details and architecture

Table S.1: Hyperparameters for training Optimizer | AdamW
---|---
$\lambda$ | 0.0001
Learning rate | 0.002
Scheduler | ReduceLROnPlateau
Patience (epochs) | 5
Factor | 0.1
Batch Size | 256
Dropout | 0.1
$Q_{\text{dim}}$ | 32
$K_{\text{dim}}$ | 32
$V_{\text{dim}}$ | 64
Normalization | Scaled
Attention | Laplace
Activation | SiLU
RQS bins | 8
Epochs (NLL) | 200
Epochs (NLL + W) | 50
Epochs (NLL + W + KL) | 20
Epochs (NLL + KL) | 10

All models were trained on a single NVIDIA A100 GPUs with the Adam optimizer
and a dropout factor of 0.1. For model that utilized GAU-RQS blocks, the
dimensionalities of the $Q$, $K$, and $V$ matrices were 32, 32, and 64,
respectively. In addition, we utilized scaled normalization (Nguyen & Salazar,
2019), the Laplace attention function (Ma et al., 2023), and SiLU activations
(Ramachandran et al., 2017). For the gated attention units, we also use the T5
relative positional bias (Raffel et al., 2020). For the rational quadratic
splines (RQS), we use a bin size of $K=8$.

Data was all standard normalized. Dihedral angles were constrained to be
within $[-\pi,\pi]$ and shifted as done by Sittel et al. (2017).

For the multi-stage training strategy, all models were trained for 200 epochs
( 12 hours) with the NLL loss, 50 epochs ( 8 hours) with NLL+W, 20 epochs ( 8
hours) with NLL+W+KL, and 10 epochs ( 3 hours) with NLL+KL. The approximate
times are for protein G, which has 56 residues.

We do no hyperparameter tuning due to the lack of compute and time. Further
implementation details are given in the code, which is available upon request.
A summary of the hyperparameters for our model are provided in Table S.1.

##  Appendix E Further ablations

In this section, we provide further ablations for Table 1 in the main text. In
particular, we provide further ablations with regards to the training strategy
with the baseline neural spline flows (NSF) architecture in Table S.2. As we
can see from the table, while our different training strategies improve upon
the baseline model with NLL training, the improvements are not as drastic as
for our architecture (Table 1).

Table S.2: Training NSF basline model with different strategies. We compute $\Delta D$, energy $u(\cdot)$, and mean NLL of $10^{6}$ generated structures after training with different training strategies for the baseline NSF model with ADP, protein G, and Villin HP35. |  | Training strategy |  |
---|---|---|---|---
System | NLL | KL | W2 |  $\Delta D$ (Å) | Energy $u(\mathbf{x})$ (kcal/mol) | $-\mathbb{E}_{p(\mathbf{x})}[\log q_{\theta}(\mathbf{x})]$
| ✓ |  |  | $0.09\pm 0.01$ | $(-1.19\pm 0.61)\times 10^{1}$ | $38.29\pm 0.19$
ADP | ✓ | ✓ |  | $0.08\pm 0.01$ | $(-1.21\pm 0.48)\times 10^{1}$ | $40.11\pm 0.20$
| ✓ |  | ✓ | $0.05\pm 0.01$ | $(-0.99\pm 0.57)\times 10^{1}$ | $41.03\pm 0.08$
| ✓ | ✓ | ✓ | $0.04\pm 0.00$ | $(-1.22\pm 0.13)\times 10^{1}$ | $39.10\pm 0.13$
| ✓ |  |  | $2.92\pm 0.80$ | $(2.15\pm 3.31)\times 10^{10}$ | $-263.46\pm 0.13$
Protein G | ✓ | ✓ |  | $18.19\pm 2.88$ | $(2.90\pm 0.82)\times 10^{2}$ | $-260.87\pm 0.51$
| ✓ |  | ✓ | $1.81\pm 0.33$ | $(6.04\pm 3.79)\times 10^{7}$ | $-261.01\pm 0.33$
| ✓ | ✓ | ✓ | $1.58\pm 0.29$ | $(-0.86\pm 2.04)\times 10^{2}$ | $-257.82\pm 0.92$
| ✓ |  |  | $0.81\pm 0.06$ | $(7.78\pm 17.4)\times 10^{7}$ | $687.95\pm 1.92$
HP35 | ✓ | ✓ |  | $0.91\pm 0.05$ | $(2.15\pm 11.4)\times 10^{3}$ | $691.41\pm 1.47$
| ✓ |  | ✓ | $0.59\pm 0.05$ | $(9.61\pm 2.55)\times 10^{7}$ | $690.87\pm 1.05$
| ✓ | ✓ | ✓ | $0.61\pm 0.07$ | $(-1.77\pm 1.49)\times 10^{2}$ | $691.10\pm 2.12$

##  Appendix F Reweighted distribution

Table S.3: Effective sample size (ESS) of the various training strategies and architectures. |  | Training strategy |
---|---|---|---
System | Arch. | NLL | KL | W2. | ESS (%)
| NSF | ✓ |  |  | $3.9\pm 0.5$
|  | ✓ |  |  | $9.1\pm 1.7$
ADP | Ours | ✓ | ✓ |  | $58.6\pm 8.4$
|  | ✓ |  | ✓ | $39.4\pm 2.7$
| Ours | ✓ | ✓ | ✓ | $\mathbf{88.4\pm 0.2}$
| NSF | ✓ |  |  | $0.0\pm 0.0$
|  | ✓ |  |  | $0.0\pm 0.0$
Protein G | Ours | ✓ | ✓ |  | $0.0\pm 0.0$
|  | ✓ |  | ✓ | $0.0\pm 0.0$
| Ours | ✓ | ✓ | ✓ | $\mathbf{62.47\pm 1.4}$
| NSF | ✓ |  |  | $0.0\pm 0.0$
|  | ✓ |  |  | $0.0\pm 0.0$
HP35 | Ours | ✓ | ✓ |  | $0.0\pm 0.0$
|  | ✓ |  | ✓ | $0.0\pm 0.0$
| Ours | ✓ | ✓ | ✓ | $\mathbf{43.9\pm 1.3}$

Typically, the output distribution of the flow model will not match exactly
with the target distribution, and previous works employ importance sample
reweighting to the target distribution (Noé et al., 2019; Midgley et al.,
2022; Wu et al., 2020). While efficient Boltzmann reweighting is feasible for
a small system like alanine dipeptide, the current work makes several modeling
assumptions to scale to larger molecules. We model a distribution on a space
with reduced dimensionality, which is not exactly the Boltzmann distribution.
To be precise, we model $p(\mathbf{\tau},\mathbf{\theta}_{bb}|L=\bar{L})$,
where $\mathbf{\tau}$ are torsion angles, $\mathbf{\theta}_{bb}$ denotes
backbone bond angles, and $L$ and $\bar{L}$ denote other internal coordinates
and their mean marginals, respectively. In addition, while the MD simulations
for the training data were trained with explicit water, we use an implicit
water model (for efficiency) when training with the energy function. This
results in a wider range of energy values for our training data and generated
samples for protein G and villin HP35 (the alanine dipeptide data is open
source and is run with an implicit solvent). Due to precision issues, it is
difficult to meaningfully compare importance weights as only the lowest energy
structures will tend to have a nonzero importance weight. For these reasons,
we consider the histogram of our training data distribution for
$e^{-u(\mathbf{x})}$, and use the bins and densities to define
$p_{data}(\mathbf{x})$, our target distribution. We define $q(\mathbf{x})$ as
the likelihood according to our flow model.

We report the effective sample size (ESS) (Martino et al., 2017) (Table S.3)
and display the reweighted energy distribution according to the energies
computed for the training data distribution of protein G and villin HP35 (Fig.
S.2). As we can see from the table, for larger systems such as protein G and
HP35, only the model that utilizes our novel architecture and multi-stage
training strategy are capable of capturing a meaningful subset of the data
distribution, as measured by ESS. This is primarily due to the fact that the
other models generate samples with atomic clashes that dramatically increase
their associated energies.

To remedy some of the issues with using the Boltzmann distribution as our
target distribution, we computed the energies of the generated structures and
the training data with force field parameters that more closely modeled the
simulating force field (specifically, we use the GBn2 implicit solvent model).
We then set the target distribution as $p(\mathbf{x})\propto
e^{-u(\mathbf{x})}$. We conduct importance sampled reweighting and display the
results for protein G in Fig S.3. As we can see, Boltzmann reweighting tends
to sample only for the lowest energy states. In fact, samples from the
training data would rarely have nonzero weights. This motivated our previous
approach for Fig. S.2.

Figure S.2: Energy distribution of the training data (orange), samples
generated from the model (blue), and the importance weight resampled energy
distribution from the flow model (green) for protein G and HP35. The target
distribution for importance weighting is set as the histogram distribution of
$p_{data}(\mathbf{x})\propto e^{-u(\mathbf{x})}$. Energy here is computed in a
vacuum. Figure S.3: Energy distribution of the training data (orange), samples
generated from the model (blue), and the importance weight resampled energy
distribution from the flow model (green) for protein G. We modify the force
field for computing the energies to be closer to the simulating force field.
The target distribution for importance weighting is the Boltzmann probability
$p(\mathbf{x})\propto e^{-u(\mathbf{x})}$.