# Coarse Grained Molecular Dynamics with Normalizing Flows

Samuel Tamagnone1 Alessandro Laio1,2 Marylou Gabrié3

(1International School for Advanced Studies (SISSA), Via Bonomea 265, Trieste,
Italy
2The Abdus Salam International Centre for Theoretical Physics (ICTP), Strada
Costiera 11, Trieste, Italy
3CMAP, CNRS, École polytechnique, Institut Polytechnique de Paris, 91120
Palaiseau, France )

###### Abstract

We propose a sampling algorithm relying on a collective variable (CV) of mid-
size dimension modelled by a normalizing flow and using non-equilibrium
dynamics to propose full configurational moves from the proposition of a
refreshed value of the CV made by the flow. The algorithm takes the form of a
Markov chain with non-local updates, allowing jumps through energy barriers
across metastable states. The flow is trained throughout the algorithm to
reproduce the free energy landscape of the CV. The output of the algorithm are
a sample of thermalized configurations and the trained network that can be
used to efficiently produce more configurations. We show the functioning of
the algorithm first on a test case with a mixture of Gaussians. Then we
successfully test it on a higher dimensional system consisting in a polymer in
solution with a compact and an extended stable state separated by a high free
energy barrier.

##  1 Introduction

Sampling from the Boltzmann distribution is a fundamental task to understand
the equilibrium properties of physical systems. For molecular systems,
molecular dynamics (MD) simulations mostly relies on local-updates inspired by
physical dynamics. Yet, in the presence of local free energy minima caused by
energetic or entropic barriers, these simulations fail to sample from the
Boltzmann distribution of metastable systems as it takes an exponential time
for physical dynamics to escape such minima.

To tackle this issue, one strategy is to resort to enhanced sampling
techniques. Some of these approaches are based on the knowledge of a low-
dimensional collective variable (CV) capturing the metastability which is used
to drive the sampling between metastable states [EVE04, RMZT02, KRB+95,
vEMB03, KPV96, CCHK89, DP01, LP02]. CV-based enhanced samplers have two
requirements. First, that for a fixed value of the CV, the remaining degrees
of freedom can be efficiently equilibrated. Second, that the CV space is small
enough to be explored efficiently. Meeting these requirements simultaneously
is not an easy task as restricting to a handful of dimension to summarize the
metastability of the system may be very difficult, if not simply inadequate.

A different line of work recently leveraged a class of generative models
called normalizing flows (NF) [PNR+21] to enhance the sampling. There, NF are
used to propose candidate configurations subsequently reweighed or
incorporated into a Markov chain with an accept-reject step; see the
pioneering works of [WWZ19, AKS19, NOKW19]. As the NF generates independent
configurations across the entire configurational space, NF-based samplers are
agnostic to energetic and entropic barriers and therefore particularly
advantageous for metastable systems. However, NFs were found to hit a limit in
the dimension and complexity of systems they can handle [DDRW21, MMLL22,
CTWZ23, GDMG23].

In this work, we present a Markov chain Monte Carlo (MCMC) sampling algorithm
combining the use of a (high-dimensional) CV and a NF.

At each iteration, a refreshed CV value is proposed by a NF and a new
configuration is constructed given the refreshed CV. The new configuration is
accepted or rejected following the framework of non-equilibrium candidate
Monte Carlo [NCMC11, CR15], which guarantees that the stationary measure of
the defined Markov Chain is the target Boltzmann distribution. The NF is
trained adaptively along the MCMC to approach the image of the Boltzmann
measure through the CV-map – related to the CV free energy –. This learning
procedure entirely avoids the requirement for a dataset of thermalized
configurations beforehand and improves the NF proposals along the way. By
requiring the NF to learn a distribution in CV space rather than the
configurational space, the method benefits from the exploration power of these
generative models while easing its training by reducing dimension. The
proposed method also revisits the idea of a CV-driven sampler while
significantly relaxing the constraint on the dimension of the CV, precisely
because the NF can drive jumps between metastable states in CV-space. Thus,
the CV dimension must be smaller than the total number of degrees of freedom
but can be significantly larger than what is allowed by a traditional CV-
enhanced sampler. We refer to a regime of _mid-size_ for the dimension of CV
handled by the proposed algorithm.

A popular choice of a mid-size CV for a molecule is a coarse-grained
representation of the chain of atoms. Coarse-grained simulations of molecular
systems have brought significant insights on large molecular systems which are
unreachable by full-atomistic simulations, in particular for events happening
on large-time scales [Cle08, PV18]. However, designing a statistically
unbiased sampler for the coarsed-grained Boltzmann measure is highly non
trivial as it relates to the CV free energy, which is computationally
intractable due to high dimension of the CV. Using smart proposals of CV
values therefore requires to ‘backmap’ this coarsed-grained configuration to a
full-atomistic configuration on which the fine-grained energy function can be
evaluated so as to reweigh, or decide whether to accept, the proposal. This
‘backmapping’ remains a significant challenge. A line of recent works explored
the ability of deep generative models to help with this task [LBP+20, WXC+22,
MMLL22, CTR23, CR24]. Yet, backmapping strategies based on conditional NFs are
eventually limited by their already mentioned restricted expressiveness
[MMLL22, CTR23]. Conversely, methods involving other types of generative
models are typically more accurate but are crucially lacking the tractable
likelihood necessary to compute reweighing factors and acceptance
probabilities. Here, we take a different route by using an NF generative model
exclusively in the CV space and resorting to non-equilibrium dynamics to move
the remaining degrees of freedom [NCMC11, CR15].

The article is organized as follows. In section 2 we detail the proposed
algorithm. In section 3, we first demonstrate our method on a control example
of mixture of Gaussians showing the robustness of the method to dimension and
energy-barrier height. We also validate the method on a molecular system of a
9-beads polymer in a Lennard-Jones solvent, confirming here again the
robustness to energy-barrier height.

##  2 Algorithm description

We are interested in sampling a physical system with configurations
$x\in\mathbb{R}^{d}$ harvested from a Boltzmann distribution with energy
$U(x)$ and inverse temperature $\beta$:

$\displaystyle\rho_{*}(x)=e^{-\beta U(x)}/Z.$ (1)

We assume that $x$ can be decomposed into a (typically multidimensional)
collective variable (CV) $\psi\in\mathbb{R}^{n}$ and transversal degrees of
freedom $x_{\perp}\in\mathbb{R}^{d-n}$, that is
$x=(\psi,x_{\perp})$111Equivalently, we can also consider the situation where
there exists an invertible mapping $f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ such
that $x=f(\psi,x_{\perp})$. We denote by
$\Psi:\mathbb{R}^{d}\to\mathbb{R}^{n}$ the mapping $x\to\Psi(x)=\psi$. The
coordinate decomposition translates into a decomposition of the Boltzmann
distribution:

$\displaystyle\rho_{*}(x)=\rho_{*}(x|\psi)\rho^{\Psi}_{*}(\psi)$ (2)

with

$\displaystyle\rho^{\Psi}_{*}(\psi)=\int_{\mathbb{R}^{d}}\rho_{*}(x)\delta(\psi-\Psi(x))\mathrm{d}x=e^{-\beta
F_{*}(\psi)},$ (3)

where $\rho_{*}(x|\psi)$ is the density of the Boltzmann distribution (1)
conditioned on the CV value $\psi$ and we define the CV free energy
$F_{*}:\mathbb{R}^{n}\to\mathbb{R}$.

The proposed algorithm combines ideas from the non equilibrium candidate Monte
Carlo (NCMC) framework [NCMC11] and from the adaptive Monte Carlo augmented
with normalizing flows (NF), sometimes referred to as flowMC [GRVE22]. The
latter consists in an MCMC iterating three basics steps: Metropolis-adjusted
Langevin dynamics (MALA updates), non-local moves proposed by a NF (non-local
updates) and training steps of the NF to reproduce the distribution of
observed MCMC samples (NF training). The procedure is called adaptive as the
NF involved in the non-local update kernel learns from the previous samples
produced by the MCMC. The method exploits the two specific characteristic of
NF among the available generative models, they are easy to sample from and the
density of the probability distribution they represent is easily evaluated
[PNR+21].

Here, we follow the three steps of flowMC but amend the algorithm by using a
NF to propose non-local moves in a mid-sized CV space rather than in high-
dimensional configuration-space. From a refreshed CV value, a move in
configuration space is then proposed by steering the CV towards this value
while relaxing accordingly the transversal degrees. The steered proposal is
then accepted with a probability following the NCMC framework. The NF is
trained to reproduce the statistics of the CV values visited by the Markov
chain. The algorithm resulting from the iteration of MALA updates, steered
updates, for multiple parallel walkers, and NF training is summarized in
algorithm 1 and in fig. 2. We explain in detail the steered update and NF
training below.

### Steered update

The steered moves are decomposed in three steps. Assume we start from a
configuration $x=(\psi,x_{\perp})$. First the NF proposes the new CV value
$\psi^{\prime}$ independently from $x$; the flow density is denoted by
$\rho_{\theta}^{\Psi}(\psi)$ where $\theta$ stands for the trainable
parameters. Second, we update the CV in such a way that in $N$ steps it
reaches the target value. We present the algorithm assuming that this protocol
is a simple linear interpolation, but our derivation is valid also for a
generic protocol. In the linear interpolation case

$\displaystyle\psi^{t+1}=\psi^{t}+\frac{\Delta\psi}{N},$ (4)

where $\Delta\psi=\psi^{\prime}-\psi$. In each step of the interpolation we
perform MALA relaxation steps for the remaining degrees of freedom, with
proposals

$\displaystyle\widetilde{x}_{\perp}=x^{t}_{\perp}-\frac{\tau}{\gamma}\nabla_{x_{\perp}}U\left(\psi^{t}+\Delta\psi/2N,\,x^{t}_{\perp}\right)+\sqrt{\frac{2\tau}{\beta\gamma}}\eta$
(5)

where $\tau$ is a finite time step, $\gamma$ a friction coefficient,
$\nabla_{x_{\perp}}$ the gradient operator at fixed CV and $\eta$ a sample
from the standard normal distribution $\mathcal{N}(0,1)$. If the proposal
$\widetilde{x}_{\perp}$ is accepted – according to the Metropolis-Hasting
criteria for the conditional target $\rho_{*}(x_{\perp}|\psi^{t})$ – then
$x_{\perp}^{t+1}=\widetilde{x}_{\perp}$, otherwise
$x_{\perp}^{t+1}=x_{\perp}^{t}$; see algorithm S2 in the Supplementary
Information for a reminder of the MALA algorithm. Third, the full move
$(\psi,x_{\perp})\rightarrow(\psi^{\prime},x_{\perp}^{\prime})=(\psi^{N},x_{\perp}^{N})$
is accepted with probability $p=\min[\,1,\,R\,]$ with

$\displaystyle
R=\frac{\rho_{\theta}^{\Psi}(\psi)}{\rho_{\theta}^{\Psi}(\psi^{\prime})}\exp(-\beta
W).$ (6)

The first factor in $R$ accounts for a ratio of proposal probabilities, while
the second depends on the work performed on the system during the steered
move:

$\displaystyle
W=U(\psi^{\prime},x_{\perp}^{\prime})-U(\psi,x_{\perp})+\sum_{t=0}^{N-1}\Big{(}U(\psi^{t}+\Delta\psi/2,x_{\perp}^{t})-U(\psi^{t}+\Delta\psi/2,x_{\perp}^{t+1})\Big{)}.$
(7)

This procedure samples the Boltzmann distribution (1) as the acceptance
criterion satisfies detailed balance [NCMC11, Fre04]. A parallel work,
[SGLS60] discusses further the conditions for the reversibility of a CV-
steered update. In practice, the number of steps $N$ has implications on the
acceptance probability and the computational cost of each steered update. It
can be chosen in a system-dependent optimal way by running a preliminary study
on a series of steered moves and choosing, as a function of the steering
distance $\|\Delta\psi\|_{2}$ in CV space, the $N$ that yields the maximal
$p/N$ ratio, with $p$ the acceptance probability (6). Intuitively, the larger
the distance the more the transversal degrees of freedom are perturbed and the
higher is the optimal $N$ to use. The approximate law for the optimal
$N(\|\Delta\psi\|_{2})$ is then used at run time.

### NF training

At each iteration, the NF parameters are updated to learn the marginal
probability distribution on the collective variables. By analogy with flowMC
[GRVE22], we seek to minimize the Kullback-Leibler divergence between
$\rho_{*}^{\Psi}$ and $\rho_{\theta}^{\Psi}$,

$\displaystyle D_{\rm
KL}(\rho_{*}^{\Psi}\mid\mid\rho_{\theta}^{\Psi})=-\int_{\mathbb{R}^{n}}\log(\rho_{\theta}^{\Psi}(\psi))\rho_{*}^{\Psi}(\psi)d\psi+C_{*}$
(8)

where $C_{*}$ is a constant independent of the parameters $\theta$ of the NF.
For each gradient descent step, a batch of configurations
$\\{x^{k}\\}_{k=1}^{n_{\rm batch}}$ is randomly selected from the states
previously visited by the parallel walkers of the algorithm. The empirical
loss function then is the negative log-likelihood in CV space

$\displaystyle
L[\theta]=-\frac{1}{n}\sum_{i=k}^{n}\log(\rho_{\theta}^{\Psi}(\Psi(x^{k}))).$
(9)

In the literature of NF-assisted samplers, this maximum likelihood training
has been called “training by example”, as we use a set of states, which, at
convergence, are Boltzmann-distributed, to compute the loss function. We do
not perform “training by energy”, where the reverse Kullback-Leibler
divergence is minimized [RM15, AKS19, NOKW19, WWZ19]. For a NF defining a
density $\rho_{\theta}$ in full configuration space the reverse Kullback-
Leibler divergence $D_{\rm KL}(\rho_{\theta}\mid\mid\rho_{*})$ is an
expectation on $\rho_{\theta}$ and can be estimated by a Monte Carlo average
from the knowledge of the energy function $U$. Note that in our case this is
unfeasible as it would require evaluating the intractable free energy
$F_{*}(\psi)$ defined in (3).

Algorithm 1 CV-space adaptive Monte Carlo

$n_{\rm iter}$ number of iterations of the algorithm,
$\\{x_{i}^{0}(0)\\}_{i=1}^{n_{\rm walkers}}$ initial set of configurations,
$\beta$ inverse temperature, $U(\cdot)$ energy function, $\Psi(\cdot)$
collective variable mapping, $\theta$ initial NF parameters, $\epsilon$
learning rate, $n_{\rm batch}$ batch-size, $n_{\rm local}$ number of local
steps, $N(\cdot)$ steering-schedule number-of-steps rule, $\tau$ time step for
MALA, $\gamma$ friction coefficient.

$\mathcal{B}=\\{\\}$ $\triangleright$ initialize buffer

for $n=0,\dots,n_{\rm iter}$ do

for $i=1,\dots,n_{\rm walkers}$ do

$\\{x^{i,\alpha}(n)\\}^{\alpha=1,\dots,n_{\rm local}}={\rm
MALA}(U(\cdot),\beta,\gamma,\tau,x_{i}^{0}(n),n_{\rm local})$ $\triangleright$
update the walkers with local moves222In SI algorithms S2 and S3 the reader
may find the precise description of the STEERED update and of the MALA update.

$x^{i,0}(n+1)={\rm STEERED}(U(\cdot),\beta,\gamma,\tau,x_{i}^{n_{\rm
local}}(n),N(\cdot),\rho_{\theta}^{\Psi}(\psi))$ $\triangleright$ use NF for a
non-local move

$\mathcal{B}\leftarrow\mathcal{B}\cup\\{x^{i,0}(n+1)$}

end for

for $k=1,\dots n_{\rm batch}$ do

$x^{k}\sim\mathcal{U(\mathcal{B})}$ $\triangleright$ sample training examples
uniformly over the buffer

end for

$L[\theta]=\frac{-1}{n_{\rm batch}}\displaystyle\sum_{k=1}^{n_{\rm
batch}}\log(\rho_{\theta}^{\Psi}(\Psi(x^{k})))$ $\triangleright$ compute
negative log-likelihood

$\theta\leftarrow\theta-\epsilon\nabla_{\theta}L[\theta]$ $\triangleright$
make one stochastic gradient descent update on the NF parameters333One could
decide to do more than one training step each non-local move, as this is the
more computationally costly part of the simulation. This is what we have done
in our numerical examples.

end for

return $\theta$ trained NF parameters,
$\\{x^{i,\alpha}(n)\\}_{i=1,\dots,n_{\rm walkers}}^{n=1,\dots,n_{\rm iter};\
\alpha=0,\dots,n_{\rm local}}$ all visited configurations444Typically, one
does not save the configurations reached at all the local steps, but only a
sub-sample of those, as they can be strongly correlated.

##  3 Numerical results

### Mixture of Gaussians

As a first proof of concept, we test our procedure on mixture of Gaussian
distributions. We show that the NF correctly learns the free energy of the
system and that, at convergence, the average acceptance of the steered moves
does not depend on the height of the free energy barrier in CV space. As a
consequence, if the free energy barrier depends only weakly on the transversal
degrees of freedom, namely if the CV are well chosen, the cost of the
simulation is approximately independent of the barrier height.

We first consider a mixture of two Gaussians in 3D. Here, the full states are
described by vectors $x=(x_{0},x_{1},x_{2})\in\mathbb{R}^{3}$ and the chosen
CVs are the first two coordinates $x_{0}$ and $x_{1}$
($\Psi(x)=(x_{0},x_{1})$), leaving one transversal degree of freedom $x_{2}$.
The example is set up so that the CVs allow distinguishing between the two
different modes (see Materials & Methods below for details). The results are
presented in fig. 1. In panel a) we show the true 2D free energy of the
mixture that we are reproducing with the algorithm. In b) we show the free
energy obtained by the states saved alongside the CV-MC iterations. In panels
c) and d) we show the free energy learnt along the proposed adaptive algorithm
by two different NF architectures: the Rational Quadratic Spline (RQS)
[DBMP19] in c) and the Real NVP (RNVP) [DSB17] in d). As shown in panels a),
b), c) and d) there is a good agreement between the target free energy, the
learnt one and the one computed with the output of the CV-space adaptive Monte
Carlo. Because of its smoothness, the normalizing flow has a tendency to
connect the two separated modes. This occurs especially for the RNVP, as shown
in d), because it is less expressive than the RQS. This small systematic error
in the free energy learned by the flow is well corrected by the accept/reject
step in the non-local updates. Indeed the CV of the samples produced by the
algorithm are distributed correctly. In panel f), we further check that the
cumulative distribution of the marginal law on $x_{0}$ for the target and
simulated distribution (with the RQS) match almost perfectly. In panel e) we
show the behavior of the acceptance and of the loss as a function of the
algorithm’s iterations (again using RQS). As the NF learns the underlying free
energy, the algorithm becomes more and more efficient in making proposals and
the acceptance grows. As the plot shows, 2000 training iterations are already
enough to reach convergence.

We then consider the same mixture of two Gaussians example, but increase the
free energy barrier by pulling further apart the mixture components. The plots
in panel g) illustrate the most important advantage of the procedure: the
average acceptance at convergence depends very weakly on the free energy
barrier. The small dependence is caused only by the slightly varying ability
of the NF to learn the underlying probability. Indeed, when the proposal is
the exact target marginal probability (green line), as if the NF has perfectly
learnt the underlying marginal, the acceptance of NF-proposed moves is
constant. We have repeated the test with the two different NF architectures
mentioned above: the RNVP, which is less expressive, and the more expressive
RQS. The decrease in the acceptance that we have with the RNVP (dark blue
line) follows exactly the increase in the loss (orange line), which means that
this decrease is caused by proposals that become less and less accurate
because the RNVP has more difficulties learning modes that are further apart,
as it involves greater non-linearities. This does not occur with the RQS,
which is more expressive. Hence we have decided to use that architecture in
the tests on the polymer model presented below.

Finally, in panel h) we repeat the same numerical experiment adding more
transversal degrees of freedom (up to 28), but keeping the same 2D free energy
landscape. As shown by the average loss at convergence in CV space but also by
the average acceptance in full space, the result is independent of the number
of transversal dimensions if the CVs are properly selected.

Figure 1: a), b), c) and d) maps of the true, output of the CV-MC, learnt by
RQS or RNVP free energy landscapes. e) acceptance and loss function of a test
case with the RQS NF architecture during the first 2k steps of the algorithm.
f) comparison of the cumulative distribution on x for the target distribution
and the CV-MC one for checking if the weights of the modes are correctly
reproduced. g) performance in the learning task obtained with different
architectures and different energy barriers. h) acceptance and loss at
convergence for different numbers of transversal degrees of freedom.

### A model polymer in a Lennard-Jones solvent

We then benchmarked the algorithm on a more realistic and challenging test
system, a polymer of 9 beads immersed in a bath of particles. The interactions
considered are: a) A Lennard-Jones (LJ) potential acting among the particles
with coupling $\epsilon=1$, b) a LJ interaction between polymer and particles
with a smaller coupling $\epsilon=0.5$ (this is to favor configurations in
which the polymer is compact ), c) a LJ interaction between polymer beads with
$\epsilon=1$. The chosen LJ radius is always $2^{\frac{1}{6}}\sigma$ with
$\sigma=1$. To control the level of metastability we also include in the
energy function a double well potential $V(d)$, where $d$ is the end-to-end
distance of the polymer (see sketch in fig. 2). The barrier height can be
tuned by changing $V(d)$. With this setup the polymer has two metastable
states: a compact one and an extended one, with different weights due to
entropic contributions and solvation effects. We test two settings for the
barrier between the two states by changing the barrier in $V(d)$ from
$7\,K_{\rm B}T$ to $28\,K_{\rm B}T$. For the low barrier case, we produced a
MD trajectory of $10^{9}$ steps that we used as reference. For the high
barrier case we estimated the unbiased probability distribution by umbrella
sampling (see Methods).

We choose as collective variables the set of internal distances of the
polymer, namely the distance $\ell_{1}$ between monomer 1 and monomer 4, the
distance $\ell_{2}$ between monomer 2 and monomer 5 and so on up to
$\ell_{6}$. Therefore we have a 6-dimensional CV space. These distances fully
describe the conformation of the polymer, and have the same information
content as the dihedral angles, given the bond-lengths and the angles between
adjacent monomers are fixed. In fig. 2 we present a schematic representation
of algorithm 1 applied to this test system.

Figure 2: Cartoon representing the functioning of the algorithm. At each
iteration, the state of the system is updated through some local steps with a
sampler, MALA in our case. Then the next step consists in a steered update:
the NF proposes a new value of the internal distances, the polymer is then
stretched or compacted through steered MD and the solvent is relaxed
accordingly. The newly obtained configuration is rejected or accepted
according to the work performed and the ratio of NF-proposal probabilities.
After the local and non local update of the system, the NF undergoes some
training steps.

In fig. 3 we show the results. As mentioned above, the free energy landscape
of this test case is in 6D, with two minima, one corresponding to the compact
state, the other to the extended state. These two minima cannot be spotted in
any marginal of the probability density as a function of a single CV, but
appear on 2D marginals. In panel a) we show some of these 2D marginals for the
low barrier case, estimated from an ensemble of 100,000 independent
configurations. On each column, the samples were produced in a different way.
The first corresponds the reference MD trajectory, which is long enough to
have a large number of transitions from one minimum to the other. The second
one is the visited configuration of algorithm 1, while the third is obtained
by sampling in one shot from the NF at the end of the simulation, therefore
without the Monte Carlo correction of the visited configurations. One can
notice that MD and CV-MC marginals are more alike than MD and pure NF
marginals, but in general there is a good agreement between the three
marginals, showing that the algorithm works well.

In panel d) and e) we provide a quantitative measure of the difference in the
free energies in 6D estimated by the different approach, for the low and high
barrier case. For this, we used the PAk density estimator from [GMD+22] on the
CV-MC and MD datasets using the method developed in [CL21]. For the NF, the
normalized density is available analytically. The figures show the mean
squared error (mse) between the free energy estimated by the brute-force MD
simulation (considered as ground truth) and the one estimated by NF (green
line) and by the CV-MC datasets (blue line), along the iterations of the
algorithm. Since algorithm 1 alternates local MALA steps and non local moves
with steered MD, each iteration has a cost in terms of what we have called
equivalent MD steps, which is calculated multiplying the number of walkers
running in parallel and the number of local steps at each iteration and
summing that to the number of walkers times the average number of steps taken
by the steered MD step. In this way, we can monitor the convergence of the
free energy of the NF and CV-MC samples to the ground truth as a function the
computational cost. As a baseline for the expected error on the free energy,
we computed the mean squared difference between estimators obtained by two
independent MD trajectories (which could in principle both be considered as
ground truth), and reported it as a black dashed line in fig. 3d). This line
can be considered as a lower bound for the best possible result for the NF and
CV-MC mse. As shown in the plots of fig. 3d), the free energy estimated by CV-
MC (blue line) converges quickly to the estimate of the MD simulation. This is
not the case for the free energy estimated from NF samples (green line). In
fact after an initial phase of learning, the NF free energy does not improve
anymore. This is mainly due to the poor amount of points on the energy
barrier, which result in a impossibility for the NF to learn its proper shape,
and, possibly, to the insufficient expressiveness of the NF. The red and
orange lines represent the probability mass in the free energy minimum
corresponding to the extended state of the polymer obtained with the NF and
with the CV-MC samples as a function of the number of the number of equivalent
MD steps. The reference value obtained with MD is represented as a grey dashed
horizontal line. In contrast to what happens with the mse of the 6D free
energy, here both the NF and the CV-MC converge to the expected result. This
means that the failure of the NF to represent the barrier, and the unnatural
connection that it draws between the two free energy minima, causes errors in
the free energy landscape, but not in the aggregated relative probability of
the two minima.

In fig. 3e) we show the results of the same analysis, but relative to the high
barrier case. In this case the lower bound for the mse is produced by taking
the biased dataset produced by umbrella sampling, correcting it for the bias
and then splitting it randomly in two parts and estimating the mse between the
two halves. The sample from CV-MC has a free energy that converges to the
ground truth with an accuracy compatible with the low barrier case. The NF
sample, instead, has a free energy which is much more distant from the true
one. Here, there are no samples in the very unlikely region aound the barrier,
making it impossible to learn for the NF. Yet again, the weight of the
extended state, instead, is well represented even in this case (red and orange
lines). This shows the power of the procedure, as we can efficiently estimate
free energy differences between states separated by a $28\,K_{\rm B}T$ barrier
at the same cost as with a $7\,K_{\rm B}T$ barrier.

In fig. 3b) and c) we show the marginal free energy of the system calculated
with respect to the end-to-end distance $d$ of the polymer. Notice that this
variable is not part of the CV that we have chosen for the coarse graining,
but it is useful to visualize the two distinct minima. In fig. 3b) we show the
result of the application of algorithm 1 to the low barrier case, we compare
the result of the output of the CV-MC sampling (orange line) and a sample from
the trained NF without MC correction (green line) with the long MD trajectory
(blue line). The CV-MC and the MD curves are very well overlapped, showing
that the new algorithm is working as expected and can be used to have a good
estimate of free energy differences. The NF curve, instead, shows that the
barrier is poorly represented and, as a consequence of overall normalization,
also the shape of the minima does not perfectly match with the ground truth
especially for the higher minimum. In figure fig. 3c) we present the same
result, but in a case in which the barrier was set to $28K_{\rm B}T$. Here, we
exploit the real power of algorithm 1, as we are able to sample efficiently
enough jumps between one minimum and the other even with a very high barrier
between them, at the same computational cost as the low barrier case. The
overlapping between the free energy obtained by umbrella sampling and the CV-
MC free energy is excellent. For what concerns the NF sample, instead, we face
again the same problems as in the low barrier case.

As a final remark we mention that the average acceptance of the non-local NF
proposal in this augmented Monte Carlo is greater than 20% both in the low and
in the high barrier case. In the Supplementary Information we show a plot of
the average acceptance computed after each resampling step.

Figure 3: a) 2D slices of the 6D probability density in CV space, with respect
to some couples of collective variables. The first column is realized with MD
data, the second with the output of algorithm 1, the third with a sample from
the trained normalizing flow. b) 1D projection of the free energy on the end-
to-end distance of the polymer, in the case where we set a $7\,K_{\rm B}T$
barrier on $V(d)$. The three lines are realized with a long MD simulation
(blue), the output of algorithm 1 (orange) and a sample from the trained NF
(green). c) 1D projection of the free energy on the end-to-end distance of the
polymer, with a $28\,K_{\rm B}T$ barrier on $V(d)$. The three lines are
realized with a biased MD trajectory (blue), with the output of algorithm 1
(orange) and a sample from the trained NF (green). d) Mean squared error
between the 6D free energy obtained from the long MD simulation and the sample
from the trained NF (green line) and the sample from algorithm 1 (blue line).
The dashed black horizontal line represents the error obtained comparing two
distinct long MD trajectories. The orange and red line represent the weight of
the extended configuration of the polymer with the sample obtained both with
the CV-MC (orange) and with NF (red), this is all compared with the target
weight (grey dashed line) e) The results are the same as in d), but for the
high barrier case.

##  4 Discussion

In this work we introduced an algorithm which exploits a NF to perform
enhanced sampling in a coarse-grained space defined by a mid-size CV vector,
while generating Boltzmann-distributed configurations of the whole fine-
grained system. We first tested this approach on a mixture of Gaussians,
considering only some of the coordinates as CVs. The algorithm has shown a
good performance in sampling from this multimodal distribution, even with
highly separated modes. The relative weights of the different modes were
perfectly reproduced and the marginal distribution on the CV space was
correctly learnt by the NF and well represented by the samples generated by
the algorithm. This simple test case allowed us to verify that the efficiency
is only weakly affected by the energy barrier separating the different modes
and by the number of transversal degrees of freedom, if the CV set is properly
selected.

We then used the same approach to sample the configurations of a much more
complex toy model, a small polymer in a Lennard-Jones bath. This test case
added two main complexities to the sampling procedure. First of all, the total
number of degrees of freedom passed from being on the order of tens to several
hundreds. Secondly, the interaction between the CV and the transversal degrees
of freedom became non trivial, since moving the internal distances (or
equivalently the dihedrals) of a polymer in a bath causes LJ repulsion between
the bath particles and the beads of the polymer that is being moved. These
complexities forced us to develop a non trivial scheme to produce the steered
moves as in algorithm S3, by finding an approximate law for the optimal
$N(\|\Delta\psi\|_{2})$. We have shown that even in this complex case the
algorithm works remarkably well both with a low and with a high energy barrier
between the two modes of the system, and produces results that are compatible
with those of standard methods of simulation.

Our approach offers some important advantages for sampling multimodal
distributions in molecular systems. First of all, the use of a NF for making
proposals in the CV space allows us to use a large number of collective
variables. This is very useful because in most of the enhanced sampling
problems choosing the proper minimal set of CV is typically one of the hardest
tasks. Moreover the efficiency of the present enhanced sampling techniques is
high only when the CV space is low dimensional, while with the NF we can
afford to include tens of collective variables. Another very important
advantage of our approach comes from the fact that, if the CV set is properly
chosen, the acceptance of a move between two points does not depend on the
height of the free energy barrier between them. In the case of the polymer,
the dependence of the average acceptance on the barrier is very weak even
though there is a non trivial coupling between transversal degrees of freedom
and CV, confirming again the validity of this result. One more important
advantage of the algorithm is the good scaling that it shows with the number
of transversal degrees of freedom, this is proven by the Gaussian mixture
example and by the large number of solvent particles used in the polymer toy
model. Finally, one great advantage is that, once trained, the NF can be used
at will to have a very efficient Monte Carlo to produce rapidly uncorrelated
configurations for the system of interest.

The procedure still shows weak points that will have to be improved in a
future work. One downside is given by the fact that the NF is not able to
perfectly learn the underlying marginal probability density function, as one
can see for instance from fig. 3b) to e). This causes some small biases in the
CV proposal distribution that on one side affect the overall acceptance, on
the other side do not allow us to use directly the trained NF to produce
samples without the need of a MC step to re-weight them. A more important
drawback is that moves that require travelling a large distance in CV space
have a good acceptance only if performed using a large number of MD steps
during the steering (see fig. S4 a)). This affects the overall efficiency of
the algorithm, making it competitive with other enhanced sampling methods only
when we have to overcome large free energy barriers between different minima,
or when we cannot choose a small set of collective variables.

For future developments, we will need to improve how the chosen map reproduces
the underlying marginal probability density function of the system. This
improvement can be addressed in two directions. On one side, we could test the
performance of more sophisticated generative models in the CV space, e.g.
using diffusion models and working directly in Cartesian coordinates with
equivariant models [YTDB+23, MSA+23, ARB+21]. On the other side, we could
improve the chosen loss function, and make it more suitable for learning
better also the shape of the free energy barrier.

##  5 Material & Methods

### Mixture of Gaussians

Let $x=(\psi,x_{\perp})$ and $g(x)=\mathcal{N}\left(x-\mu,\Sigma\right)$ be a
multivariate Gaussian with mean $\mu\in\mathbb{R}^{n}$ and covariance
$\Sigma\in\mathbb{R}^{n\times n}$, the tested distributions were built as
follows:

$\displaystyle
p(x)=\left(w_{1}\mathcal{N}\left(\psi-\psi_{1},\Sigma_{1}\right)+w_{2}\mathcal{N}\left(\psi-\psi_{2},\Sigma_{2}\right)\right)\times\mathcal{N}\left(x_{\perp}-\vec{0},\mathbb{I}\right)$
(10)

In all the tests we considered $\psi\in\mathbb{R}^{2}$, $w_{1}=\frac{1}{4}$,
$w_{2}=\frac{3}{4}$, $\Sigma_{1}$ has off-diagonal elements
$\sigma_{12}=\sigma_{21}=-3.5\times 10^{-2}$ and diagonal elements
$\sigma_{ii}=5\times 10^{-2}$, $\Sigma_{2}$ is diagonal with
$\sigma_{ii}=0.2$. For fig. 1a) to f) we have considered only one transversal
degree of freedom and $\psi_{1}=(-m,m)$, $\psi_{2}=(m,m)$ with $m=1.84$. For
fig. 1g) we have used one transversal degree of freedom but varied $m$ in the
set $\\{1.0,1.21,1.42,1.63,1.84,2.05,2.26,2.47,2.68,2.89\\}$. Since there are
two modes with different weights, the free energy has two minima, one lower
than the other and a single maximum in between. The energy barrier is
calculated taking the difference between the maximum and the highest minimum.
For fig. 1h) we kept $m=1.84$ and changed the number of transversal degrees of
freedom from $1$ to $28$.

In all simulations, the local steps were performed using MALA as described in
algorithm S2 with parameters $\beta=1$, friction $\gamma=1$, and time-step
$\tau=0.005$. The steered moves had $N=20$ steps and 60 independent random
walkers were initialized in each minimum.

The architectures of the tested NFs were as follows: for the RealNVP we have
used 40 affine coupling layers where the scaling and translation networks were
multilayer perceptrons with 3 hidden layers, all with width equal to 10. The
base distribution was chosen to be a Gaussian with average and covariance
taken from a population of independent walkers evenly distributed among the
minima and evolved for some MC steps inside the minima to have them
decorrelated. For the RealNVP, it was essential to have a number of training
iterations on the order of tens of thousands as the convergence of the loss
function was relatively slow. We used 40000 training iterations with learning
rate 0.0025. For the RQS, we leveraged the library nflows[DBMP20] and used
splines with 10 bins between $-5$ and 5 and 3 layers, the parameters of the
spline were obtained with multilayer perceptrons of depth 6 and 12 hidden
features per layer. The base distribution, in this case, was a standard
normal. The RQS not only has a much better representational power, but also
converges faster, in about 2000 iterations, using the same learning rate as
the one for the RealNVP.

### The polymer in a Lennard-Jones solvent: MD simulation

The simulation was performed using MALA as described in SI algorithm S2. For
the interaction between different particles we have used a Lennard Jones
potential: $U_{\rm
LJ}(r)=4\epsilon[\left(\frac{\sigma}{r}\right)^{12}-\left(\frac{\sigma}{r}\right)^{6}]$.
We have chosen for solvent-solvent interaction $\epsilon_{ss}=1$, for polymer-
polymer interaction $\epsilon_{pp}=1$, and for polymer-solvent interaction
$\epsilon_{ps}=0.5$, while we have set $\sigma=1$ in all the cases. All the
energy units were then set relative to $\epsilon_{ss}$. As parameters for the
dynamics we have chosen $K_{\rm B}T=1.3\,\epsilon_{ss}$, $\gamma=1$,
$\tau=0.0002$. The time step $\tau$ was chosen in such a way to have the
fastest possible decorrelation among successive MALA steps, which occurs when
the average MALA acceptance is roughly 70%. The boundary conditions were
chosen to be periodic, in a cubic box of size $7.5$ (all the lengths are
relative to the polymer bond length, which is chosen to be 1). The potential
on the end-to-end distance has the form
$V(d)=\frac{v}{4}\left(\frac{d-d_{0}}{\alpha}\right)^{4}-v\left(\frac{d-d_{0}}{\alpha}\right)^{2}$,
with $d_{0}=4.62$ and $\alpha=0.8$. We set $v=9.1$ or $v=36.4$ yield
respectively $\sim 7\,K_{\rm B}T$ and $\sim 28\,K_{\rm B}T$ barriers between
the extended and the compact state. The particles are evolved with an
overdamped Langevin dynamics, while the polymer is evolved through the
integration of a second order Newtonian dynamics, with a chosen monomer mass
$m=0.1$. There is no thermostat linked to the polymer, as it already
thermalizes with the surrounding bath of LJ particles. After every step of the
system evolution we set a Metropolis acceptance rule. The dynamics of the
polymer is constrained: the bond lengths are all fixed to 1 and all the angles
between adjacent monomers are fixed to $70.5^{\circ}$. Moreover the frame of
reference of the simulation is chosen in such a way that the central monomer
of the polymer is fixed, therefore the only relevant degrees of freedom of the
polymer are the dihedral angles or, equivalently, the internal distances, as
there is a one to one mapping between the two. For this reason, all the forces
are projected on the internal distances and the dynamics is done directly on
them. Even though the dynamics of the polymer is calculated directly on the
internal distances, at each iteration, we need to know the 3D coordinates of
every monomer to be able to calculate the forces coming from the LJ
interactions. To do so, we use the NeRF algorithm [AlQ19] to pass from the
internal coordinates of the polymer to the Cartesian ones. To ease the
computation of forces we compute derivatives with the autodifferentiation
library Zygote [Inn51]. The reference MD trajectory used in the analysis in
fig. 3b) was obtained with $10^{9}$ steps, saving one configuration each
$10^{4}$ steps. To obtain the free energy in the high-barrier case used as a
reference in fig. 3c) we used umbrella sampling, with a bias potential
$V_{B}(d)=-V(d)$ (which corresponds to switching off the potential on the end-
to-end distance). Also in this case, we run the simulation for $10^{9}$ steps.

### The polymer in a Lennard-Jones solvent: CV-MC simulation

The simulation was performed with 200 random walkers running in parallel, half
of them initialized in the compact state and half initialized in the extended
state. The local moves corresponded to 30 local MD updates with the schedule
explained above at each iteration of algorithm 1.

For setting up properly the non local moves we have done some preliminary
study on the system.

We first estimated the optimal number of steered steps $N(\cdot)$ for any
proposed steered update, as a function of the size of the CV jump
$\|\Delta\psi\|_{2}$. $N(\|\Delta\psi\|_{2})$ is the value that maximises the
ratio between the probability of acceptance $p$ (6) over the computational
cost of the steered update, which is given by the number of intermediate steps
$N$ that we take in a steered move. We consider a for a set of values
$\|\Delta\psi\|_{2}$ spanning well the CV space, and for each estimate
$N(\|\Delta\psi\|_{2}$) from a set of independent transitions between multiple
pairs of points $(\psi,\psi^{\prime})$ that have the same distance
$\|\psi-\psi^{\prime}\|_{2}$ in CV space. More precisely, we consider for this
study loop transitions: from $\psi$ value to $\psi^{\prime}$ and back to
$\psi$. Thanks to this construction the free energy difference of the loop
transition is zero, and the acceptance of the transition is influenced only by
the work performed on the system, it is also independent of the direction of
the transition $\psi\to\psi^{\prime}$ versus $\psi^{\prime}\to\psi$. The
results of the optimal number of steps in the steered update for a given
$\|\psi-\psi^{\prime}\|_{2}$ are presented in the SI fig. S4a). There are two
main regimes with a sharp transition. The first regime, where the optimal
number of steps is on the order of a few hundreds, corresponds to small scale
transitions where the NF proposes new configurations inside the same mode,
namely transitions from compact state to compact state or from extended to
extended. The second one, where the optimal number of steps is on the order of
tens of thousands corresponds to large scale transitions, from one stable
state to the other were the polymer is stretched or compacted and all the
particles around it have to adapt to the new conformations.

To accelerate the learning of the NF, we have pretrained the network on
configurations obtained by local MD simulations with an equal number of
walkers in the two minima of the potential. The NF was then learns first the
presence of two modes but considering their weights are equal. This explains
why the acceptance of the resampling steps is quite high even at the beginning
of the algorithm iterations, as one can see from fig. S4b). We use the same
RQS architecture as in the Gaussian example. After the pretraining phase, the
full algorithm was run using as initialization for the NF the pretrained RQS,
the learning rate was set to $0.0045$, the total number of training iterations
was $12\,000$, and the total number of resampling steps was 400.

## Acknowledgements

We thank Tony Lelièvre, Christoph Schönle, Pilar Cossio, Gabriel Stoltz and
Eric Vanden-Eijnden for insightful discussions. The work of M.G. is supported
by the Hi! Paris center. Work partially funded by NextGenerationEU through the
Italian National Centre for HPC, Big Data, and Quantum Computing (grant number
CN00000013).

## References

  * [AKS19] M. S. Albergo, G. Kanwar, and P. E. Shanahan.  Flow-based generative models for Markov chain Monte Carlo in lattice field theory.  Physical Review D, 100(3):034515, August 2019.
  * [AlQ19] Mohammed AlQuraishi.  Parallelized natural extension reference frame: Parallelized conversion from internal to cartesian coordinates.  Journal of Computational Chemistry, 40, 2019.
  * [ARB+21] Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong.  Iterated Denoising Energy Matching for Sampling from Boltzmann Densities, arXiv:2402.06121.
  * [CCHK89] E A Carter, G Ciccotti, J T Hynes, and R Kapral.  Constrained reaction coordinate dynamics for the simulation of rare events.  Chem. Phys. Lett., 156:472–477, 1989.
  * [CL21] Matteo Carli and Alessandro Laio.  Statistically unbiased free energy estimates from biased simulations.  Molecular Physics, 119(19-20):e1899323, 2021.
  * [Cle08] Cecilia Clementi.  Coarse-grained models of protein folding: Toy models or predictive tools?  Current Opinion in Structural Biology, 18(1):10–15, February 2008.
  * [CR15] Yunjie Chen and Benoît Roux.  Generalized Metropolis acceptance criterion for hybrid non-equilibrium molecular dynamics—Monte Carlo simulations.  The Journal of Chemical Physics, 142(2):024101, January 2015.
  * [CR24] Shriram Chennakesavalu and Grant M. Rotskoff.  Data-Efficient Generation of Protein Conformational Ensembles with Backbone-to-Side-Chain Transformers.  The Journal of Physical Chemistry B, 128(9):2114–2123, March 2024.
  * [CTR23] Shriram Chennakesavalu, David J. Toomer, and Grant M. Rotskoff.  Ensuring thermodynamic consistency with invertible coarse-graining.  The Journal of Chemical Physics, 158(12):124126, March 2023.
  * [CTWZ23] Simone Ciarella, Jeanne Trinquier, Martin Weigt, and Francesco Zamponi.  Machine-learning-assisted Monte Carlo fails at sampling computationally hard problems.  Machine Learning: Science and Technology, 4(1):010501, March 2023.
  * [DBMP19] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  Neural spline flows.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. dAlché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [DBMP20] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  nflows: normalizing flows in PyTorch, November 2020.
  * [DDRW21] Luigi Del Debbio, Joe Marsh Rossney, and Michael Wilson.  Efficient Modelling of Trivializing Maps for Lattice $\phi^4$ Theory Using Normalizing Flows: A First Look at Scalability.  Physical Review D, 104(9):094507, November 2021.
  * [DP01] E Darve and A Pohorille.  Calculating free energies using average force.  J. Chem. Phys., 115:9169–9183, 2001.
  * [DSB17] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using Real NVP.  In International Conference on Learning Representations, 2017.
  * [EVE04] W E and E Vanden-Eijnden.  Metastability, conformation dynamics, and transition pathways in complex systems.  Lecture notes in computational science and engineering, Ed. S Attinger and P Koumoutsakos, Springer, 2004.
  * [Fre04] Daan Frenkel.  Speed-up of monte carlo simulations by sampling of rejected states.  Proceedings of the National Academy of Sciences, 101(51):17571–17575, 2004.
  * [GDMG23] Louis Grenioux, Alain Oliviero Durmus, Eric Moulines, and Marylou Gabrié.  On Sampling with Approximate Transport Maps.  In Proceedings of the 40th International Conference on Machine Learning, pages 11698–11733. PMLR, July 2023.
  * [GMD+22] Aldo Glielmo, Iuri Macocco, Diego Doimo, Matteo Carli, Claudio Zeni, Romina Wild, Maria d’Errico, Alex Rodriguez, and Alessandro Laio.  Dadapy: Distance-based analysis of data-manifolds in python.  Patterns, page 100589, 2022.
  * [GRVE22] Marylou Gabrié, Grant M. Rotskoff, and Eric Vanden-Eijnden.  Adaptive Monte Carlo augmented with normalizing flows.  Proceedings of the National Academy of Sciences, 119(10):e2109420119, March 2022.  arXiv:2105.12603 [cond-mat, physics:physics].
  * [Inn51] Michael Innes.  Don’t unroll adjoint: Differentiating ssa-form programs, arXiv:1810.07951.
  * [KPV96] S Kumar, P W Payne, and M Vásquez.  Method for free-energy calculations using iterative techinques.  J. Comput. Chem., 17:1269–1275, 1996.
  * [KRB+95] S Kumar, J M Rosenberg, D Bouzida, R H Swendsen, and P A Kollman.  Multidimensional free-energy calculations using the weighted histogram analysis method.  J. Comput. Chem., 16:1339–1350, 1995.
  * [LBP+20] Wei Li, Craig Burkhart, Patrycja Polińska, Vagelis Harmandaris, and Manolis Doxastakis.  Backmapping coarse-grained macromolecules: An efficient and versatile machine learning approach.  The Journal of Chemical Physics, 153(4):041101, July 2020.
  * [LP02] Alessandro Laio and Michele Parrinello.  Escaping free energy minima.  Proc. Natl. Acad. Sci. U.S.A., 99:12562–12566, 2002.
  * [MMLL22] 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):1602–1617, April 2022.
  * [MSA+23] Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato.  SE(3) Equivariant Augmented Coupling Flows.  Advances in Neural Information Processing Systems, 36:79200–79225, December 2023.
  * [NCMC11] Jerome P. Nilmeier, Gavin E. Crooks, David D. L. Minh, and John D. Chodera.  Nonequilibrium candidate monte carlo is an efficient tool for equilibrium simulation.  Proceedings of the National Academy of Sciences, 108(45):E1009–E1018, 2011.
  * [NOKW19] 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):eaaw1147, September 2019.
  * [PNR+21] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing Flows for Probabilistic Modeling and Inference.  Journal of Machine Learning Research, 22(57):1–64, 2021.
  * [PV18] Alexander J Pak and Gregory A Voth.  Advances in coarse-grained modeling of macromolecular complexes.  Current Opinion in Structural Biology, 52:119–126, October 2018.
  * [RM15] Danilo Rezende and Shakir Mohamed.  Variational Inference with Normalizing Flows.  In Proceedings of the 32nd International Conference on Machine Learning, pages 1530–1538. PMLR, June 2015.
  * [RMZT02] L Rosso, P Minary, Z.W Zhu, and M.E Tuckerman.  On the use of the adiabatic molecular dynamics technique in the calculation of free energy profiles.  J. Comput. Phys., 116:4389–4402, 2002.
  * [SGLS60] Christoph Schönle, Marylou Gabrié, Tony Lelièvre, and Gabriel Stoltz.  Sampling metastable systems using collective variables and Jarzynski-Crooks paths, arXiv:2405.18160.
  * [vEMB03] TS van Erp, D Moroni, and PG Bolhuis.  A novel path sampling method for the calculation of rate constants.  J. Comput. Phys., 118(17):7762–7774, MAY 1 2003.
  * [WWZ19] Dian Wu, Lei Wang, and Pan Zhang.  Solving Statistical Mechanics Using Variational Autoregressive Networks.  Physical Review Letters, 122(8):080602, February 2019.
  * [WXC+22] Wujie Wang, Minkai Xu, Chen Cai, Benjamin K. Miller, Tess Smidt, Yusu Wang, Jian Tang, and Rafael Gomez-Bombarelli.  Generative Coarse-Graining of Molecular Conformations.  In Proceedings of the 39th International Conference on Machine Learning, pages 23213–23236. PMLR, June 2022.
  * [YTDB+23] Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola.  SE(3) diffusion model with application to protein backbone generation.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 40001–40039. PMLR, 23–29 Jul 2023.

## Supplementary information

Algorithm S2 Metropolis-Adjusted Langevin Algorithm (MALA)

$U(\cdot)$ system energy function, $\beta$ inverse temperature, $\tau$
integration time step, $\gamma$ friction coefficient, $x(0)$ starting
configuration, $n_{\rm steps}$ number of total steps.

for $n=1\dots n_{\rm steps}$ do

$\widetilde{x}=x\left(n\right)-\frac{\tau}{\gamma}\nabla
U\left(x\left(n\right)\right)+\sqrt{\frac{2\tau}{\beta\gamma}}\eta,\quad\eta\sim
N\left(0,1\right)$ $\triangleright$ propose new configuration with overdamped
Langevin step

$\tilde{\eta}=\sqrt{\frac{\beta\tau}{2\gamma}}\left(\nabla
U\left(x\left(n\right)\right)+\nabla U\left(\widetilde{x}\right)\right)-\eta$

$\log
R=-\frac{1}{2}\left(\tilde{\eta}^{2}-\eta^{2}\right)-\beta\left(U\left(\widetilde{x}\right)-U\left(x\left(n\right)\right)\right)$
$\triangleright$ calculate acceptance probability

$u\sim\rm\mathcal{U}([0,1])$$\triangleright$ where $\mathcal{U}$ is the
uniform law.

if $u<\min\\{1,R\\}$ then $\triangleright$ accept or reject the proposal

$x\left(n+1\right)=\widetilde{x}$

else

$x\left(n+1\right)=x\left(n\right)$

end if

end for

return $\\{x\left(n\right)\\}_{n=1,...,n_{\rm steps}}$

Algorithm S3 Steered Update

$U(\cdot)$ system energy function, $\beta$ inverse temperature, $\tau$
integration time step, $\gamma$ friction coefficient, $x=(\psi,x_{\perp})$
starting configuration, $N(\cdot)$ steering schedule discretization rule,
$\rho_{\theta}^{\Psi}(\psi)$ proposal probability parametrized by the flow.

$(\psi^{0},x_{\perp}^{0})=(\psi,x_{\perp})$ $\triangleright$ store starting
point

$\psi^{\prime}\sim\rho_{\theta}^{\Psi}(\psi)$ $\triangleright$ draw new set of
collective variables

$\Delta\psi=\frac{\psi^{\prime}-\psi}{2N(\|\psi^{\prime}-\psi\|_{2})}$

$\Delta S=0$ $\triangleright$ initialize action produced

while $n<N(\|\psi^{\prime}-\psi\|_{2})$ do

$\psi\leftarrow\psi+\Delta\psi$ $\triangleright$ update CV

$x_{\perp}\leftarrow{\rm
MALA}(U(\psi,\cdot),\beta,\gamma,\tau,x_{\perp},n_{\rm steps}=1)$
$\triangleright$ update transversal d.o.f. with 1 MALA step at fixed $\psi$

$\Delta S\leftarrow\Delta S+U(\psi,x_{\perp})-U(\psi,x_{\perp}^{\prime})$
$\triangleright$ save the change in energy

$\psi\leftarrow\psi+\Delta\psi$ $\triangleright$ update CV

$n\leftarrow n+1$

end while

$R=\frac{\rho_{\theta}^{\Psi}(\psi^{0})}{\rho_{\theta}^{\Psi}(\psi)}\exp\left({-\beta(\Delta
S+U(x_{\perp},\psi)-U(x_{\perp}^{0},\psi^{0}))}\right)$ $\triangleright$
calculate acceptance probability

$u\sim\rm\mathcal{U}([0,1])$$\triangleright$ where $\mathcal{U}$ is the
uniform law

if $u<\min\\{1,R\\}$ then $\triangleright$ accept or reject the proposal

return $(\psi,x_{\perp})$

else

return $(\psi^{0},x_{\perp}^{0})$

end if

Figure S4: a) Optimal number of steps for a steered transition as a function
of the distance $\|\Delta\psi\|_{2}$ travelled in CV space. b) Acceptance of
the non-local steered moves averaged among the 200 parallel random walkers for
the two tried barriers.


[1,3,4]Gottfried Schlaug

1]Institute of Applied Life Sciences, University of Massachusetts, Amherst,
01003, MA

2]Manning College of Information and Computer Sciences, University of
Massachusetts, Amherst, 01003, MA

3]Department of Biomedical Engineering, University of Massachusetts, Amherst,
01003, MA

4]Department of Neurology, Baystate Medical Center, and UMass Chan Medical
School - Baystate Campus, Springfield, 01107, MA

# SAM vs BET: A Comparative Study for Brain Extraction and Segmentation of
Magnetic Resonance Images using Deep Learning

Sovesh Mohapatra  Advait Gosai<EMAIL_ADDRESS>[  [  [  [

###### Abstract

Brain extraction is a critical preprocessing step in various neuroimaging
studies, particularly enabling accurate separation of brain from non-brain
tissue and segmentation of relevant within-brain tissue compartments and
structures using Magnetic Resonance Imaging (MRI) data. FSL’s Brain Extraction
Tool (BET), although considered the current gold standard for automatic brain
extraction, presents limitations and can lead to errors such as over-
extraction in brains with lesions affecting the outer parts of the brain,
inaccurate differentiation between brain tissue and surrounding meninges, and
susceptibility to image quality issues. Recent advances in computer vision
research have led to the development of the Segment Anything Model (SAM) by
Meta AI, which has demonstrated remarkable potential in zero-shot segmentation
of objects in real-world scenarios. In the current paper, we present a
comparative analysis of brain extraction techniques comparing SAM with a
widely used and current gold standard technique called BET on a variety of
brain scans with varying image qualities, MR sequences, and brain lesions
affecting different brain regions. We find that SAM outperforms BET based on
average Dice coefficient, IoU and accuracy metrics, particularly in cases
where image quality is compromised by signal inhomogeneities, non-isotropic
voxel resolutions, or the presence of brain lesions that are located near (or
involve) the outer regions of the brain and the meninges. In addition, SAM has
also unsurpassed segmentation properties allowing a fine grain separation of
different issue compartments and different brain structures. These results
suggest that SAM has the potential to emerge as a more accurate, robust and
versatile tool for a broad range of brain extraction and segmentation
applications.

###### keywords:

Brain Extraction, FMRIB Software Library, Segment Anything Model, Deep
Learning

##  1 Main

Brain extraction from MRI or other imaging techniques is an important
preprocessing step to advance various further image segmentation tasks [1].
This process involves isolating brain tissue from non-brain tissues, such as
the skull, scalp, and meninges, while preserving brain tissue that might have
been altered due to various types of lesions or injuries, enabling accurate
and effective subsequent processing [2, 3, 4]. Automatic and user independent
brain extraction is challenged when brain tissue or non-brain tissue have
similar or inhomogeneous signal intensity, leading to confusion around what
constitutes brain tissue. The accurate removal of non-brain tissue enables
automatic brain segmentation models to achieve significantly higher accuracy,
and the reduction in information being processed enhances computational
efficiency, making this an indispensable component in the preprocessing
pipeline of brain imaging studies [5, 6].

The current gold standard in automatic brain extraction is a tool offered by
the FMRIB Software Library (FSL). However, this Brain Extraction Tool (BET)
has several limitations, such as over-extraction leading to loss of brain
matter, inaccurate extraction resulting in residual skull or non-brain tissue
surrounding the brain to be falsely identified and included as brain tissue,
and sensitivity to image quality, all of which complicates its application on
clinical images [7, 8, 9]. Consequently, manual intervention is often required
to correct over and under-cropping errors, especially when dealing with scans
exhibiting atypical features such as lesions, injuries, or uncommon image
sequences [10, 11, 12]. Furthermore, BET’s restricted scope confines its
utility to general brain extraction, without providing the capability to
isolate or segment specific regions or structures of the brain to be studied
individually. Some areas of interest might include the corpus callosum, the
grey and white matter, subcortical structures like the basal ganglia, vascular
structures within the brain, and any kind of brain lesions.

As a result, numerous deep learning-based models have been developed for brain
extraction and other region specific extraction related tasks. Despite this,
there is still a significant demand for enhancing the accuracy of their
results [13, 14]. These models heavily rely on training data and have
exhibited a lack of robustness [15]. As a consequence, BET remains the most
widely used whole brain extraction tool, largely due to its extraction speed
and familiarity within the field.

The Segment Anything Model by Meta AI (SAM) has emerged as a state-of-the-art
tool for image segmentation of any kind, being trained over a vast dataset of
1 billion masks and 11 million images [16]. SAM has already shown remarkable
potential in accurately segmenting objects in real-world scenarios; its
extensive training and zero-shot learning allow it to respond appropriately to
any prompt at inference time [17, 18]. Its demonstrated speed, accuracy and
versatility position it as a promising tool for overcoming the limitations of
BET and other deep learning-based models for whole brain extraction, as well
as segmenting isolated brain regions or structures.

To explore the potential of SAM in neuroimaging research, we conducted a
comparative analysis of brain extraction techniques and tested the capability
of SAM to segment structures within the brain. Our study evaluated both SAM
and BET on a diverse range of brain scans for brain extraction, which included
varying image qualities and brain challenges caused by lesions and injuries.
Additionally, we highlight various additional benefits that SAM could offer to
neuroimaging researchers over existing techniques in rapidly segmenting
within-brain structures of translational research and clinically applied
relevance.

##  2 Methodology

###  2.1 Selection of MRI data

We examined 45 anonymized MR brain images curated from various databases
(ATLAS [19], white matter hyperintensities (WMH) Challenge Segmentation [20],
and our anonymized and de-identified in-house dataset of neuroimages). To
encompass a wide range of factors, including the various types of MRI
sequences utilized, quality of the images, presence of any brain
abnormalities, and use of both normalized and non-normalized data, we sampled
five scans from nine customized brain image categories each (as presented in
Section 3).

###  2.2 Preprocessing of MRI data

To enable more accurate and reliable comparisons between the results obtained
from BET and SAM, we resampled all the MRI data into MNI152 space [21]. This
approach was intended to ensure all data is aligned to a standardized
reference template, which is essential when performing group-level analyses
and comparing results across different outputs.

###  2.3 Brain Extraction using BET

BET, a command-line tool, offers the convenience of using various parameters
to enhance brain extraction in difficult or unusual cases, resulting in a
customizable and semi-automated brain extraction process. The primary
parameter which is used is the fractional intensity threshold ($f$) that
determines the intensity level to distinguish between brain and non-brain
tissue. By adjusting this value, we optimized the brain boundary estimation
for the various data sets used in this paper. Notably, BET is able to perform
brain extraction directly on 3D data, without requiring conversion into 2D
slices.

In certain instances, we found that none of the $f$ values effectively
isolated the brain without either including a significant amount of non-brain
tissue or excessively truncating the brain. In such cases, we reported the
extraction results obtained using the default value ($f=0.5$).

Figure 1: Comprehensive SAM workflow – visualizing the end-to-end process with
custom box algorithm, highlighting blue dots (exclusion mask) and yellow dots
(inclusion mask)

###  2.4 Brain Extraction using SAM

SAM facilitated the customization of the brain extraction process by enabling
us to utilize various parameters, such as intensity threshold and spatial
location, to selectively retain or remove different regions of the image. As a
result, we were able to achieve a fully automated and highly customizable
extraction process, not only for the brain but also for specific regions
within the brain, such as the corpus callosum, subcortical structures (e.g.,
basal ganglia), vascular structures, and any other distinct brain structures
of interest such as brain lesions (refer to 3.3).

In our case, we initially transformed each MRI scan into a set of 2D slices.
Subsequently, we implemented a custom bounding box algorithm to accurately
position markers, allowing us to effectively include a particular region or
structure within the brain for extraction while simultaneously excluding
unwanted regions based on their corresponding and neighboring voxel
intensities. These inclusion and exclusion markers, alongside the 2D image
slices, served as the set of inputs enabling precise mask generation from SAM.

####  2.4.1 3D MRI scan preprocessing and reconstruction

SAM requires 2D images as input, which we obtained by converting the processed
3D MRI data into various 2D slices using the med2image Python tool. The
predicted slices were then converted into the required brain extracted data
using the nibabel Python library. The end-to-end workflow can be found in
Figure 1.

###  2.5 Evaluation metrics for brain extraction

We have evaluated our results over five different metrics: Dice coefficient,
Jaccard Index (referred to as IoU), Accuracy, Precision, and Recall, in order
to understand how the outputs from BET and SAM have performed against the
standard MNI152 1mm Brain images (as shown in the figure in SI). These images
are used as ground truths for an extracted brain since they comprise of only
relevant brain tissue. Detailed definitions of each of the chosen evaluation
metrics and their relevance can be found in the supplemental information.

Table 1: Comparison of the evaluation metrics for 9 customized categories of
brain images. The source of the data is mentioned in parentheses next to the
type of MRI sequence. The numbers in bold indicate the value of the metric,
and the superior model for the specific brain image type. (Acc = Accuracy,
Prec = Precision)

| BET | SAM
---|---|---
Category | Dice | IoU | Acc | Recall | Prec | Dice | IoU | Acc | Recall | Prec
T1-weighted images |  |  |  |  |  |  |  |  |  |
from Acute Stroke Phase (ATLAS [19]) | 0.942 | 0.891 | 0.919 | 0.945 | 0.940 | 0.956 | 0.918 | 0.936 | 0.958 | 0.957
T2-weighted images |  |  |  |  |  |  |  |  |  |
from Chronic Stroke Phase (in-house) | 0.966 | 0.936 | 0.951 | 0.985 | 0.957 | 0.969 | 0.940 | 0.955 | 0.971 | 0.967
FLAIR images |  |  |  |  |  |  |  |  |  |
from Chronic Stroke Phase (in-house) | 0.949 | 0.907 | 0.928 | 0.945 | 0.954 | 0.955 | 0.915 | 0.935 | 0.954 | 0.957
Normalized T2-weighted images |  |  |  |  |  |  |  |  |  |
from Chronic Stroke Phase (in-house) | 0.919 | 0.854 | 0.860 | 0.990 | 0.862 | 0.977 | 0.955 | 0.965 | 0.990 | 0.965
Normalized FLAIR images |  |  |  |  |  |  |  |  |  |
from Chronic Stroke Phase (in-house) | 0.853 | 0.744 | 0.743 | 0.951 | 0.792 | 0.962 | 0.927 | 0.941 | 0.951 | 0.962
FLAIR images |  |  |  |  |  |  |  |  |  |
from WMH Challenge Set [20]* | 0.937 | 0.883 | 0.904 | 0.941 | 0.934 | 0.928 | 0.868 | 0.893 | 0.905 | 0.953
3DT1-weighted images |  |  |  |  |  |  |  |  |  |
from WMH Challenge Set [20] | 0.628 | 0.518 | 0.642 | 0.549 | 0.956 | 0.914 | 0.842 | 0.871 | 0.944 | 0.888
T1-weighted images |  |  |  |  |  |  |  |  |  |
from WMH Challenge Set [20] | 0.898 | 0.815 | 0.857 | 0.863 | 0.936 | 0.928 | 0.867 | 0.898 | 0.923 | 0.934
Fractional Anisotropy (FA) images |  |  |  |  |  |  |  |  |  |
from Chronic Stroke Phase (in-house) | 0.461 | 0.305 | 0.480 | 0.306 | 0.988 | 0.916 | 0.845 | 0.880 | 0.893 | 0.941

*Results discussed in Section 3.1. Refer to the SI document for more detail.

##  3 Results

###  3.1 Evaluation of Extractions by BET and SAM

In our evaluation of brain extractions performed by BET and SAM, we compared
their performance across five key metrics against ground truth. Our analysis
demonstrated that SAM outperformed BET in the majority of cases, as evidenced
by the evaluation metrics and accompanying figures from various imaging
modalities. The supplemental information contains the visual comparisons and
raw scores for each of the 45 scans used to generate our results.

The results presented in Table 1 indicate that both SAM and BET exhibit high
scores across all five metrics when processing high-quality T2-weighted images
without lesions, or images that have undergone spatial normalization to a
standard space. While normalization does not lead to a significant improvement
in BET’s performance, SAM’s metrics exhibit a minor improvement compared to
its non-normalized samples.

For our in-house datasets, BET’s metrics experience a decline when dealing
with lower quality FLAIR images (both normalized and non-normalized). This is
also observed for Acute Stroke Phase images from the ATLAS dataset [19],
potentially due to lesion size and positioning. In contrast, the custom
bounding box algorithm implemented with SAM enables high performance across
most image qualities and cases of brain abnormalities, as it allows dynamic
inclusion and exclusion of specific smaller regions based on their
intensities. The most evident examples of this include our in-house Fractional
Anisotropy (FA) images and 3DT1 weighted images from the WMH challenge set.
For instance, in the WMH 3DT1 dataset, our SAM based algorithm achieved a Dice
coefficient of 0.914, an IoU of 0.842, an accuracy of 0.871, and a recall of
0.944, This is in contrast to BET’s performance, which yielded a Dice
coefficient of 0.628, an IoU of 0.518, an accuracy of 0.642, and a recall of
0.306.

For both of these datasets, we observe a notably higher BET precision, even
though there is poor performance on other metrics. This can be attributed to
the excessive cropping by BET, which inadvertently removes large portions of
actual brain tissue. Consequently, the substantially smaller segmented region
does not contain any out-of-brain tissue, leading to a minimal number of false
positives which results in deceptively high precision value (as shown in Table
1).

SAM also presents room for improvement. This is particularly seen for the
FLAIR images from the WMH dataset [20] where SAM slightly underperforms BET,
as highly specific exclusion markers near the skull region result in
occasional over-extraction (as shown in Figure I of the supplemental
information).

Figure 2: Comparison of original MRI images and extraction outputs by BET and
SAM for A. sagittal, B. coronal, and C. axial anatomical planes. BET
extraction results in sagittal, coronal, and axial planes, respectively. SAM
extraction results in sagittal, coronal, and axial planes, respectively. The
data is from ATLAS v2.0 dataset [19].

###  3.2 Visual Comparison

Our findings from the figures are consistent with those presented in Table 1,
whereby BET performed highly consistent and desirable extractions in cases
involving high-quality images and normalized data, however, SAM consistently
outperformed BET across most other variations. For instance, Figure 2 depicts
the T1 scan from the acute stroke phase in the ATLAS v2.0 dataset, along with
its corresponding BET output and SAM output across three different planes:
sagittal, coronal, and axial. Here, we notice that BET’s performance is
compromised since the scan contained lesions proximal to the brain’s outer
contour or in regions that are challenging to distinguish from the surrounding
meninges, despite good image quality.

Figure 3: Comparison of original MRI images and extraction outputs by BET and
SAM for A. T2 scan from a chronic stroke patient in the coronal plane, B.
FLAIR scan from a chronic stroke patient in the sagittal plane, C. 3DT1 scan
from a patient having WMH in the coronal plane, and D. FA scan from a patient
having chronic stroke in the coronal plane. Data for A., B. and D. are from
in-house dataset, and C. is from WMH challenge dataset [20].

In Figure 3, we can observe one of the three planes of the original brain,
along with its corresponding BET and SAM outputs for a random sample of the
MRI modalities we studied. We observe that both BET and SAM produced
comparable brain extractions for high resolution normalized T2-weighted images
(3A), and for the in-house FLAIR image (3B) where BET leaves part of the
skull, whereas SAM performs minor over-cropping. However, the limitations of
BET related to lesion proximity to the outer parts of the brain and image
quality are clearly evident in 3C and 3D, where the extracted brains are
severely over-cropped, rendering them unsuitable for any future segmentation
or subsequent analysis. This highlights a disadvantage of using BET in such
cases, whereas SAM’s consistently accurate extractions offer a more reliable
option.

###  3.3 Applications of SAM segmenting particular brain structures and
lesions

Given its versatility, the potential applications of SAM in neuroimaging
research are vast. This tool could be used for automating the various
preprocessing pipelines required in imaging data for large scale studies like
those investigating effects and changes in neurodegenerative diseases or brain
development. Additionally, its ability to accurately segment and extract brain
structures, as well as identifying specific regions of interest, could
accelerate a lot of the studies involving structure-function relationships as
well as quantitative lesion volume and lesion load determination. This could
also contribute and enhance precision medicine and personalized medicine in
neurology and psychiatry. As demonstrated in Figure 4, we show examples of how
SAM could also be utilized for precise segmentation of brain regions and brain
lesions demonstrating its utility in segmenting ischemic and hemorrhagic
lesions, specific brain structures of interest such as the midsagittal corpus
callosum size, and subcomponents of brain that might be affected in particular
neurological disorders such as the size/volumes of ventricles or size/volume
of basal ganglia structures.

Figure 4: A., B. Ventricles extracted from a T1 scan from acute stroke phase.
C. Corpus callosum extracted from a 3DT1 scan from chronic stroke phase. D.
Lesions extracted from a DWI scan from chronic stroke phase. E. Microbleeds
extracted from a SWAN scan from chronic stroke phase. Data in A. and B. are
from ATLAS v2.0 [19], and C., D., and E. are from in-house dataset.

##  4 Discussion

The results of our study provided valuable comparative insights into the
effectiveness and various potential applications and challenges of a new
method of brain extraction, SAM, with BET, a gold-standard in the field of
brain extraction. The primary objective of our study was to examine the
performance of these tools in the context of brain extraction and provide an
overview of how effective SAM is in segmenting the specific regions of
interest. Our findings support the utility of these tools for specific tasks
of image pre-processing as well as detailed within-image processing and
analysis.

In accordance with our research question, the study has shown that both the
SAM and BET are capable of accurately extracting brain structures from a
variety of brain types to a certain degree of error. In addition, the
comparison of these tools has revealed that SAM has outperformed BET in terms
of Dice coefficient, Jaccard Index and Accuracy for a majority of our diverse
set of tested cases. SAM’s robustness has been highlighted based on both
quantitaive and qualitative analysis, suggesting its potential to be a more
suitable tool for brain extraction, segmentation, and identification of
specific regions in neuroimaging research when paired with an appropriate
inclusion/exclusion algorithm.

###  4.1 Advantages and Limitations of SAM and BET

One of the key advantages of SAM is its flexibility and adaptability to
different types of data. Since the model has been trained on a huge number of
images, it enables it to adapt to various imaging modalities and contacts,
making it a versatile tool for a wide range of neuroimaging applications.
Moreover, SAM’s impressive performance on all evaluation parameters attests to
its efficacy in brain extraction. Its ability to identify specific brain
regions, such as the corpus callosum, subcortical structures, vascular
structures, and other specialized areas, makes it a highly versatile tool.
Nonetheless, it is important to note that the conversion of 3D data to 2D
slices and subsequent reconstruction, which is required for using SAM,
represents a potential drawback and source of error in neuroimaging analysis.

BET’s strength lies in its simplicity, speed, and ease of use, with it being a
reliable tool for brain extraction particularly when brains are in a spatially
normalized space. However, our study validated that its extraction accuracy is
quite inferior to that of SAM in certain scenarios, such as when the image
quality is suboptimal, has poor contrast, unnormalized, or when images contain
brain lesions that are located near the outer portions of the brain. Moreover,
unlike SAM, BET is unable to efficiently segment and extract specific brain
sub-regions, which limits its application to primarily full brain extraction.

Finally, it is worth noting that processing the data using SAM requires high
computational resources. The entire pipeline, from converting 3D data to
running SAM on 2D slices and per- forming reconstruction, can take
approximately 30 seconds, which is considerably longer than BET.

###  4.2 Future Work

Future research in this area should concentrate on addressing the existing
limitations of SAM, such as exploring various reconstruction techniques to
optimize its performance for 3D image datasets. Section 3.3 provides a
demonstration of possible applications of SAM based segmentation algorithms in
neuroimaging research outside of brain extraction which are vital to be
explored in greater detail. These algorithms can be further refined to involve
a variety of relevant parameters to improve the quality of inclusion,
exclusion markers, or generate extremely precise bounding boxes for automatic
mask extraction from SAM. To enhance the robustness and generalizability of
SAM inherently across a broad range of contextual tasks, including fast and
precise lesion segmentation, segmentation of clinically relevant structures
within the brain, as well as lesion progression and therapy responses, fine-
tuning could be performed on its baseline model. This would require a diverse
and comprehensive training dataset that covers various neuroimaging modalities
and pathologies.

##  5 Conclusion

Our findings provide evidence of the significant additional potential of SAM
in comparison with BET in pre-processing as well as detailed within-brain
image analysis, in particularly brain extraction from the skull and meninges
even in challenging cases with lesions involving the outer parts of the brain,
brain tissue segmentation allowing the differentiation of lesioned from non-
lesioned tissue components, and rapid identification and delineation of
specific brain regions of interest. Our evaluation metrics indicate that SAM
consistently outperforms BET, demonstrating its superior performance across a
broad range of imaging modalities and pathologies. Our ongoing research aims
to further improve the effectiveness and generalizability of the algorithm
across a wide array of neuroimages and to utilize SAM’s versatility for other
brain structure segmentation and extraction tasks. We believe that
advancements in the fields of image preprocessing and segmentation can lead to
a deeper understanding of brain structure and function and can facilitate the
development of more precise and effective diagnostic approaches, which has the
potential to advance therapeutic options. Overall, our study highlights the
value of cutting-edge deep learning techniques to advance the field of applied
and translational neuroimaging.

## Supplemental Information

Here is the link to the supplementary document referenced in this paper:
shorturl.at/bACZ9

## Acknowledgement

This research was supported by NIMH (Brain-Initiative) (7R01MH111874-05). GS
also acknowledges support from NINDS (U01NS102353). SM acknowledges the
support from IALS summer internship program (direct summer stipend support).
Conflict of interest statement: GS is the co-founder of Brainify, LLC, and has
ownership interest in Brainify, LLC. Brainify LLC financially supported the
research conducted in this study, but had no influence on data processing or
the derived results or the presentation or interpretation of the results.

## References

  *   * [1] Despotovic, I., Goossens, B. & Philips, W.  Mri segmentation of the human brain.  _Challenges, Methods, and Applications, Computational and Mathematical Methods in Medicine_ 2015 (2015).
  * [2] Ashburner, J. _et al._ Computer-assisted imaging to assess brain structure in healthy and diseased brains.  _The Lancet Neurology_ 2, 79–88 (2003).
  * [3] Somasundaram, K. & Kalaiselvi, T.  Fully automatic brain extraction algorithm for axial t2-weighted magnetic resonance images.  _Computers in biology and medicine_ 40, 811–822 (2010).
  * [4] Thompson, P. M. _et al._ Detection and mapping of abnormal brain structure with a probabilistic atlas of cortical surfaces.  _Journal of computer assisted tomography_ 21, 567–581 (1997).
  * [5] Van der Kouwe, A. J., Benner, T., Salat, D. H. & Fischl, B.  Brain morphometry with multiecho mprage.  _Neuroimage_ 40, 559–569 (2008).
  * [6] Fein, G. _et al._ Statistical parametric mapping of brain morphology: sensitivity is dramatically increased by using brain-extracted images as inputs.  _Neuroimage_ 30, 1187–1195 (2006).
  * [7] Zhuang, A. H., Valentino, D. J. & Toga, A. W.  Skull-stripping magnetic resonance brain images using a model-based level set.  _NeuroImage_ 32, 79–92 (2006).
  * [8] Klein, A. _et al._ Evaluation of volume-based and surface-based brain image registration methods.  _Neuroimage_ 51, 214–220 (2010).
  * [9] Boesen, K. _et al._ Quantitative comparison of four brain extraction algorithms.  _NeuroImage_ 22, 1255–1261 (2004).
  * [10] Smith, S. M.  Fast robust automated brain extraction.  _Human brain mapping_ 17, 143–155 (2002).
  * [11] Jenkinson, M., Beckmann, C. F., Behrens, T. E., Woolrich, M. W. & Smith, S. M.  Fsl.  _Neuroimage_ 62, 782–790 (2012).
  * [12] Battaglini, M., Smith, S. M., Brogi, S. & De Stefano, N.  Enhanced brain extraction improves the accuracy of brain atrophy estimation.  _Neuroimage_ 40, 583–589 (2008).
  * [13] Pei, L. _et al._ A general skull stripping of multiparametric brain mris using 3d convolutional neural network.  _Scientific Reports_ 12, 10826 (2022).
  * [14] Fatima, A., Shahid, A. R., Raza, B., Madni, T. M. & Janjua, U. I.  State-of-the-art traditional to the machine-and deep-learning-based skull stripping techniques, models, and algorithms.  _Journal of Digital Imaging_ 33, 1443–1464 (2020).
  * [15] Thakur, S. _et al._ Brain extraction on mri scans in presence of diffuse glioma: Multi-institutional performance evaluation of deep learning methods and robust modality-agnostic training.  _NeuroImage_ 220, 117081 (2020).
  * [16] Kirillov, A. _et al._ Segment anything.  _arXiv preprint arXiv:2304.02643_ (2023).
  * [17] He, K. _et al._ Masked autoencoders are scalable vision learners.  _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ 16000–16009 (2022).
  * [18] Xian, Y., Lampert, C. H., Schiele, B. & Akata, Z.  Zero-shot learning—a comprehensive evaluation of the good, the bad and the ugly.  _IEEE transactions on pattern analysis and machine intelligence_ 41, 2251–2265 (2018).
  * [19] Liew, S.-L. _et al._ A large, open source dataset of stroke anatomical brain images and manual lesion segmentations.  _Scientific data_ 5, 1–11 (2018).
  * [20] Kuijf, H. J. _et al._ Standardized assessment of automatic segmentation of white matter hyperintensities and results of the wmh segmentation challenge.  _IEEE transactions on medical imaging_ 38, 2556–2568 (2019).
  * [21] Evans, A. C., Janke, A. L., Collins, D. L. & Baillet, S.  Brain templates and atlases.  _Neuroimage_ 62, 911–922 (2012).