JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensemble Generation
Authors: Ameya Daigavane
Bodhi P. Vani
Darcy Davidson
Saeed Saremi
Joshua A Rackers
Joseph Kleinhenz
Venue: NeurIPS 2025 poster ()
Keywords: Conformational Ensembles
Molecules
Protein Dynamics
Drug Discovery
Statistical Mechanics
Generative Models
Primary area: machine_learning_for_sciences
TL;DR: Molecular dynamics in smoothed noisy spaces to accelerate conformational ensemble generation

Abstract: Conformational ensembles of protein structures are immensely important both for understanding protein function and drug discovery in novel modalities such as cryptic pockets. Current techniques for sampling ensembles such as molecular dynamics (MD) are computationally inefficient, while many recent machine learning methods do not transfer to systems outside their training data. We propose JAMUN which performs MD in a smoothed, noised space of all-atom 3D conformations of molecules by utilizing the framework of walk-jump sampling. JAMUN enables ensemble generation for small peptides at rates of an order of magnitude faster than traditional molecular dynamics. The physical priors in JAMUN enables transferability to systems outside of its training data, even to peptides that are longer than those originally trained on. Our model, code and weights are available at https://github.com/prescient-design/jamun.

JAMUN: Bridging Smoothed Molecular Dynamies and Score-Based Learning for Conformational Ensemble Generation
Ameya Daigavanc MIT ameyad@mit.com edu
Bodhi P. Vani* Prescient Design, Genentech vanibSgene. com
Darcy Davidson
Prescient Design, Genentech davidsd50gene.com com
Saeed Saremi Prescient Design, Genentech saremis@gent com
Joshua A. Rackers Prescient Design, Genentech
Joseph Kleinhenz Prescient Design, Genentech kleinhej @gene. com
Abstract
Conformational ensembles of protein structures are immensely important both for understanding protein function and drug discovery in novel modalities such cryptic pockets. Current techniques for sampling ensembles such as molecular dy- namics (MD) are computationally inefficient, while many recent machine learning methods do not transfer to systems outside their training data. We propose JAMUN which performs MD in smoothed, noised space of all-atom 3D conformations of molecules by utilizing the framework of walk-jump sampling. JAMUN enables ensemble generation for small peptides at rates of order of magnitude faster than traditional molecular dynamics. The physical priors in JAMUN enables trans- ferability systems outside of its training data, than those originally trained on. Our model, code and weights are available at https://githhbb com/preecientt-desiign///
peptides that longer
Introduction
Proteins are inherently dynamic entities constantly in motion, and these movements can be vitally important. They are not well characterized as single structures as has traditionally been the case, but rather as ensembles of structures drawn from the Boltzmann distribution (30). Protein dynamics required for the function of most proteins, for instance the global tertiary structure motions for myglobin to bind oxygen and move it around the body (57), the beta-sheet transition to disordered strand for insulin to dissociate and find and bind to its receptor [5]. Similarly. drug discovery on protein kinases depends characterizing kinase conformational ensembles (26 In general the
search for druggable 'cryptic pockets' requires understanding protein dynamics 15), and antibody design is deeply affected by conformational ensembles (23). However, while machine learning (ML) methods for molecular structure prediction have experienced enormous success recently, ML methods for dynamics have yet have similar impact. ML models for generating molecular ensembles
widely considered the 'next frontier' 11 57 91.. In this work, we present JAMUN (Walk-Jump Accelerated Molecular ensembles with Universal Noise), generative MIL model which advances this frontier by demonstrating improvements in both speed and transfcrability over previous approaches.
While the importance of protein dynamics is well-established, it can be exceedingly difficult to sufficiently sample large biomolecular systems. The most common sampling method is molecular
" "Equal contribution.
1 Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

b)
Noise
a)
(0)
g0
Smoathhred Brongy Landscape Fy (2))
Walk
Walk
2n a/1
Noise
Walk
Jump
Langevin Dynamics with Jogay (V/1)
g
Molecuter Dynamics Brongy Lantsoape (x)
Jump
+o lo
y
Figure 1: (a) Overview of walk-jump sampling in JAMUN. (b) Adding noise to an initial conformation (r()) to obtain ( py- (c) One iteration of BAOAB ddiscreiiie Langevin dynamics (Equation and Equation 68 starting from (+)) py leads to sample y(t+1) py. (d) Denoising of y()) according to Equation gives us new samples ((())
dynamics (MD), which will sample the Boltzmann distribution in the limit of infinite sampling time, limited by the need for very short timesteps of 1-2 femtoseconds in the numerical integration scheme. Many important protein dynamic phenomena occur on the timescale of milliseconds. As described by Borhani and Shaw [10], simulating with this resolution .. equivalent to tracking advance and retreat of the glaciers of the last Ice Age tens of thousands of years by noting their locations each and every second.' Importantly, there nothing fundamental about this small time-step limitation; it is artifact of high-frequency motions, such bond vibrations, that have little effect protein ensembles (5T). Enhanced sampling methods have been developed an attempt to accelerate sampling, but they often require domain knowledge about relevant collective variables, and, more importantly, do not address the underlying time-step problem [87].
A large number of generative models have been developed to address the sampling inefficiency problems of MD using machine learning, which we discuss in greater detail in Section The key requirement is that of transferability; any model must be able to generate conformational ensembles for molecules that are significantly different from those in its training set. To benchmark this transferability, focus small peptides whose MD trajectories reasonable amount of time, unlike those for much larger proteins (78).
be run to convergence in
Here, we propose a new method, JAMUN, that bridges molecular dynamics with score-based learning in a latent space. This physical prior enables JAMUN to transfer well - just like force fields for molecular dynamics can unseen systems. The key idea is to run Langevin molecular dynamics in a noisy 'latent' space E gg-x33 x3, instead of the original space X E g all-atom 3D positions. Indeed, Y is constructed by adding small amount of i.i.d Gaussian noise to X:
Y X +o
To run this Langevin MD over Y, the crucial component the score function V, log py (v), which needs to be modelled. (This is identical to how force field needs to be parametrized in classical MD). Once the MD trajectory over is run, the resulting samples need to be mapped back to clean data via denoising procedure.
This framework is actually mathematically described by Walk-Jump Sampling (WIS), as first intro- duced by Saremi and Hyvärinen (70). WJS has been used in voxelized molecule generation 62 63

---

and protein sequence generation [24]. In particular, JAMUN corresponds to ( Walk-Jump sampler of point clouds.
The WJS framework actually tells us that the score function Vy, log (v) can be used for denoising as well, removing the need learn separate model. Indeed, this mathematical connection is used to learn the score function, similar to the training of diffusion models. Unlike diffusion models such as DDPM [32], however, we only need to learn the score function at single noise level. The choice of this noise level is important; we aim to simply smooth out the distribution enough to resolve sampling difficulties without fully destroying the information present in the data distribution.
In short, the score function V, y logpy py (v) is learned by adding noise to clean data æ, and denoising neural network is trained to recover the clean samples x from y. This denoiser defines the score function of the noisy manifold Y which we sample using Langevin dynamics (walk step) and allows us periodically project back to the original data distribution (jump step). Crucially, the walk and jump steps are decoupled from each other.
Rather than starting over from an uninformative prior for each sample is commonly done in diffusion [32 79) and low-matching [54] 47), JAMUN is able to simply denoise samples from the slightly noised distribution, enabling much greater sampling efficiency. Appendix A contains a comparison of JAMUN at a single noise level, against full diffusion which begins the sampling process at high noise level.
We train JAMUN on a large dataset of MD simulations of small peptides. We demonstrate that this model can generalize to holdout set of unseen peptides. In all of these cases, generation with JAMUN yields converged sampling the conformational ensemble faster than MD with a standard force field, even outperforming several state-of-the-an baselines. These results suggest that this transferability is a consequence of retaining the physical priors inherent in MD data. Significantly, we find that JAMUN performs well even for peptides longer than the ones seen in the training set.
2 Methods
2.1 Representing Peptides Point Clouds
Each point cloud of N atoms can be represented by the tuple (z,h) where æ E N-3 represents the 3D coordinates of each of the N atoms and h IR"" represents atom and covalent bonding information. h can be easily computed from the amino acid sequence for each peptide, and hence is not learned sampled. For clarity of presentation, and models below. We discuss how our model uses à in Section 2.4
omit the conditioning on h in the distributions
At sampling time, we assume access to an initial sample (0) E RR sampled from the clean data distribution px. Similarly to how MD simulations of small peptides are commonly seeded, we use the sequence command in the LEaP program packaged with the Amber force fields to procedurally generate (0))) In theory, (0) could also be obtained from experimental data, such as crystallized structures from the Protein Data Bank 9.. We plan explore this approach in the future to seed
JAMUN simulations for larger proteins.
2.2 Walk-Jump Sampling
JAMUN operates by performing walk-jumm sampling molecular systems represented as 3D point clouds. A conceptual overview the process is illustrated in Figure I]
Given the initial sample (0) px. walk-jump sampling performs the following steps:
1. Noise the initial structure x (0) to create the initial sample (0) from the noisy data distribution py FFigure lak:
y (0) (0) + (0) where (0) (0,,TTxx))
(2)
2. Walk to obtain samples y( "....... (N) from py using Langevin dynamics which conists of numerically solving the following Stochastic Differential Equation (SDE) Figure lb); dy u,dt,
du = v, Py Py (v)dt - yo,dt + M- /2dB,
(4)
3

---

where by represents the particle velocity, Vy ly (y) is the gradient of the log of the probability density function (called the score function) of py, friction, M is the mass. and B3 is the standard Wiener process in N x 3-dimensions: B, NiO,, II xx). In practice, we employ the BAOAB solver (Appendix
integrate L'quation Inuagrat/Equatioon numerically.
3. Jump back to px to obtain samples 81, , 2N Figure lcb:
i (vx)
X Y wi]:
(5)
where 2(-) B[X Y = | is called the denoiser. It corresponds to the minimizer Section H.I of the E2-loss between clean samples X and samples denoised back from X + ae.
a()) = arg min Ni0 (lI -
(6)
Y=X+os
where f: Nx3 - Nx3, As shown by Robbins 666
Miyasawa
(and Section H.2p.
the denoiser á is closely linked to the score V, logpy:
i(y)
R log
Importantly, the score function V. log py shows up in both the walk and jump steps, and we need to approximate this quantity.
2.3 Learning Denoise
In order to run Walk-Jump Sampling as outlined above, we have the choice of modelling either the score V, log py the denoiser they are equivalent by Eqqation 7 Following trends in diffusion models [1 42 we model the denoiser as a neural network Tig (3) R i(y) parameterized by model parameters
Importantly, we only need to learn a model at a single, fixed noise level a. This is unlike training diffusion now-matching models where wide range of noise levels are required for sampling. In particular, the choice noise level a for WJS important because mode-mixing becomes faster as is increased, but the task asked of the denoiser becomes harder.
The denoiser de thus takes in noisy point clouds y formed by adding noise (at a fixed noise level o) to clean point clouds æ. The denoiser is tasked to reconstruct back given y. To be precise, training the denoiser dg consists solving the following optimization problem:
0- = arg min E Xx ((.. (o, (l) Y= xta
X
(8)
to obtain 0", the optimal model parameters. As is standard in the empirical risk minimization (ERM) [84] setting, we approximate the expectation Equation 8) by sampling X px and Ixx33)... We minimize the loss as function of model parameters 0 using the first-order optimizer Adam in PyTorch 2.0 22
2.4 Parametrization of the Denoiser Network
We summarize the key features of the denoiser network Ro which will approximate ((v) in this section. Note that a is fixed in our setting, but we explicitly mention in this section for clarity. A diagrammatic overview of our model along with specific hyperparameters are presented in Appendix F
We utilize the same parametrization of the denoiser as originally proposed by Karras et al. (in the context of image generation):
foly: .)) Cskip (a yy + Coun (a) Fo (Cin (a)y Craise (o));
where F9 represents a learned network parameterized by parameters 8. In particular, Fe (Figure 14 geometric graph neural network (GNN) model similar to NequlP g 81
Importantly, Fg is 46 that utilize the E(3)-
chosen to be S -equiivauiya in contrast to existing methods 33 equivariant EGNN model [7T). rightly pointed out by Dumitrescu et al. 20 and Schreiner et al.
45
733, E(3) -euuivarra models are equivariant under parity, which means that are forced to transform

---

CYS
CYS
ACE
ALA
ALA
NME
Figure 2: A side-by-side comparison of uncapped (left) compared to capped (right) ALA-CYS. The acetyl (ACE) and N-methyl (NME) capping groups provide steric hindrance and prevent local charge interactions on the N-terminal and C-terminal ends.
mirrored structures identically. When we experimented with E33 equivaaraa architectures, we found symmetric Ramachandran plots which arise from the unnecessary parity constraint of the denoising network. For this reason, TBG [5 and Timewarp [46] use a 'chirality checker' to post-hoc fix the generated structures from their model. For JAMUN, such post-pree protessocssngg unnecessary because our model can distinguish between chiral structures.
The coefficients (o), (o), (n) Cnoise (o) in Equation 9 are normalization functions (from R to IR) which adjust the effective inputs and outputs to They are chosen to encourage re-use of the input y at low noise levels, but the opposite high noise levels. Importantly, based on the insight that Fa uses relative vectors the message passing steps, we adjust the values of these coefficients instead of simply using the choices made in Karras al. 42). Wohlwend et al. 89 Abramson et al. [T. as discussed in Appendix G
In Fa. edges between atoms are computed using radial cutoff of 10Å over the noisy positions in y. The edge features are a concatenation of a one-hot feature indicating bonded-ness and the radial distance embedded using Bessel functions. As obtained from h, atom-level features computed using the embedding of the atomic number (eg. C and N), and the atom name following PDB notation (eg. CA, CB for alpha and beta carbons). Similarly, residue-level features obtained using the embedding of the residue code (eg. ALA, CYS) and concatenated to each atom in the residue. Importantly, do generalization to longer peptide lengths.
the sequence index of the residues (eg. 0, .....) found that it hurts
3 Datasets
For development, demonstration, and benchmarking against existing models, we use three different datasets consisting of peptides from 2 to 5 amino acids (AA) long: TIMEWARP 2AA-LARE LLAGGE and TIMEWARP 4AA-LARGE from Klein al. [46). MDGEN 4AA EXPLICIT from Jing et al. [37], and our own MD data simulated with OpenMM (2]] A summary of these datasets is presented in Table
The differing simulation conditions across these datasets allows us to test the broad applicability of our approach. We describe these, well as an internal dataset we use for more realistic setting analysis [n
We also use the implicit MD code from Timewarp to generate trajectories for three randomly picked 5AA peptides with codes NRLCQ, VWSPF and KTYDI. We call this dataset UNCAPPED 5AA. Further details on MD simulation conditions can be found in Appendix C
Related Work and Baseline Models
The goal of building machine learning models that can generate conformational ensembles of molecular systems is not new. While full overview of this field is beyond the scope of this work
see 66) for a recent review we note a few relevant previous efforts. Boltzmann Generators [60] introduced the idea that a neural network could be used to transform the underlying data distribution into an easier-to-sample Gaussian distribution. DiffMD 900 learns a diffusion model over conformations of small organic molecules from MD17 [14] showing some level of transferabiliity
across 1o isomers [76]. Timewarp 46 uses normalizing flow as proposal distribution in MCMC sampling of the Boltzmann distribution approximate the conditional distribution of future
5

---

conformational states (1(++1t) conditional on the present state (2) ITO [73 modelled this distribution using diffusion with eqquivaaan PaiNN architecture. Equijump extended this idea with protein-specifer message-passing neural network with reweighting to sample rarer conformations of fast-folding proteins. Importantly, Timewarp was the first truly transferable Boltzmann generator, but was still too slow relative to molecular dynamics on unseen systems. Later, Transferable Boltzmann Generators (TBG) [45] built upon Timewarp by using fow-matching instead of maximum likelihood estimation and more efficient continuous normalizing flow architecture.
We discuss a few more related methods and studies
related problem is that of protein structure prediction, as successfully tackled by the AlphaFold models 777 38 [ Here, compare to Boltz-1 [89] an open- reproduction of AlphaFold3 1 Boltz-1 was trained exclusively on static crystalline structures of folded states, without any dynamics conformational information, allowing us evaluate how effectively the conformational landscape can be inferred from structural data alone.
By fixing covalent bond lengths along the backbone and side chains, AlphaFold2 introduced a 'frames' parametrization of protein structures consisting of a roto-translation together with torsion angles for each residue. MDGen (37 cleverly builds on this parametrization by creating S (E)33) (3))niinarii tokenization of the backbone torsion angles, relative to known initial conformation (here, x:))). Then, they learn a stochastic interpolant [3] (a generalization of diffusion and flo-mmatching) over the trajectories of these tokens. While their overall objective is different from ours MD trajectory generation as opposed to Boltzmann distributions we can compare to their 'forward simulation' model. Similarly to AlphaFold2, their 3E(3)-invarian tokenization is limited to single-chain proteins and peptides, but allows for more efficient architectures. On the other hand, JAMUN models atoms explicitly with (3(3) equiivarra network, which makes it far more flexible and easily extendable to molecules other than proteins, as discussed in the supplement.
Several models build upon AlphaFood22 to sample conformational ensembles, discussed further in One such method is BioEmu- technically only backbone- oonly model, but their repository provides an additional side-chain reconstruction step using H-Packer [86]. allowing comparison to the all-atom models. As seen in Figure 3] the side-chain reconstruction can be quite expensive, because of the lack of support for batched inference with H-Packer.
5 Metrics
We adopt the analysis methods of MDGen [37] In particular, we compare models by projecting their sampled distributions of all-atom positions onto a variety of variables: pairwise distances, dihedral angles of backbone (known as Ramachandran plots) and sidechain torsion angles, TICA (time-lagged independent coordinate analysis) projections, and metastable state probabilities as computed by Markov State Models (MSMs) fit with PyEMMA (72]. TICA [59] is a popular dimensionality reduction method for larger molecules which aims to extract slow collective degrees of freedom from a trajectory [61 (55) As is standard practice, all TICA projections and MSMs are estimated using the
reference MD) data.
6 Experiments
Here, compare several of these state-of-the art methods on our benchmark datasets: 1. TBG [45] on TIMEWARP 2AA-LARGE, 2. MDGen [37] on MDGEN 4AA-EXPLICIT, and 3. Boltz-1 89
and BioEmu 52 on UNCAPPED 5AA.
Figure 3 contains a summary of the sampling efficiencies for different models, averaged over the corresponding test sets peptides. While the TBG model can technically produce reweighted samples, we run it in the un-reweighted mode, making it a Boltzmann emulator which is almost 10x faster in practice. This allows for fair comparison to JAMUN.
Due to the different simulation conditions across the datasets shown Table we train a different JAMUN model for each dataset. However, the same noise level of a = 0.4A is applied for training and sampling on all datasets. In fact, all training hyperparameters are kept identical across datasets. We trained each JAMUN model for 3 days on NVIDIA RTX A100 GPUs with 40 GB memory,

---

Model
Time per Sample
Number of Samples
Total Time
TBG
720 ms
5. 000
60 min
MDGen
6 ms
10, 000
1 min
Boltz-1 BioEmu
360 ms
10, 000
60 min
15 ms
10, 000
2.5 min
BioEmu + H-Packer
4320 ms
10, 000
720 min
JAMUN (2AA)
2 ms
100,0000
3 min
JAMUN (4AA)
3 ms
100,000
5 min
JAMUN (5AA)
8 ms
100, 000
12.5 min
CAPPED 2AA
40 ms
60, 000
40 min
MDGEN 4AA-EXPLICTI
11 ms
1,000, 000
180 min
UNCAPPED 5AA
108
100, 000
180min
Figure 3: Comparison of (approximate and batched) sampling times per test peptide for baseline models (top), baseline MD simulations (middle) and JAMUN. All models and baselines a single NVIDIA RTX A100 GPU, except for MDGEN 4AA 4- A-EXPLICIT which was simulated on a single NVIDIA T4 GPU, mentioned in Jing et al. 377
1.0
R² 0.881
R 0.873
0.8
JAMUN
0.6
0.4
0.2
0.0
0.00
0.25
0.50
0.75
1.00 0.00
0.25 0.50
0.75
1.00
Reference
Reference
Figure 4: Across both TIMEWARP 2A (left) and TIMEWARP 4AA-LARGE (right), MSM state probabilities for JAMUN samples (on the y-axis) and those for the reference MD trajectories (on the axiss) across all test peptides are strongly correlated. The perfect sampler will obtain R² of 1.
although competitive results can be obtained by training for only day. Our denoiser model is built with the e3nn library (25). and contains approximately 10.5M parameters.
This scale of noise (a = 0.4A) large enough to result in significant disruption of structure, leading to the smoothed Gaussian convolved 'walk' manifold. However, the scale is also small enough to avoid atoms 'swapping' position, for instance, or pairs of bonded atoms ending up very far from each other with reasonable probability. Indeed, shown in Figure 8a find that higher noise levels (such as = 0.8Á) result in samples with broken topologies, while lower noise levels (such as 0.2A) require many more sampling steps explore the entire conformational landscape.
6.1 Results on TIMEWARP 2AA-LARGE and TIMEWARP 4 4AA- A-LARGE
First, we show that JAMUN samples similar states as the reference MD data. Indeed, Figure 4 shows that the metastable state probabilities over JAMUN sampled trajectories match very well with those over the reference MD data on the TIMEWARP 2AA-LARGE and TIMEWARP 4AA- datasets.
Table shows that TIMEWARP -LAGGE, JAMUN outperforms TBG when for equal amounts time (based on Figure 3). and only slightly worse when TBG is run for 20x longer.
In Figure 5a and Figure 5b we visualize the TICA-0,1 projections and Ramachandran plots for randomly chosen test peptides from TIMEWARP 2AA-LARGE, highlighting that TBG misses certain basins that JAMUN is able to sample.

---

Trijectory
Torsions
JAMUN
0.130 1 0.020
0.185 0.044
0.165 1: 0.030
0.177 1. 0.053
0.260 0.052
0.155 1: 0.063
0,028
0.105
TBG (20x shortec)
0,203 = 0.070
0.235 + 0.073
0.225 + 0.071
0.210 + 0.072
0,481 0.073
0.121 = 0.051
JAMUN
0.159
0.000
0.210 0.057
0.187 0.054
0.257
0.353 0.122
0.262
Reference x10x shorter)
0.100 0.0835
0.082 0,027
0.095 1: 0.025
0.234 0.088
0.832 0.087
0.286 1: 0.066
Reference (100x shorter)
0.227 0.062
0.251 0.060
0:210 + 0.051
0.411 + 0.131
0. 0.569 0.18
0.462 = 0.138
0.039
JAMUN
0.196 =: 01.027
0.196 0.013
0.197 +: 0.010
0.336 1: 0.049
0.440 1:00::000 0.048
0.250 1: 0.075
shorter)
Ref. (100x shorter)
0.272 002
0.307 0.023
0.250 + 0.039
0.565 + 0.070
0.678 0.034
0.601 +: 0. 112
Boltz-1
0.125 0,033
0.4012 0.036
0.411 + 0.029
0.457 0.050
0.581 0,026
0.483 = 0.017
BioEmu
0.329 0.013
0,489 4 0.024
0,420 + 0.018
0.415 + 0.082 0.507 0.026
0.321 + 0.018
Table I: Comparison of Jenson-Shannon distances between JAMUN and baselines, both benchmark models and reference MD trajectories, averaged over test peptides in each dataset. Baselines are shortened represent comparable sampling time Jamun. First, compared with TBG for TIMEWARP 2AA-LARGE. Note that TBG (20x shorter) has similar sampling time as JAMUN. Second, with MDgen and MD trajectories (shortened by factor of 10 and 100) for the MDGEN 4AA-EXPLICIT. Third, for generalization, reference MD (shortened by factor of 10 and 100), Boltz-1 and BioEmu, for three test peptides in UNCAPPED 5AA.
a)
KG
LW
QG
ST
b)
1/2
Reference
11
JAMUN
71/2
H
TBG
-1/2
9
(20x chortord
T8G
-102
-2-22
m/2
-2-02
m/2
-2-02
m/2
-0-12
m/2
Figure 5: Results for 4 randomly chosen test peptides TIMEWARP SAA-LARGE: (a) Ramachan- dran plots, and (b) TICA-0,1 projections.
6.2 Results on MDGEN JAA-EXPLICH
Table I shows that JAMUN is very competitive with MDGen the JSD metrics. In fact, the TICA projections in Figure 6) shows that MDGen missing some basins that JAMUN is able to sample, while the Ramachandran plots show example where JAMUN hallucinates basin.
Figure 7 shows the significant speedups in backbone and sidechain torsion decorrelation times for JAMUN compared to the reference MD data. This matches our intuition about taking 'larger' integrator steps in the smoothed manifold of the noisy latent space.
We perform an analysis of the physical validity and the energy of JAMUN samples in Appendix B In summary, we find that JAMUN samples pass all Posebusters [12] checks at of 94.7% and have energies close to the ground truth MD for randomly selected 20 unseen peptides.
6.3 Assessing Generalization over Peptide Lengths with UNCAPPED 5AA
JAMUN's graph neural network architecture enables to operate on molecules of larger sizes than it was originally trained on. Thus, test whether JAMUN generalize to peptides of lengths

---

LIRH
Reference
0
R/2
Reference (100x shorter)
m/2
D
1/2
f
n/2
JAMUN
n/2
MDGen
11/2
-11
121120 m/2 n P:
1/2
m/2
n/2 m/2
Figure 6: MDGEN 4AA- EXPLICT results: (left) TICA-0,1 projections for randomly chosen test peptides; (right) Ramachandran plots for JAMUN and MDGen on a randomly chosen test.
125
150
100
Frequency
Frequency
50
25
0
10 10-2
10-1
109
10a
102
10a
10
3
10
10
109
10a
102
10a
Speedup Factor
Speedup Factor
Figure 7: MDGEN 4AXIIIT results: Speedup defined ratio of decorrelation time his- togrammed for all test backbone (left) and sidechain (right) torsions.
beyond its training set. This is challenging task, and one that believe conformational generation models have not been adequately benchmarked on. Unfortunately. neither TBG nor MDGen transfer to UNCAPPED 5AA due to fixed-length positional embeddings, despite our best efforts. In particular, we initialized the positional embeddings to support longer peptides, but this resulted in broken topologies in the resulting samples. Instead, choose Boltz-1 and BioEmu which support sampling on UNCAPPED 5 A to compare against JAMUN trained on TIMEWARP 4-A-LARGE.
Surprisingly, we find that the JAMUN model trained only 4AA peptides accurately predict ensembles for 5AA peptides. Figure 8b and Figure 11 show that JAMUN able to recover most states and even reproduce relative probabilities. Interestingly, the same experiment does not work if we train on TIMEWARP AAR-LARGE instead, suggesting that the 2AA reference MD data may not be informative enough generalize from.
the other hand, we find that Boltz-1 unable to sample the diversity of peptide conformations. This is not entirely surprising Boltz-1 was not trained on any MD data, as we noted before. Further, Boltz-1 also utilizes a common pair representation across all diffusion samples, as computed by its Pairformer stack. The pair representation intuitively represents the residue-wist distance matrix,

---

Refeeence
JMUN
Holmu
Reforence
m12
112
Figure 8: (left) Comparing noise sensitivity for an example test peptide GCSL for JAMUN, sampled identically, showing the tradeoff between slower mode mixing at 0.2Å) and broken topologies 0.&A. (right) TICA-0,1 projections for three test peptides in UNCAPPED 5AA.
and thus encodes significant portion of the geometry. Keeping this representation fixed possibly prevents the sampling of large conformational changes.
Surprisingly, BioEmu also seems struggle in this setting, even when considering distributions of backbone torsion angles only. This suggests that BioEmu cannot capture the relative flexibility of smaller peptides, even when trained on MD data for much larger proteins.
Quantitatively, Table I shows that JAMUN significantly outperforms Boltz-1 and BioEmu on the metrics from Section 5 As seen in Figure 3 JAMUN is roughly 5x faster than Boltz-1, and is roughly 60x faster than BioEmu when we perform side-chain reconstruction with H-Packer.
7 Conclusion
We present JAMUN, walk-jump sampling model for generating ensembles of molecular confor- mations, outperforming the state-of-the-art TBG model, and competitive with the performance of MDGen with no protein-specife parametrization. This represents an important step toward the ultimate goal of transferable generative model for protein conformational ensembles. Performing MD in the noised space gives the model clear physics interpretation, and allows faster decorrelation (and hence, sampling) than classical MD.
The model has some limitations that motivate future work. While is highly transferable in the space of two to five amino acid peptides, scaling is likely to require more exploration and intensive data generation in future work. Additionally, while the current -qquivvrra denoiser architecture works well, further development of the denoising network could speed up sampling. Alternative jump methods, such as multiple denoising steps (à la diffusion), could also serve to sharpen generation. Lastly, promising direction that has not yet been explored is the application of classical enhanced sampling methods, such metadynamics, for traversing the noisy space.
10

---

References
[1] J. Abbammson, J. Adler, J. Dunger, R. Evans, T. Green, A. Pritzel, O. Ronneberger. L. Willmore, J. Ballard, J. Bambrick, S. W. Bodenstein, D. A. Evans, C.-C. Hung. M. O'Neill, D. Reiman, K. Tunyasuvunakool, Z. Wu, A. Zemgulyté, E. Arvaniti, C. Beattie, O. Bertolli, A. Bridgland, A. Cherepanov, M. Congreve, A. I. Cowen-Rivers, A. Cowie, M. Figurnov, F. B. Fuchs, H. Glad- man, R. Y. A. Khan, C. M. R. Low, K. Perlin, A. Potapenko, P. Savy, S. Singh, A. Stecula, A. Thillaisundaram, C. Tong, S. Yakneen, E. D. Zhong, M. Zielinski, A. Zídek, V. Bapst, P. Kohli, M. Jaderberg, D. Hassabis, and J. M. Jumper. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493 500, Jun 2024. ISSN 1476-4687 doi: 1210588/41186--04-07444533
httphtttt:///bi.rrrg/// 1ttp////111111111111111111 /t0p:/0000000011116--00-00-77770
[2] V. Agarwal and A. C. McShan. The power and pitfalls of alphafold2 for structure predic- tion beyond rigid globular proteins. Nature Chemical Biology, (008)955-9959) Aug 2024. ISSN 1552-4469. doi: 10. 10.000/441199-02---6666333 941689-024-01638-
URL. http::::/ddoo..rr////// 1038/
[3] M. S. Albergo, N. M. Boffi, and E. Vanden- Eijiden. Stochastic interpolants: A unifying framework for flows and diffusions, 2023. URL
[4] Ansel, E. Yang, H. He, N. Gimelshein, A. Jain, M. Voznesensky, B. Bao, P. Bell, D. Berard, E. Burovski, G. Chauhan, A. Chourdia, W. Constable, A. Desmaison, Z. Vito, E. Ellison,
W. Feng, J. Gong, M. Gschwind, B. Hirsh, S. Huang, K. Kalambarkar, L. Kirsch, M. Lazos, M. Lezcano, Y. Liang, J. Liang, Y. Lu, C. Luk, B. Maher, Y. Pan. C. Puhrsch, M. Reso, M. Saroufim, M. Y. Siraichi, H. Suk, M. Suo, P. Tillet, E. Wang, X. Wang, W. Wen, Zhang. X. Zhao, K. Zhou, R. Zou, A. Mathews, G. Chanan, P. Wu, and S. Chintala. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '2). ACM, Apr. 2024. doi: 10.1145/3620665. pdf
3640366. URL http:://ptttoooeee
org/ wwsets/pytorchhe.......oo
[5] A. Antoszewski, C.-J. Feng, B. P. Vani, E. H. Thiede, L. Hong, J. Weare, A. Tokmakoff, and A. R. Dinner. Insulin dissociates by diverse mechanisms of coupled unfolding and unbinding. The Journal of Physical Chemistry B, 124(27)55711-558 2020. doi: ww..///cs.jpcb.eee/ee PMID: 32515958.
URL httth:ttt/////t /doi. org/ 10. 1021/ acs. app.ac...pcbb............. 0c03521
[6] A. Aranganathan, X. Gu, D. Wang. B. Vani, and P. Tiwary. Modeling Boltzmann weighted structural ensembles of proteins using Al based methods. 2024.
[7] M. Arts, V. Garcia Satorras, C.-W. Huang, D. Zugner, M. Federici, C. Clementi, F. Noé, R. Pinsler, and R. van den Berg. Two for one: Diffusion models and force fields for coarse- grained molecular dynamics. Journal of Chemical Theory and Computation, 1911666155,66599 2023.
[8] S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. Smidt, and B. Kozinsky. E(3)-qquivariant graph neural networks for data- and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022. ISSN 2041-1723 doi: URL https://:///d orgg 10. 1038/ /414//4442772222 299399-
10. 1..1038/41446--02229933
[9] H. M. Berman, J. Westbrook, Z. Feng, G. Gilliland, N. Bhat, H. Weissig, N. Shindyalov, and P. F. Bourne. The Protein Data Bank. Nucleic Acids Revearch, 28 1 2422 01 2000. ISSN 0305-1048.2
doi: 10.1093/nar/28.1 1.235. URL httphhtss::// ftttps://dol: org/ 10. /0w0///////// 1.25
[10] D. W. Borhani and D. E. Shaw. The future of molecular dynamics simulations drug discovery. Journal of computer-aided molecular design, 26: 16:15-26, 2012.
[11] G. R. Bowman. AlphaFold and Protein Folding: Not Dead Yet! The Frontier Is Conformational Ensembles. Annual Review of Biomedical Data Science, 7, 2024.
[12] M. Buttenschoen, G. M. Morris, and C. M. Deane. Posebusters: Ai-based docking methods fail generate physically valid poses generalise to novel sequences. Chemical Science, 15(9): 3130-3139, 2024. ISSN 2641-6539, doi: 10. 10.1003/dscc04130 10. 10.1009/D33004185A
URL/http://dl doi org/
[13) D. Chakravarty and I.. I.. Porter. AlphaFold2 fails to predict protein fold switching. Protein Sci, 31(6):e433, June 2022.
11

---

[14] S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky. K. T. Schütt, and K.-R. Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, (5)x1600015, 2017. doi: 10.1 126/sciadv. 1603015. URL htt.httpss:/// htopp:://www... http:///wwwwsc.nceeee.c org/doi/
abs/10. /116//sciadi 1603015
[15] G. Colombo. Computing allostery: from the understanding of biomolecular regulation and the discovery of cryptic sites to molecular design. Current Opinion in Structural Biology, 83:
102702, 2023.
[16] T. Darden, D. York, and I.. Pedersen. Partiele Mesh Ewald: An N log(N) method for Ewald sums in large systems. The Journal of Chemical Physics, 98( 12:1100000111099, June 1993. doi: 00.1063/1.464399 URL https://doi org/ 10. .0t/:/1.66... 146377
[17] D. del Alamo, D. Sala, H. S. Mchaourab, and J. Meiler. Sampling alternative conformational states of transporters and receptors with alphafold2. eLife, 11:e75751, mar 2022. ISSN 2050-084X. doi: 10. 7/755///Life.75511 URL httpss///// org/ 10. 7554/eLife 75751
[18] A. R. Dinner, J. C. Mattingly. J. O. B. Tempkin, B. Van Koten, and J. Weare. Trajectory stratification of stochastic dynamics. SIAM Rev Soc Ind Appl Math, 60(4):909-99) Nov. 2018.
[19] A. dos Santos Costa, I. Mitnikov, F. Pellegrini, A. Daigavane, M. Geiger, Z. Cao, K. Kreis, T. Smidt, E. Kucukbenli, and J. Jacobson. Equijump: Protein dynamics simulation via so(3)- equivariant stochastic interpolants, 2024. htttttttt//xxiii///
24100006677
[20] A. Dumitreseu, D. Korpela, M. Heinonen, Y. Verma, V. Iakovlev, V. Garg, and H. Lähdesmäki. Field-based Molecule Generation, 2024. URL httph:tt:///ttt http:///axxvvv.. org/ abs///////22/// 15864
[21] P. Eastman, J. Swails, J. D. Chodera, R. T. McGibbon, Y. Zhao, K. A. Beauchhmp, L.-P. Wang. A. C. Simmonett, M. P. Harrigan, C. D. Stern, R. P. Wiewiora, R. Brooks, and V. S. Pande. OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. PLOS Computational Biology, 13(7)xe July 2017. doi:
10056.0
1005659. URL
https://doi org 10. 1371/journal pcbi. 1005659
[22] W. Falcon and The Torch Lightning team. PyTorch Lightning, Mar. 2019. URL https http://iiiaaiiiinnAA//////gele
[23] M. L.
N. D. Pomarici, A.-L. M. Fischer, V. J. Hoerschinger, K. B. Kroell, J. R. Riccabona, A. S. Kamenik, J. R. Loeffler, J. A. Ferguson, H. R. Perrett, et al. Structure and dynamics guiding design of antibody therapeuties and vaccines. Antibodies, 12(4):67, 2023.
[24] N. C. Frey, D. Berenberg, J. Kleinhenz, I. Hotzel, J. Lafrance- ipal, S. Ra, R. Bonneau, K. Cho, A. Loukas, V. Gligorijevic, and S. Saremi. Protein discovery with discrete walk-jump sampling. In International Conference on Learning Representations, 2024.
R. L. Kelly, Y. Wu, A. Ra-
[25] M. Geiger and T. Smidt. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022.
[26] N. R. Gough and C. G. Kalodimos. Exploring the conformational landscape of protein kinases. Current Opinion in Structural Biology, 88: 102890, 2024.
[27] C. A. Grambow, H. Weir, C. N. Cunningham, T. Biancalani, and K. V. Chuang. Ringer: Rapid inference of macrocyclic peptide conformational ensembles. arXiv preprint, arXiv:2310.01234,
2023. URL https: //axivv. org/ httg//b//231000 01234
[28] C. A. Grambow, H. Weir, C. N. Cunningham, T. Biancalani, and K. V. Chuang. Cremp: Conformer-rotamer ensembles of macrocyclic peptides for machine learning. Scientific Data, 11, 2024. doi: 10. 10.10388/441577-024-0000
URL https://doi. org/10. 1038/
1415977024-0369887
[29] X. Guan, Q.-Y. Tang, W. Ren, M. Chen, W. Wang, P. G. Wolynes, and W. Li. Predicting protein conformational motions using energetic frustration analysis and alphafold2. Proceedings of the National Academy of Sciences, 121 121(35))22410662222 2024. doi: 10. pnas. http///i/aab///// 1073/ /ptas//ass.... 2410666222
[30] K. Willman and D. Kern. Dynamic personalities of proteins. Nature, 450(7172); 964-972, Dec. 2007.
[31] B. Hess, H. Bekker, H. J. Berendsen, and J. G. Fraaije. Lines: a linear constraint solver for molecular simulations. Journal of computational chemistry, (8(12):14633147) 1997.
12

---

[32] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models, 2020. URL https //arxiv. wwrggabs/2006... 11239
[33] E. Hoogeboom, G. Satorras, C. Vignac, and M. Welling. Equivariant Diffusion for Molecule Generation in 3D, 2022. URLhttpss: //t:xi//aiivvee// orgo 2abs/2203 072033170033 17003.
[34] T. Hsu, B. Sadigh, V. Bulatov, and F. Zhou. Score dynamics: Scaling molecular dynamics with picoseconds time steps via conditional diffusion model. Journal of Chemical Theory and Computation, (006))23335-3444 2024.
[35] X. Huang, R. Pearce, and Y. Zhang. FASPR: open-source tool for fast and accurate protein side-chain packing. Bioinformatics, 36( (0(11)33378-3311 3765, June 2020.
[36] B. Jing, B. Berger, and T. Jaakkola. AlphaFold meets flow matching for generating protein ensembles. arXiv preprint urXiv:2402.04844 2024.
[37] B. Jing, H. Stärk, T. Jaakkola, and B. Berger. Generative modeling of molecular dynamics trajectories, arXiv preprint arXiv:2409. 17008, 2024.
[38] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool,
R. Bates, A. Zidek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, J. Ballard, A. Cowie, Nikolov, R. Jain, J. Adler, T. Back, S. Petersen, D. Reiman, E. Clancy, M. Zielinski, M. Steinegger, M. Pacholska, T. Berghammer. S. Bodenstein, D. Silver, O. Vinyals, A. W. Senior, K. Kavukcuoglu, P. Kohli, and D. Hassabis. Highly accurate protein structure prediction with alphafold. Nature, (96(7873)5583-588) Aug 2021. ISSN 1476-4687 doi: 10.1038/s41586- 021-00819---00 URL URL httpss://oo. org/10. 1038/ s41586- 021 001--666--01-------222
B. Romera-Paredes.
[39] A. Jussupow and structure prediction models. Journal of Chemical Theory and Computation, 19(7):1965 1975, 2023. doi: 10. 102100224
I. Kaila. Effective molecular dynamics from neural network-
URL
2c01027] PMID: 36961997.
[40] W. Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A, (2(3)22(222299)))
Sep 1976. doi: 10.1107/505677347700018733 URL https://doi
org/10. 1107 01077//0056773976001188888
[41] T. Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the Design Space of Diffusion-Bart Generative Models. In Proc. NeurlPS, 2022.
[42] T. Karras, M. Aittala, J. Lehtinen, J. Hellsten, T. Aila, and S. Laine. Analyzing and Improving the Training Dynamics of Diffusion Models. In Proc. CVPR, 2024.
[43] S. Kieninger and B. G. Keller. GROMACS Stochastic Dynamics and BAOAB Are Equivalent Configurational Sampling Algorithms. Journal of Chemical Theory and Computation, 18(10): 5792-5798, 2022.
[44] D. Kingma and J. Ba. Adam: method for stochastic optimization, 2017. URL https
7/ /arxiv.. www//rr/////eeeeeeee.eeeeeeeeeee abs/ 6980
[45] L. Klein and F. Noé. Transferable Boltzmann Generators. arXiv preprini arXiv:2406.1 14426,
2024.
[46] L. Klein, A. Foong, T. Fjelde, B. Mlodozeniec, M. Brockschmidt, S. Nowozin, F. Noé, and R. Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time- coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
[47] I.. Klein, A. Krämer, and F. Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024.
[48] M. Lazou, O. Khan, T. Nguyen, D. Padhorny, D. Kozakov, D. Joseph-McCarthy, and S. Va- jda. Predicting multiple conformations of ligand binding sites in proteins suggests that al- phafold2 may remember too much. Proceedings of the National Academy of Sciences, 121(48): e2412719121, 2024. doi: 10.
URL https://: pnas. org/doi/
abs/10. 1073 1 2412719121
[49] B. Leimkuhler and C. Matthews. Rational Construction of Stochastic Numerical Methods for Molecular Sampling. Applied Mathematics Research eXpress, 2013(1):34-56 June 2012. ISSN 1687-1200. doi: 10. /m URL htt.thtttt://ttttoo org/10../1. 1093/amrx/ abs010 -eprint:
13

---

[50] B. Leimkuhler and C. Matthews. Robust and efficient configurational molecular sampling via langevin dynamics. The Journal of Chemical Physics, 17):111174,,,, May 2013. doi: N0.1063//1.48 .4802990. Ittpp:///:/ /ttttpp::/ddo..or////0. 1063/10611111 44802990
[51] B. Leimkuhler and C. Matthews. Molecular Dynamics: With Deterministic and Stochastic Nu- merical Methods. Number 39 in Interdisciplinary Applied Mathematics. Springer International Springer, Cham, 1st ed. 2015 edition, 2015. ISBN 978-3-319-16375-8
Publishing: Imprint:
[52] S. Lewis, T. Hempel, J. Jiménez-Luna, M. Gastegger, Y. Xie, A. Y. K. Foong, V. G. Satorras, O. Abdin, B. S. Veeling, I. Zaporozhets, Y. Chen, S. Yang. A. Schneuing. Nigam, F. Barbero, V. Stimper, A. Campbell, J. Yim, M. Lienen, Y. Shi, S. Zheng, H. Schulz, U. Munir, C. Clementi, and F. Noé. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv, 2024, doi: 10. 1 10111024...11 10.100/2224..2.055 2.665.2.
[53] Z. Lin, H. Akin, R. Rao, B. Hie. 2. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. dos Santos Costa, M. Fazel-Zarandi, T. Sercu, Candido, and A. Rives. Exolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6676:)666 1130, 2023. doi: 10. 1 226/sciiece..dee25544 URL https: //www. science. oraa/////orr////////////// 1126/ sersncceenc..... ade2574
[54] Y. Lipman, R. T. Q Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. URL
https://openrevian
[55] J. Lu, B. Zhong. 7. Zhang, and J. Tang. Str2str: A score-based framework for zero-shot protein conformation sampling. In The Twelfth International Conference on Learning Representations, wett/oormmiddCABBi..
2024. URL L httpstt/
[56] S. J. Marrink, H. J. Risselada, S. Yefimov, D. P. Tieleman, and A. H. de Vries. The martini force field: Coarse grained model for biomolecular simulations. The Journal of Physical Chemistry 1,1 111(27)7812--8224 Jul 2007, ISSN 1520-6106. doi: 10. 10.1021/jp071099 hhtpps//doi. 0tg/10.1100 10210 1021///7797
URL
[57] M. D. Miller and G. N. Phillips. Moving beyond static snapshots: Protein dynamics and the Protein Data Bank. Journal of Biological Chemistry, 296, 2021.
[58] K. Miyasawa. An Empirial Bayes Estimator of the Mean of a Normal Population. Bulletin de I' I'Institt international de statistique., 38(4):181-188 1960.
[59] L. Molgedey and H. G. Schuster. Separation of mixture of independent signals using time delayed correlations. Phys. Rev. Lett., 72:3634-3377 Jun 1994. doi: 10. 1 103/PhysRevleett 3634. Uttpttsss:: //1ink. aps
72.3634
[60] Noé, S. Olsson, J. Köhler, and H. Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457)) 147, 2019.
[61] G. Pérez-Hernández, F. Paul, T. Giorgino, G. De Fabritiis, and F. Noé. Identification of slow molecular order parameters for markov model construction. J Chem Phys, 139(1):015102, July 2013.
[62] P. O. Pinheiro, A. Jamash, O. Mahmood, V. Sresht, and S. Saremi. Structure-based Drug Design by Denoising Voxel Grids. arXiv preprint arXiv:2405.0334 2024.
[63] O. Pinheiro, J. Rackers, Kleinhenz, M. Maser, O. Mahmood, A. Watkins, S. Ra, V. Sresht, and S. Saremi. 3D molecule generation by denoising voxel grids. Advances in Neural Informa- tion Processing Systems, 36, 2024.
[64] W. Ponder and D. A. Case. Force fields for protein simulations. ISSN 0065-3233.
66:27-85, 85, 2003. URL https
doi: http://dooooyy/////ooo00oo222210
/www. httpa//ceeeeeeeo.eee.ee.ree////n
[65] Pracht and S. Grimme. Automated exploration of the low-energy chemical space with fast quantum chemical methods. Physical Chemistry Chemical Physics, 22(14):7169- 7192, 2020. doi: 10. C09/C/669CC
URL htttps://pube www.bb..rc.orro.. wwtpg//////////////////o
articlellanding222220 cp/ cp/c9c068699
[66] H. Robbins. An Empirical Bayes Approach to Statistics. In Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, Volume I: Contributions to the Theory of Statistics, volume 3.1, 1956.
14

---

[67] J. P. Roney and S. Ovchinnikov. State-of-the-ane estimation of protein model accuracy using alphafold. Phys. Rev. Lett. 129:238101, Nov 2022. doi: 10. 10.100//hhyseewccco 1ett 129.238 URL. Jtt.pttss::///ink aps org/ wwrgg/di//10 1103/hhysRevLettr 129. 2381011
[68] M. Sachs, B. Leimkuhler, and V. Danos. Langevin dynamics with variable coefficients and nonconservative forces: from stationary states to numerical methods. Entropy, 19(12):644
2017.
[69] Saldaño, N. Escobedo, J. Marchetti, D. Zea, J. Mac Donagh, A. J. Velez Rueda, E. Gonik, A. Gareía Melani, Novomisky Nechcoff, M. N. Salas, T. Peters, N. Demitroff, S. Fernandez. Al- berti, N. Palopoli, M. S. Fornasari, and G. Parisi. Impact of protein conformational diversity on alphafold predictions. Bioinformatics, 38(10):7742 doi: 10.1093. (093//oonffrmmatissta22222222222 URL httpt:ttt/////t/ org/ 10. http:///ioiooooaaiicccaoooo btacc202
04 2022. ISSN 1367-4803.
[70] S. Saremi and A. Hyvärinen. Neural Empirical Bayes. Journal of Machine Learning Research, 20(181):1-23, 2019.
[71] V. G. Satorra, F. Hoogeboom, and M. Welling. E(n) Equivariant Graph Neural Networks, 2022.
URL http:///aaxx.i..rrr////2/022... 09844
[72] M. K. Scherer, B. Trendelkamp-Schronk F. Paul, G. Pérez-Hermández, M. Hoffmann, N. Plattner, C. Wehmeyer, J..H. Prinz, and F. Noé. Pyemma 2: A software package for estimation, validation, and analysis of markov models. Journal of Chemical Theory and Computation, 11(11):5525- 5542. 2015, doi: 10. 102 (022lacc.jccc.550000 URL htt.://0044.....ttts///ooooooo// 10t1//00...11//..... jctc
5b00743 PMID: 26574340.
[73] M. Schreiner, O. Winther, and S. Olsson. Implicit transfer operator leamning: Multiple time- resolution surrogates for molecular dynamics, 2023. URL http::t/aaaiii..rrrrrre http://arrii..orrr///2/200
18046]
[74] M. Schreiner, O. Winther, and S. Olsson. Implicit transfer operator learning: multiple time- resolution surrogates for molecular dynamics. arXiv preprint arXiv:2305. 2023.
[75] C. R. Schwantes and V. S. Pande. Improvements in markov state model construction reveal many Non-Native interactions in the folding of NTL9. J Chem Theory Comput, 244:22000- Apr. 2013.
[76] K. T. Schütt, F. Arbabzadah, S. Chmiela, K. R. Müller, and A. Tkatchenko. Quantum-chemical insights from deep tensor neural networks, Nature Communications, 8(1), Jan. 2017. ISSN 2041- 1723. doi: 10. 10.1038//comms 13890. URL //dp::/tt//xx..../ //xx. www.ooogo........
org/10.. htte//0.1038//110 1038//commms18899
[77] A. W. Senior, R. Evans, J. Jumper, J. Kirkpatrick, L. Sifre, T. Green, C. Qin, A. Židek, A. W. R. Nelson, A. Bridgland, H. Penedones. S. Petersen, K. Simonyan, S. Crossan, P. Kohli, D. Jones, D. Silver, K. Kavukcuoglu, and D. Hassabis. Improved protein structure prediction using potentials from deep learning. Nature, 577(7792):706-770 Jan 2020. ISSN 1476-4687, doi: 10.1038/441586-11 (100:4411866-91--111 119-197777199 URI. httpt::////ooii org/ 10. 1038/ 101111008118800
[78] D. E. Shaw, P. Maragakis, K. Lindorffi-Lareen, S. Piana, R. O. Dror, M. P. Eastwood, J. A. Bank, J. M. Jumper, J. K. Salmon, Y. Shan, and W. Wriggers. Atomic-level characterization of the structural dynamics of proteins. Science, 330066002 2010. doi: 10.1126/science. 1187409
1187409. JtLpsttp:: http:://wwwwweoo www.://cience.c wtt.://cicce....ro//aaaaaa//////
[79] J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models, 2022. URL https
02502
[80] M. Steinegger and J. Söding. MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechnology, 35(11):1026-1028, Nov 2017. ISSN 1546-1696, doi: 10. 1038//8/n..33988. URL /to. htt.:dii.ooooggo//ee org//0. www0//nt........ 3988
[81] N. Thomas, T. Smidt, S. Kearnes, I.. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rntation-and
-enuslationeuuuariiinn neural networks for 3d point clouds. arXiv 2018.
preprint
[82] S. Umeyama. I Least-qquares estimation of transformation parameters between two point patterns. IEEE Transactions Pattern Analysis and Machine Intelligence, 13(4)337-380, 1991. doi: 10.1109/34.883 [83] B. P. Vani, J. Weare, and A. R. Dinner. Computing transition path theory quantities with trajectory stratification, J Chem Phys, 157(3):034106, July 2022.
15

---

[84] V. Vapnik. Principles of risk minimization for learning theory. In J. Moody, S. Hanson, and
R. Lippmann, editors, Advances in Neural Information Processing Systems, volume 4. Morgan- Kaufmann, 1991. http///ttttss/ppcceeeingg nww..is......pp...... httpp/ppppppllepeeeeeeeaaaapepee www.s//pppeeeeeeeeee/eeeaee 1991/ fia2ff444fffafffffcccoo2222aaaaf
pdff
[85] M. Varadi, D. Bertoni, P. Magana, U. Paramval, I. Pidruchna, M. Radhakrishnan, M. Tsenkov, S. Nair, M. Mirdita, J. Yeo, O. Kovalevskiy, K. Tunyasuvunakool, A. Laydon, A. Zidek, H. Tomlinson, D. Hariharan, J. Abrahamson, T. Green, J. Jumper, E. Birney, M. Steinegger, D. Hassabis, and S. Velankar. Alphafold protein structure database in 2024: providing structure coverage for over 214 million protein sequences. Nucleic Acids Research, 5201)):3688D33 11 2023. ISSN 0305-1048. doi: 10. 1..09/9//r/gaad0011.. URL https:://doi.org/100 1093/ nar/ /k1d/1111011
[86] G. M. Visani, W. Galvin, M. Pun, and A. Nourmohammad. H- H-packer: Holographic rotationally equivariant convolutional neural network for protein side-chain packing. In D. A. Knowles and Mostafavi, editors, Proceedings of the 18th Machine Learning in Computational Biology meeting, volume 240 of Proceedings of Machine Learning Research, pages 230-249. PMLR, html
30 Nov-01 Dec 2024. URI. httia:::/pproceeddnnn
mlr
[87] A. Vitalis and R. Pappu. Methods for monte carlo simulations of biomacromolecules. Annual reports in computational chemistry, 5:49-76, 2009.
[88] H. K. Wayment-Steee -taeele, A. Ojoawo, R. Otten, J. M. Apitz. W. Pitsawong, M. Hömberger, S. Ovchinnikov, L. Colwell, and D. Kern. Predicting multiple conformations via sequence clustering and alphafold2. Nature, 625(7996):332-83) Jan 2024. ISSN 1476-4687. doi: 10. 10.108///4586-000 0300://41118800--686683 023-02232--2-- URL patps://///8 org/00 1038/s41586 023 023-06832-9
[89] Wohlwend, G. Corso, S. Passaro, M. Reveiz, K. Leidal, W. Swiderski, T. Portnoi, I. Chinn, J. Silterra, T. Jaakkola, and R. Barzilay, Boltz-1 democratizing biomolecular interaction modeling. bioRxiv, 2024. doi: 10.1101/2024.1111..2441111
httptttst//:/w.bbii
19. 624167
[90] Wu and Z. Li. Diffimd: A geometric diffusion model for molecular dynamics simulations,
2023. URL URL httpsttt/tx/xi..../ http://www.rr/abb//2..oooo
[91] L.-E. Zheng, S. Barethiya, E. Nordquist, and J. Chen. Machine learning generation of dynamic protein conformational ensembles. Molecules, 28(10)44047, 2023.
[92] Zheng, He, C. Liu, Y. Shi, Z. Lu, W. Feng, F. J. Wang, J. Zhu, Y. Min, et al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, pages 1-10, 2024.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes] In the abstract, our claims that JAMUN is able generate ensembles faster than traditional methods and recent ML methods, and can generalize outside its training data.
Figure 3 shows the comparison of sampling times for JAMUN and all baselines. Table I and Table shoow that JAMUN is able to outperform TBG, Boltz-1 and BioEmu in terms of efficiency in sampling the conformational space in total wall-clock time. Table I shows that JAMUN is competitive with the recently proposed peptide-specific MDGen model in terms of sampling quality. [Section 6.3 shows that JAMUN is able to generalize to 5AA peptides, even when trained only on 4AA peptides, which is out of the reach of models such as TBG and MDGen. For these 5AA peptides, JAMUN is able outperform Boltz-1 and BioEmu which were trained on crystal structures and MD data of much larger protein structures respectively.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors? Answer: [Yes] The key limitation of the paper is described in Section 7 that JAMUN is transferable across 2-5 amino acid long peptides, but the model has not been tested on longer peptides and full proteins. The -quivaaquiaa denoiser works well, but further development of the denoising network could speed up sampling to match up to MDGen.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
[Yes] The main contribution of our work is not theoretical. We largely build upon the theoretical framework of Walk-Jum Sampling developed by [70]] For completeness, we re-derive all key theoretical results in Apppnddix H and re-derive the Langevin dynamics update in Appendix J
Answer:
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes] We provide details about the training and sampling of JAMUN in Section 2.2] with details of the denoiser in [Section 2.4] and Appendix We also provide code to reproduce our results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes] The code was submitted in accordance with the (https://nips cc/ httpe//guidee//oSSumiiiieeeeeeee
which includes scripts to reproduce the re- sults in the paper. Section 3 includes the description of the data used within the paper. The TIMEWARP 2 A-LARGE, TIMEWARP 4AA-LARGE, and MDGEN 4-X-IITLI datasets are available to the public under an MIT licence. We are the legal process of sharing the CAPPED 2 dataset for public use.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes] The data splits are described in (Section 3 and were chosen to allow comparison to existing baseline models. The parametrization and the hyperparameters of our model are described in Section 2.4 2.4|and Appendixx F We also provide code to reproduce our results.
17

---

7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Ys]] Our metrics averaged over the test dataset, where report the mean and I-sigma errors for all JSD metrics.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer:
[Yes] We provide the compute details for training JAMUN in Section and sampling JAMUN and all baseline models in Figure 3
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with NeurlPS Code of Ethics https:: ctcp:ppp.ccpppllll cctpubicc/cccccccccccccccccccclt htt.bl/ic//hhccccuudleeeeeeeee
Answer: [Yes] We have made our best effort to conform to NeurIPS Code of Ethics.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA] This paper presents work whose goal is to advance the field of accelerated sampling of protein conformations. Due the niche field of study, we do not believe that there any negative societal impacts that arise directly from our work.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA] The authors do not foresee any ways in which this model could be misused.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes] Yes, we credit all authors for creating the benchmark datasets: Klein et al. [46] for TIMEWARP 1AA-LARGE and TIMEWARP 4AA-LE, and Jing et al. [37] for MDGEN 4AA-EXPLICIT available under the MIT license. We have respected their terms of use.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes] The new CAPPED 2AA dataset is described in Section and Appendix C
We are in the legal process releasing this dataset for public use.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA] This paper did not involve crowdsourcing or research with human subjects..
15. Institutional review board (IRB) approvals or equivalent for rescarch with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
18

---

Answer: [NA] This paper does not involve crowdsorcing or research with humans. The data was all collected in silico.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA] The authors did not use LLMs the core research methods.
19

---

A Comparison of Walk-Jump Sampling Against Full Diffusion
Section contains a comparison of JAMUN with existing baselines, but does not directly address the comparison of the techinique of walk-jump sampling with standard diffusion. We perform this comparison here; we fix the JAMUN architecture and directly compare walk-jump sampling (at single noise level) to full diffusion (which samples range of noise levels from very large to very small).
We train a diffusion model with the same JAMUN architecture on TIMEWARP 2AA. We then sample from this model using both diffusion (specifically, the ODE sampler with noise schedule of 64 steps from 0.01 Á to 10Å using the Heun second order method as recommended by EDM 441) and walk-jump sampling. We compare the averaged over the TIMEWARP 2AA test set as function of number of samples and number of function evaluations (NFE):
divergence (JSI) of the backbone torsions
dffusion
walkjump
walkjmpp
ji Divergencee
10-1
Divergence
10-1
10n2
102
10a
10
10a
102
10a
10a
10a
107
samples
NFFE
Figure 9: Comparison between diffusion, walk-jump and the ground truth MD in terms of JS divergence of backbone torsions to the full ground truth MD distributions, as a function of (left) number of samples and (right) number of function evaluations.
Sampler
Number of Samples
NFE
JSD-Backhone Torsions +
Walk-Jump 3149
6298
0. 0.1501
Diffusion 3149
399923
0. 0.1363
Walk-Jump 200000
200000
400000
0.0496
Diffusion
25400000
0.0460
Table 2: Comparison of walk-jump and diffusion sampling with the same model. Diffusion obtains a better JSD metric than walk-jump sampling, but the cost of x: 30 more sampling time.
Table 2
contains a summary of the key comparison. Essentially, we find that walk-jump sampling is faster with only minor loss in fidelity; because works in partially noised space, instead of having to generate every sample from an uninformative Gaussian prior over many steps.
B Physical Validity and Energy Analysis
Here, we perform a physical validity analysis on the generated JAMUN samples, using the popular Posebusters [12] package. We randomly selected 20 unseen test peptides from the MDGEN 4AA- EXPLICIT dataset for this analysis.
We see that the bond lengths are correctly captured with high probability by JAMUN. Furthermore, the overall quality of the JAMUN samples is high. For finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:
Next, we compute the force field energies for JAMUN samples. We add hydrogen atoms using OpenMM's 's PDBFixer (21), and compute energies using the amber14 force field. We find that the energies of the JAMUN samples overlap well with those of the reference MD samples:
20

---

Posebusters Metric
Average Pass Rate
Valid Bond Lengths
97.0%
Valid Bond Angles
.9.4%
Internal Steric Clash
100.0%
Internal Energy
97.5%
Overall
94.7%
Table 3: Average pass rates for Posebusters metrics JAMUN samples.
Posebusters Pass Rate
90%
92%
94%
96%
98%
100%
Valid Bond Lengths
18
17
14
20
20
20
Valid Bond Angles
20
20
20
20
20
20
20
20
Internal Steric Clash
20
20
20
20
20
20
20
20
20
20
20
20
Internal Energy
20
19
19
17
10
Overall
20
20
20
20
20
20
Table 4: Empirical CDF pass rates across 20 test peptides.
C
Data Details
The TIMEWARP and MDGEN datasets consist of "uncapped" peptides, whose termini are zwitterionic amino and carboxyl groups, shown in the left panel of Fiiure 2) These are not ideal analogues of amino acids in proteins due local charge interactions well as lack of steric effects.
We also create a similar dataset called CAPPED 2AA of 2AA peptides by adding ACE (acetyl) and NME (N-methyl amide) caps, common practice in molecular dynamics simulations of very small peptides. As illustrated in the right panel of Figure 2] these caps introduce additional peptide bonds with the first and last residues. These peptide bonds remove the need for the zwitterion, while the methyl group provides some steric interactions. These capping groups increase the complexity of the modelling task, but ensure a more realistic distribution conformations. We choose the same splits as in TIMEWARP 2AA-LARGE. LARGE. Since we simulated this data ourselves, we can also measure the wall-clock speed-ups of JAMUN relative to MD on this dataset.
We ensure that our unbiased molecular dynamics runs are converged or representative by comparing using -quiluuiii Umbrella Sampling (NEUS) [88 83 a trajectory stratification based enhanced sampling algorithm. The protein is represented by the amber99abilda force field (64). The simulations are performed at 300 K with the BAOAB integrator 550 in OpenMM (21]: LINCS is used to constrain the lengths of bonds to hydrogen atoms [31]: Particle Mesh Ewald is used calculate electrostaties [16]: the step size
against biased molecular dynamics
fs. The systems solvated with TIP3P water models and equilibrated under NVT and NPT ensembles for 100ps each.
For the 2AA datasets, the training set consists of 50% of all possible 2AA peptides. For the 4AA datasets, the generalization task is much harder, because the number of 4AA peptides in the training sets is less than 1% and 2% respectively of the total number of possible 4AA peptides.
D Results for CAPPED 2AA
D.I Further Results on internally simulated data
Here, compute the ratio of the decorrelation times for the backbone and sidechain torsions in JAMUN and the reference MD data. Figure 10 highighiss how sampling in the smoothed space Y compared to the original space X enables much faster decorrelation.
Table we compare JAMUN to the reference MD, shortened by factor 10 along the variables mentioned in Section using the Jensen-Shannol distance to the full reference MD data. JAMUN
21

---

Sequence
MDGEN 4AA-EXPLICIT
JAMUN
FHSE
-75.6 20.1
FKKL
-99.0 24.0
-562.5 192.1
FLRH
-1272.7
-1198.5 129.3
FSDP
-697.5
1687.6991..
FSRK
-1333.0
-1349.9
GCIC
#...5.21.44
-538.8
GGHN
21.9
129.
GLIL
-743.3
-
HELI
25.5
HENV
1156.9+177
1123.4
HTIQ
-662.9 + 17.0
=
IAMI
428.0 + 15.3
4
IDRH
-4416.8 18.1
IHNV
- 21.1
-864.4 48.8
IMRY
12300. 23.6
-1100.6 + 207.0
INVH
-793.6 + 22.9
-745.3 128.6
IPGD
-611.5 + 15.0
-582.5 + 55.2
Table 5: Force field energies (in (kJ mol EXPLICIT and JAMUN.
comparison between samples from MDGEN 4AA-
Dataset
Peptide Longth
Capped?
Foree Field
Solvent Model
Tempenture
# Train
Validation
# Test
ZAA-LARGE CAPFED 2AA
anberld anber99sbsled anber 14
Explicit Water
300
200
80
100
TIMEWARD AMEEARP-AA-LLN 4AA-LARGE
Implicit Water
310
1459
379
182
MIGEN 4AA- EXPLICIT
anber14 14
Explicit Water
350 K
3109
100
100
UNCAPEED SAA
anberti
Implicit Water
CREMP 4AA
X
GPW2 xTB
Explicit Chloroform
300
15842
1000
1000
Table A short description of the simulation conditions across the different datasets.
outperforms this shortened MD trajectory across all metrics, though it takes approximately 2x longer to sample than JAMUN, from Figure 3
Backbone Torsions
Sidechain Torsions
All Torsions
Metastable Probs
Trijectory
TICA-0
TICA-0,1
JAMUN
0,291 + 0.119
0.320 + 0.108
0.304 + 0.112
0.351 = 0.130
0.438
0.117
0.264 + 108 0.543
Reference (10x shorter)
0.447
0.071
0.056
0.557
0.564
Table 7: Comparison of Jenson-Shannon Shannon distances between JAMUN and the reference MD (shortened
by a factor of 10). averaged over the test peptides in CAPPED A. Note that this shortened reference MD takes 2x longer to sample JAMUN.
E
Related Work
Protein structure prediction only requires prediction of a few folded states, and is usually trained crystallized structures of proteins which can be quite different from their native states. The size of these proteins are significantly larger than the peptides we study here. On the other hand, conformational ensemble generation requires many samples from the Boltzmann distribution to capture the effects of solvent atoms and intramolecular interactions, even if dynamics is not modelled explicitly. Research has found that protein structure prediction models tend struggle with capturing this diversity, especially when sampling the conformations of flexible domains [69 13 2 and rarer conformations [48] not found in the Protein Data Bank (PDB) Another issue is that the quality of these models' predictions can be dependent a preprocessing step where sequence databases are queried similar protein sequences which indicate evolutionary conservation patterns as supplementary input to the model. In fact, some of the first attempts to sample alternative conformations of proteins with AlphaFold2 were performed
multiple sequence alignment (MSA) information,
by subsampling [17], clustering [88 manipulating MSA information
[29]. Scquence-to-structie
22

---

60
80
Frequency
Frequency
40
40
20
20
0
10 3
10
-
10-1
100 10t 102
103
10
10 10-1 100
10a 102 Speedup Factor
10a
Speedup Factor
Figure 10: Speedups defined as the ratio between decorrelation times between the reference MD and JAMUN for backbone (left) and sidechain (right) torsions for all test peptides in CAPPED 2AA.
JACU Nowe
Natenacc
JAMAN
H422
Baota
DxOra
Figure 11: Ramachandran plots for test peptides NRLCQ. VWSPF and KTYDI in UNCAPPED 5AA.
which do not not require MSA information can faster but tend to produce less physically accurate structures, as noted by Lu et al. [55]. Unfortunately. MSA information can be quite unreliable for small peptides due to the presence
models such as ESMFold [53]
many hits. In fact, popular MSA software suites such as MMseqs2 [80 querying will, by default, simply return an empty MSA for the peptides we study here. Nevertheless, previous research (67 has shown that A Alpa2 can still learn an accurate energy function for protein structures without MSA co-evolution information.
Jussupow and Kaila [39] finds that AlphaFold2's predicted local distance difference test (pLDDT) and predicted aligned error (PAE) scores correlate with local protein dynamics and global conformational flexibility respectively. They use these scores parametrize an additional harmonic potential for coarse-graings MD with the MARTINI [56] force field. AlphaFlow [36] develops flo--machhnn over the quotient space of 3D positions modulo rotations learn a distribution over 3D positions of the 3-carbon atoms, outperforming MSA subsampling with AlphaFold2 over the Protein Data Bank (PDB) [9. Distributional Graphormer 99 parametrizes diffusion model over o-carbon positions,
also trained on the PDB. Str2Str (55 proposes a noising-denoising process along range of noise levels for sampling backbone atom coordinates, followed by regression of side-chain coordinates with the rotamer-based FASPR (35] packgae. BioEmu [52] is diffusion model built using the
EvoFormer stack from AlphaFold2 (38), BioEmu is pretrained on 200 million protein structures from the AlphaFold Protein Structure Database [85] and finetuned on over 200ms of MD data, which are orders of magnitude larger than the datasets we benchmark here.
A predecessor of many of the models discussed in [] "Two for One' 7 showed that the score learned by diffusion models can be used for running molecular dynamics simulations. However, as they choose the noise level for the score function close to 0, the molecular dynamics is effectively run in the original space X, not in the latent space as JAMUN does, which again limits the effective timestep of simulation.
23

---

There have also been efforts to build ML models for taking longer MD time-steps [46] 74 34 and for approximating conformations of large proteins [92 36]. These methods rely on hand-crafted featurizations (eg. backbone torsion angles). In practice, this has made generalization to unseen molecules challenging for these models as well. The above models are often classified either as Boltzmann Generators or Boltzmann Emulators. Boltzmann Generators are guaranteed to draw unbiased samples from the Boltzmann distribution, while Boltzmann Emulators do not have this guarantee. Strictly speaking, JAMUN is Boltzmann Emulator.
JAMUN is walk-jump sampling method which uses an -quiivrqui neural network for denoising. WJS is built on the seminal work of Neural Empirical Bayes [70], and has been used in voxelized molecule generation 62[63 and protein sequence generation [24]. Our work the first to our knowledge to apply walk-jump sampling to point clouds.
F Overview of Denoiser
The denoiser is a -quuivrrraa graph neural network. The graph is defined by a radial cutoff of 10Å in y. The overall computation performed by the denoiser shown in Figure 12 with the initial embedding, message-passingg and output head blocks shown in Figure 13 Figure 14 and Figure 15) respectively.
For all datasets, we train and sample with o 0.4 A. For the Langevin dynamics EEquation 68 we M 1, friction of 1.0 and step size of At G.
Fg
Cai
SEE3 Equiiarrant Genmetric
y/
Atom Embeddde
Output Head
Sege goc
Message Module
Cout(o)) ((r)
y
blocksi
Eoly. σ)
Cuoise (7)
Noise Level
Figure 12: Overview of the denoiser network dg. The submodule Fg sees input atom coordinates (a) and outputs predicted atom coordinates y', which gets scaled and added to noise- conditional skip connection to finally obtain fe
The hidden features (n) for O,.... 4 contain 120 scalar and 32 vector features per atom. We use spherical harmonics up 1 for the tensor product.
We use the Adam optimizer with learning rate .002. Models are trained with a batch size 32 over NVIDIA RTX A100 GPUs.
24

---

Figure 13: Overview of the initial embedder in the denoiser network, creating initial features h each atom and edge.
h (r))
h h
Nobe Conditioral Scaling
Nase Conditornl Skip
Spherccll Farmenia
Cate Activation
g
Noise Level
Figure 14: Overview single (B33))-quivarian message-pane passaagggg block (indexed by n)) in denoiser network. There are four such blocks iteratively updating the atom features from (t to h(4), The atom coordinates denoted by ğ Cn(o) (and hence, the edge features) are unchanged throughout these blocks.
Figure 15: Overview of the output head, which predicts the coordinates y Fo (Cin (oll(()))) (o))s
25

---

G
Normalization
As the noise level is increased, y + where (Nxss expands in space. Let represent the "normalized" input y, as seen by the network Fg:
ÿ = (o)y
(10)
To control the expansion of y, (o)(o) chosen such that the following property holds:
E (((.))-nifoom) (E) (llis 03 || = 1 at all noise levels a.
(11)
Note that here, we define E such that Note that this is distinct from the normalization chosen by 41 42 which normalizes |!yl directly. The intuition behind this normalization that the GNN model Fe does not operate on atom positions y directly, but instead uses the relative vectors y; y4 to account for translation invariance, and controlling this object directly ensures that the topology of the graph does not change with varying noise level Ø.
To achieve this, we compute:
Cm(l)
(12)
VC + 6
where C' = (i.) (Uni)--nil || Tz || can be easily estimated from the true data distribution. The
full derivation can be found in Section G.I.
As the input now appropriately normalized, the target output of the network Fg should also be appropriately normalized. A full derivation, found in [Section G.2 leads to:
Cskip(o) C + 60²
(13)
C.
Came (0) = C+ 6
(14)
Cnoise = log10 σ
(15)
The noise normalization is a scaled version of the recommendation of !
z In for images in Karras
al. 4 42]
G.1 Input Normalization
Fix an (i,j) E E from Equation 11 As Ex,E lid (N,0, Iz), we have E ~ (0, 211)) from the closure of the multivariate Gaussian under linear combinations. Thus, for each component
= 1, 2and3, we have: (e; - E))(a) ~ NN(02) and hence:
EeLVOLLNs - 2z5²² (s: {;)1 Σ (s. -(er-- =0
(16)
d=1
3
BeL - >
[l[(s - =EDiaa =
(17)
d=1
We can now compute:
Eg [llis - ع
(o)(a) 2 yı |
( 2 x, + a(e; e
cin (11:; =,,1² + 20lE (I=; (s; (s: - =5]] +o2Ballle - eg1l1)
cin(o)² (1::; - =,,1²² + o Eellle - Ey 111))
=cn( 2 [|:; - =r,1² 6o
(18)
26

---

Now, taking the expectation over all (i,j) E E uniformly:
E(0.j)-Dnffomm))llll أا1l ((x.j)- EE Illis "y : (mo((0Ix
cin( (((ts)))ttmttmm())))) E) || zz 66o²
(19)
Let C' = (i)) -Uifio-U E) ||,; 2yl1 which we estimate from the true data distribution. Then, from Equation 19ans and our intended normalization given by Equation II
Cin(a) VC 6
(20)
G.2 Output Normalization
The derivation here identical that of 41 42 but with our normalization. The denoising loss at single noise level is:
C(ro, σ) Ex-px lo.. ae,a) -
(21)
which gets weighted across a distribution Pa of noise levels by (unnormalized) weights X(o):
((#o) Bawp, ((o))
Bowp, Exxpx B ((a) [X + 0) x1
Eaoa, Ex Ex~px By [A(r) (Y, 1
Rowp, Ex-px By N
((a) (() (o)Y + Coma Fo(Cn(a)) Cnoise (n))se((a)
2r|1²1
Cskip (0)
{or)cou(a) Fe Fo(Cm(o)),, Guoise (a)) 20
Eowp, Bx
Cour (0)
Eonp, By
M(o) (o) ( (o(C Cooise(o)) -F713]
(22)
where:
Fly o) I- Cokkppp
(23)
Cout
is the effective training target for the network Fg. We want to normalize F similarly as the network input:
Uni--too E) [1] - F511 = at all noise levels a.
(24)
e~110,
Again, for a fixed (i,j) E, we have:
Bx [1x; 1.) - Cakip(() (o) y1)l
Eg F, Fj11
Goul (
Eg [1(1 Cskip (a)pp(o))(( - x,) - Cskip(o)er (E, E1
Cran Crar(a)) (0)2
c 2 ||, 1z1l + Cakip (0)
(25)
Coul (0)2
and hence:
[(t.j)) E) ( Cakip (a))2 C Cskip (0)2 .60 (o)
[1 F311
= Comr( com(d0) (0)2 = =(1 Cakip (o)) C Cakip (0)2. 2.)2.60 -602
(26)
27

---

where C' was defined above. Now, to minimize Cour to maximize reuse and avoid amplifying
network errors, as recommended by Karras et 41 42
d
deskip (o) Coul (o)= =0 0
-2(1 Ssip(() (0)) C 20skip (0) )- -60²
C
Clip(o) C + 66
(27)
Substituting into Equation 26 we get after some routine simplification:
C 6
Cour Cour( (0) = V C+ 6
(28)
The noise normalization is chosen as Cnoise (r) loogo o, a scaled version of the recommendation of In for images in Karras et al. 41 42
From Equation 22, we set A(a)
(o) to normalize the loss at at all noise levels, as in Karras et al.
G.3
Rotational Alignment
As described in Algorithm 1 we use the Kabsch-Umeyaa algorithm 40 x before calling the denoiser.
82 to rotationally align
Algorithm 1 Rotational Alignment with the Kabsch-Umeyand Algorithm
Require: Noisy Sample y E IR Nx3, True Sample x E Nx3
H x² U. U.S.VT t SVD(H) R Udiag[1, 1, de dV))))) return y(
H R3x3 D U,V E R3x3
Note that both and æ are mean-centered to respect translational equivariance:
N
Σ
y
0 R
(29)
1=1
N
Σ x,
(30)
so there is no net translation.
H Proofs of Theoretical Results
For completeness, we prove the main theoretical results here, as first established by Robbins [66
Miyasawa (58). Saremi and Hyvärinen [70].
H.1 The Denoiser Minimizes the Expected Loss
Here, we prove Equation 6
rewritten here for clarity:
a(-) =(-)=EE E =
arg min
E witl (lf((
x1|21
(31)
Y=x+os
28

---

First, we can decompose the loss over the domain R Nx3 of Y:
e0 - x11 = Ex-px, Popy [Ils(D) - x11
(32)
Y=X+os
JRNX JNN3 (1f(v) - xx1 y (z, v)dady
(33)
(1f(2) x01 prix (y x)px (x)dxdd
(34)
JRNXS
JRNX
1(fa))
I(f.y)dy
(35)
where v) 2 for all functions f and inputs y. Hence, any minimizer f* must minimize local denoising loss [([", ,y) at each point RR-x33, For a fixed R the loss I(f, y) is convex as a function fg). Hence, the global minimizer can found by finding the critical points of ((f. y) as function of fy):
son v)
(36)
V fiv)
lI(()) =|I² pr x z)px (x)dx
(37)
RNKa
2( - =)prxx( (x))
(x))da= =
(38)
3NNS
Rearranging:
SRNXA Prix (y (x)
(2(
(39)
JANKS PYJx x)px (x)de
SaNXA Py (y (x)x
(40)
(y
Py (y xx)
(41)
JRNX
(v)
(x 1vdr
(42)
RNN
E[X Y vl
(43)
i(y)
(44)
by Bayes' rule. Hence, the denoiser as defined by Equation indeed the minimzer of the denoising loss:
=x)=EXX
arg min f
E ew [II
-X]]
(45)
Y=X+os
as claimed.
H.2 Relating the Score and the Denoiser
Here, we rederive Equation 7 relating the score function log,y py and the denoiser 4.
Let X ~ px defined over Nx3 and (N, Let Y = X + oy,, which means:
|| 1
Prix = Ni INxa) (2 37 exp
(46)
20
Then:
X Y = y] = -y+o 1 loggy(()) Py
(47)
29

---

To prove this:
VXpri x) = y -ー x
o² pr(x(() xx)
(48)
= (x - v)rrix x)) = V (y xt)
(49)
(a - 1 pxx [ JRNXS o'v v px (x)de
px
(50)
JNN
By Bayes' rule:
p ())x (x,
(x)(())))) (x0
(51)
and, by definition of the marginals:
Jenis px,r (2, = =pr(u)
(52)
For the left-hand side, we have:
Jawcs (x - --))pri((v x)px (x)dr=
JANKS (x - y)pxxyyyyy (x, y)
/
(53)
[ JNNK xp y)dx S UpX, (z, v)dr
(54)
-y (v) (f xpxir(2 lv)dr JRNKS
feve px|r 1y)x
(55)
Py (v) ( =y
(56)
For the right-hand side, we have:
0² / Veprix(y =)px(x)dx= = / Prix(u x) (x)dr
(57)
RNNX
= =o'V,, INN px,
(58)
= =o'Vnpr(1
(59)
Thus,
(v) [E[X] Y v] y) 0² Q²V VVyPP pp)
(60)
E Y y) 0² VpPy (y)
(61)
py (g)
y V log (v)
(62)
as claimed.
I Numerical Solvers for Langevin Dynamics
As mentioned in Section 2.2] solving the Stochastic Differential Equation corresponding to Langevin dynamics is often performed numerically. In partieular, BAOAB [49 51 68 refers to a 'splitting
method that solves the Langevin dynamics SDE by splitting it into three different components labelled by A, B and below:
dy = vl
(63)
du,- = M- y log PY (v)dt yuydt + 2 dB,
(64)
B
30

---

where both by E Rd. This leads to the following update operators:
Asi
y
y + Un
At
(65)
0y
Byy
Ba
y/
y
(66)
Py
Byy + M
V lo (y)
y
y 00 by M 2 vI -
Osi
(67)
[yy
where B ~ (0, Td) is resampled every iteration. As highlighted by Kieninger and Keller 45] the A and updates are obtained by simply discretizing the updates highlighted in Equation 63 by the Euler method. The O update refers to explicit solution of the we rederive for completeness in Appendix J
process, which
Finally, the iterates of the BAOAB algorithm are given by composition of these update steps, matching the name of the method:
y(++1)
(uf ((+1)
= Agi D.IAA: Bor
(68)
J The Ornstein-Uhlenbeck Process
For completeness, we discuss the distributional solution of the -eattinnUhhlebeck process, taken directly from the excellent Leimkuhler and Matthews [5]]. In one dimension, the Ornsein-Uulenbeed Uhlenbeck Process corresponds to the following Stochastic Differential Equation (SDE):
dug = yy + 2 dB,
(69)
Multiplying both sides by the integrating factor eh:
du = -ue" + e² 22 bdB. (du,, + ru, ett -2yMM tdB,
(70)
(71)
and identifying:
(duy + ylyd() = vy) We get after integrating from t to 12, two adjacent time steps of our integration grid:
(72)
d(e th / dB.
(73)
Ied d(c = [( 1/ en 2 3dd
(74)
c 2 (t2) - e (t()) = 2 f" ed'dB.
(75)
Now, for a Wiener process B,, if g(t) is a deterministic function, J0tt g(t)dB, is distributed as N (o, Sar a(t)att by Itô's integral. Thus, applying this result to g(t) = ext, we get:
(t2)) - (t))))(t)) = V2TM- $o (0. 2212 227
(76)
27
= (g(22) = e 7/' - (12 (10- (11) + 2 - (0, e21z - 2yth 27tt
(77)
27
(22((1-t2)
=e-nlo-tu))((t))) + 2
N (0,
(78)
27
= -c-1(o--t) (ti) + M v1 627(11- 12 (0,1)
(79)
31

---

In the N x 3 dimensional case, as the Wiener processes are all independent of each other, we directly get:
vy(12) = =e x---(t (. M (00()-e2(ttaaa)) (0, (1,3))
(80)
Setting Al = 1/ 41, we get the form of the O operator (Equation 65) of the BAOAB integrator in Appendix
K Parallelizing Sampling with Multiple Independent Chains
Our sampling strategy batches peptides in order to increase throughput. Another potential method to increase throughput (but which we did not employ for the results in this paper) is to sample multiple chains in parallel.
This can be done by initializing multiple chains:
f/ 1.... where:
(0)
UNa
(0)
(0) Meh x (0)
+ az
(0)
(81)
~ (0, x3) for ch = l,.... db are all independent of each other. Then, the chains can be evolved independently withindependent walk steps (Equation 3) and denoised with independent jump steps (Equation 7). This independence allows batching over the Veh over all chains ch each iteration i.
where (0) iid
Note that at the number of samples per chain is large enough, the chains are no longer correlated, as they have now mixed into the stationary distribution.
the chains are correlated they are all initialized from the same (0()), However, if
L Model transferability exploration using Macrocycles
One of the most important aspects of JAMUN is its highly general input, a point cloud, unlike other protein ensemble models that are frequently more bespoke, using dihedral or frame representations. While this may be a slight disadvantage for us in the protein space, it also makes us extremely flexible and easily transferable other modalities. Here, we demonstrate that on macrocyclic peptides with several non-canonical residues.
The exploration of conformational ensembles in macrocyclic peptides is crucial due to their emerging role as therapeutic modalities. These molecules present significant challenges in computational modeling because of their conformational diversity and inherent geometric constraints. In fact, molecular dynamics trajectories for these molecules are particularly slow as good classical forcefields unavailable and it is necessary to quantum mechanical calculations to compute forces. Macrocyclic peptides are also extremely unwieldy in their "open", most common conformations, forming hydrogen bond networks with water. However, those macrocycles that are able to occupy smaller "crumpled" conformations are greasy and able to permeate through biological membranes, making them more suitable for biodelivery. Here, as an example, we use macrocycles from the CREMP dataset (28), generate ensembles with the CREST protocol [SS] and benchmark the resulting conformers against the RINGER model (7). Figure 16 illlsstrtee the transferability of JAMUN to macrocyclic peptides.
It is clear that we are able to recover most basins sampled, even though it does seem like there new basins uncovered. We note that our outputs look significantly more diffusive than the ground truth. The main reason for this that the CREST data is clustered and filtered to represent the local minima, whereas JAMUN is designed to sample entire distributions. In some sense, the data is strictly not complete for the task JAMUN designed for. It impressive that spite of this JAMUN learns enough from the denoising "jump" step with very local data to still perform the Langevin dynamics "walk" step and sample multiple basins.
We find that for 4-mers, which is what we trained our model for, we are able to recover all the sampled basins. We also attempt run inference for 5 and 6-mers with the same model test generalizability.
While this a preliminary study, it points to the potential of JAMUN being used universally, not just for proteins, and in particular shows that it is effective even in low-data regime.
32

---

(9)
Her.
18
812
-102
INEEO
-02
Figure 16: Macrocycle results trained 4AA: 3D image of the (a) CREMP (green) and (b) JAMUN (cyan) MeS. MeS. V.L macrocycle. (c) Ramachandran plot for CREMP. JAMUN, and RINGER samples of the 4AA.A McS.McS. V.L macrocycle. (d) Ramachandran plot for CREMP, JAMUN, and RINGER samples of the 5AA F.Q.L.G.M Met macrocyclc.(e) Ramachandran plot for CREMP and JAMUN the 6AA Mes. Mes.T.Q.Mei. V.W macrocycle.
33


Reviews:
[{"review_id": "GSdmGm5Rj9", "reviewer": "Reviewer_u7xj", "summary": "The authors present JAMUN, a walk-jump sampling model for generating ensembles of molecular conformations, outperforming the state-of-the-art TBG model, and competitive with the performance of MDGen with no protein-specific parametrization. This is an application paper which applys Walk-Jump sampling in MD data.", "questions": "1. How does the Walk-Jump process introduce physical priors? Both noise injection and Langevin dynamics are essentially probabilistic sampling techniques, which still differ significantly from real physical processes.\n2. Can JAMUN benefit from MD trajectory data, or is it solely designed for sampling from equilibrium states? \n3. If so, how does MD connects with Langevin Dynamics in JAMUN in physical sense?\n4. Why single noise-level is used? Following the score function will lead to decreasing noise level, how to make sure the score estimation is still correct in the decreasing noise level?", "limitations": "Yes, the authors have discussed the limitations.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strength\n\n1. The idea of connecting MD and Score-based Learning ( more specifically, Langevin Dynamics) is interesting\n2. The experimental results presented in the paper are comprehensive and demonstrate competitive performance.  \n\n### Weakness\n\n1. The overall presentation of the paper would benefit from improvement. Adding an algorithm flowchart or illustrative diagram summarizing the proposed approach would help readers better grasp the core methodology and overall workflow.\n2. The motivation for using Walk-Jump Sampling in this paper needs to be further clarified. It is recommended to supplement the introduction with a discussion of the existing limitations of current Diffusion and Flow Matching methods to better justify the proposed approach.", "was_revised": true, "final_justification": "The Walk and Jump process was initially proposed in dWJS (Discrete Walk-Jump Sampling in Protein Generative Modeling), and this paper borrows this process for molecular dynamics. The application is appropriate although with limited innovations in generative modeling. I think with very valuable MD process data, the authors could try to learn the true kinetic process to help obtain physics-favored energetics distributions. Meanwhile, many bridge processes have been adopted to simulate the MD process and obtain the energetics distributions. The motivation and advantages of WJS over diffusion bridges should be made more clear. For summary, I appreciate the detailed explanations from the authors during rebuttal and this paper is a nice application paper of WJS to MD. It is important to try out different generative models for different tasks and the authors did a good job.", "rebuttal": "# Common Response for All Reviewers\n\nWe would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our \"simple but powerful core idea\" and \"comprehensive experimentation\" demonstrating \"competitive performance\" in recovering Boltzmann distributions for unseen proteins.\n\nThe main concerns about the paper shared by the reviewers can be broadly placed in the following categories:\n- A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)\n- A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)\n- An ablation of the jump step (Reviewer E7vk)\n- A lack of results on larger systems (Reviewers iqm5 and tyev)\n\nWe have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\\approx$ 50 AA) and the shortest ATLAS proteins ($\\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.\n\n## Comparison with Diffusion Models\n\nTo highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):\n\n|Sampler|Samples|NFE|JSD-Backbone Torsions|\n|-------|-------|---|---------------------|\n|Walk-Jump|3149|6298|0.150198|\n|Diffusion|3149|399923|0.136336|\n|Walk-Jump|200000|400000|0.049639|\n|Diffusion|200000|25400000|0.046087|\n\nEssentially, we find that walk-jump sampling is $64\\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.\n\nWe would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.\n\n## Physical Validity and Energy Analysis of JAMUN Samples\n\nWe have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.\n\n|Posebusters Metric|Average Pass Rate|\n|------------------|-----------------|\n|bond\\_lengths|97.0%|\n|bond\\_angles|99.4%|\n|internal\\_steric\\_clash|100.0%|\n|internal\\_energy|97.5%|\n|overall|94.7%|\n\nWe see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the *overall quality of the JAMUN samples is high*.\nFor a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:\n\n|Posebusters Metric / Pass Rate|>90%|>92%|>94%|>96%|>98%|100%|\n|------------------------------|----|----|----|----|----|----|\n|bond_lengths|20/20|18/20|17/20|14/20|10/20|5/20|\n|bond_angles|20/20|20/20|20/20|20/20|18/20|14/20|\n|internal_steric_clash|20/20|20/20|20/20|20/20|20/20|20/20|\n|internal_energy|20/20|19/20|19/20|17/20|10/20|5/20|\n|overall|19/20|14/20|13/20|11/20|6/20|2/20|\n\nNext, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:\n\n|Sequence|MDGen 4AA-Explicit Energy (kJ/mol)|JAMUN Energy (kJ/mol)|Difference (kJ/mol)|\n|--------|----------------------------------|---------------------|-------------------|\n|FHSE|-675.6 ± 20.1|-633.2 ± 113.3|42.4|\n|FKKL|-699.0 ± 24.0|-562.5 ± 192.1|136.4|\n|FLRH|-1272.7 ± 18.9|-1198.5 ± 129.3|74.2|\n|FSDP|-697.5 ± 23.7|-687.6 ± 91.5|9.9|\n|FSRK|-1333.0 ± 21.9|-1349.9 ± 0.0|-16.9|\n|GCIC|-557.5 ± 21.4|-538.8 ± 56.4|18.6|\n|GGHN|-905.3 ± 21.9|-821.3 ± 129.6|84.0|\n|GLIL|-743.3 ± 20.7|-711.6 ± 71.8|31.7|\n|HELI|-794.2 ± 25.5|-780.6 ± 73.5|13.6|\n|HENV|-1156.9 ± 17.3|-1123.4 ± 148.7|33.5|\n|HTIQ|-762.9 ± 17.0|-726.8 ± 105.8|36.1|\n|IAMI|-428.0 ± 15.3|-426.8 ± 74.8|1.2|\n|IDRH|-1416.8 ± 18.1|-722.8 ± 2608.1|694.1|\n|IHNV|-845.4 ± 21.1|-864.4 ± 48.8|-19.1|\n|IMRY|-1230.7 ± 23.6|-1100.6 ± 207.0|130.1|\n|INVH|-793.6 ± 22.9|-745.3 ± 128.6|48.3|\n|IPGD|-611.5 ± 15.0|-582.5 ± 55.2|29.1|\n\n## Ablation of the Jump Step\nAs requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far *higher divergence metrics*. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.\n\n|Sequence|JSD-TICA of De-Noised X|JSD-TICA of Noisy Y|\n|--------|----------------------|-------------------|\n|FGGW|0.149416|0.322338|\n|FKKL|0.189984|0.490095|\n|FSDP|0.261191|0.390349|\n|HENV|0.232173|0.406107|\n|HTIQ|0.251693|0.384356|\n|IAMI|0.129338|0.399676|\n|IMRY|0.278553|0.362170|\n\n## Results on Larger Systems\nWe recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.\n\nWe have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.\n\nImportantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).\n\nFor the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware.\nFor the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating *transferability* even in this setting.\n\nWe compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.\n\n# Response to Reviewer u7xj\n\n1. and 3. The Walk-Jump process corresponds to performing Langevin dynamics in the smoothed $Y$ space, just like how classical MD performs Langevin dynamics in clean $X$ space. Langevin dynamics is an approximation for the friction and thermal fluctuations induced by solvent particles: https://en.wikipedia.org/wiki/Langevin_dynamics#Overview.\n2. We usually train on the entire MD trajectory data. Our results on cyclic peptides in Appendix J are on post-processed equilibrium states.\n4. Note that we perform Langevin dynamics to sample $Y$ at a fixed noise level. This is different from the standard diffusion SDE where the noise level is continuously reduced during sampling.\n\nWe would be happy to update Figure 1) for clarity to show the sampling scheme. We would also add our analysis that shows that diffusion models spend many sampling steps due to sampling from an uninformative Gaussian prior, whereas walk-jump sampling operates from a smoothed, partially noised distribution which avoids the need to sample many steps."}, {"review_id": "4r0DYQcK0t", "reviewer": "Reviewer_iqm5", "summary": "The paper presents JAMUN, a walk-jump sampler that couples smoothed Langevin dynamics on a latent space with an denoiser. On short-peptide benchmarks the method delivers $10^{1}$–$10^{2}$ accelerations while preserving ensemble fidelity, as verified by Jensen–Shannon distances and MSM state populations. \nThe core idea is simple but powerful, the empirical study is convincing for small systems.\nThough several challenges remain such as scaling to larger proteins, accelerated MD methods for ensemble sampling will likely become increasingly important in the near future.", "questions": "1. **Noise-level transferability.**\nThe optimal latent noise $\\sigma$ was tuned on four-residue peptides and fixed at 0.4 \\AA. \nHave you investigated whether the same $\\sigma$ remains effective when the system size increases by an order of magnitude, say, a 30-residue mini-protein, particularly when secondary structures such as $\\alpha$-helices or $\\beta$-hairpins emerge?\n\n2. **Scale up with collective variables**\nJAMUN is demonstrated in an all-atom coordinate space. Scaling to larger proteins often benefits from reduced representations, for example by propagating only selected collective variables while reconstructing atomistic detail on demand. Can the walk-jump framework be adapted to such collective-variable spaces, perhaps by learning the score in CV coordinates and performing the jump in a coarse-grained latent?", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strength\n\nThe main contributions are both conceptually neat useful. \nBy operating in a single-noise latent space, the authors avoid the heavy diffusion-model yet still retain a learned score that drives Langevin integration.\nExperiments demonstrate that JAMUN decorrelates conformations more than 10 $\\times$ faster than classical MD or recent ML baselines while maintaining high ensemble quality, as seen in low Jensen–Shannon distances to reference trajectories and accurate MSM-derived state populations.\n\n# Weakness\n\nThe molecular system is small; all quantitative results concern peptides of at most five residues. Because the core claim is “length-agnostic generalisation”, the absence of tests on mini-proteins (tens of residues) or multi-chain complexes leaves open whether the sampler remains stable, efficient, and accurate once the conformational landscape becomes vastly higher-dimensional. Until such scale-up validation is provided, the generality of JAMUN’s speed and quality advantages must be regarded as promising but unproven.", "was_revised": true, "final_justification": "The authors addressed the key points I raised. They conducted the additional experiments, especially application on larger systems, which help to support their claims of scalability. While the results are preliminary, it suggests that JAMUN can generalize beyond the peptide scale without retuning.", "rebuttal": "# Common Response for All Reviewers\n\nWe would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our \"simple but powerful core idea\" and \"comprehensive experimentation\" demonstrating \"competitive performance\" in recovering Boltzmann distributions for unseen proteins.\n\nThe main concerns about the paper shared by the reviewers can be broadly placed in the following categories:\n- A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)\n- A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)\n- An ablation of the jump step (Reviewer E7vk)\n- A lack of results on larger systems (Reviewers iqm5 and tyev)\n\nWe have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\\approx$ 50 AA) and the shortest ATLAS proteins ($\\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.\n\n## Comparison with Diffusion Models\n\nTo highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):\n\n| Sampler  | Samples | NFE | JSD-Backbone Torsions   |\n|-----------|---------|----------|----------|\n| Walk-Jump | 3149    | 6298     | 0.150198 |\n| Diffusion | 3149    | 399923   | 0.136336 |\n| Walk-Jump | 200000  | 400000   | 0.049639 |\n| Diffusion | 200000  | 25400000 | 0.046087 |\n\nEssentially, we find that walk-jump sampling is $64\\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.\n\nWe would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.\n\n## Physical Validity and Energy Analysis of JAMUN Samples\n\nWe have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.\n \n| Posebusters Metric | Average Pass Rate |\n|--------|-------------------|\n| bond\\_lengths | 97.0% |\n| bond\\_angles | 99.4% |\n| internal\\_steric\\_clash | 100.0% |\n| internal\\_energy | 97.5% |\n| overall | 94.7% |\n\nWe see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the *overall quality of the JAMUN samples is high*.\nFor a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:\n| Posebusters Metric / Pass Rate | >90% | >92% | >94% | >96% | >98% | 100% |\n|-----------------------|------:|------:|------:|------:|------:|------:|\n| bond_lengths | 20/20 | 18/20 | 17/20 | 14/20 | 10/20 | 5/20 |\n| bond_angles | 20/20 | 20/20 | 20/20 | 20/20 | 18/20 | 14/20 |\n| internal_steric_clash | 20/20 | 20/20 | 20/20 | 20/20 | 20/20 | 20/20 |\n| internal_energy | 20/20 | 19/20 | 19/20 | 17/20 | 10/20 | 5/20 |\n| overall | 19/20 | 14/20 | 13/20 | 11/20 | 6/20 | 2/20 |\n\nNext, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:\n\n| Sequence | MDGen 4AA-Explicit Energy (kJ/mol)     | JAMUN Energy (kJ/mol)     | Difference (kJ/mol) |\n|----------|--------------------|--------------------|---------------------|\n| FHSE     | -675.6 ± 20.1      | -633.2 ± 113.3     | 42.4               |\n| FKKL     | -699.0 ± 24.0      | -562.5 ± 192.1     | 136.4              |\n| FLRH     | -1272.7 ± 18.9     | -1198.5 ± 129.3    | 74.2               |\n| FSDP     | -697.5 ± 23.7      | -687.6 ± 91.5      | 9.9                |\n| FSRK     | -1333.0 ± 21.9     | -1349.9 ± 0.0      | -16.9              |\n| GCIC     | -557.5 ± 21.4      | -538.8 ± 56.4      | 18.6               |\n| GGHN     | -905.3 ± 21.9      | -821.3 ± 129.6     | 84.0               |\n| GLIL     | -743.3 ± 20.7      | -711.6 ± 71.8      | 31.7               |\n| HELI     | -794.2 ± 25.5      | -780.6 ± 73.5      | 13.6               |\n| HENV     | -1156.9 ± 17.3     | -1123.4 ± 148.7    | 33.5               |\n| HTIQ     | -762.9 ± 17.0      | -726.8 ± 105.8     | 36.1               |\n| IAMI     | -428.0 ± 15.3      | -426.8 ± 74.8      | 1.2                |\n| IDRH     | -1416.8 ± 18.1     | -722.8 ± 2608.1    | 694.1              |\n| IHNV     | -845.4 ± 21.1      | -864.4 ± 48.8      | -19.1              |\n| IMRY     | -1230.7 ± 23.6     | -1100.6 ± 207.0    | 130.1              |\n| INVH     | -793.6 ± 22.9      | -745.3 ± 128.6     | 48.3               |\n| IPGD     | -611.5 ± 15.0      | -582.5 ± 55.2      | 29.1               |\n\n## Ablation of the Jump Step\nAs requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far *higher divergence metrics*. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.\n\n| Sequence | JSD-TICA of De-Noised X | JSD-TICA of Noisy Y     |\n|----------|-----------|-----------|\n| FGGW     | 0.149416  | 0.322338  |\n| FKKL     | 0.189984  | 0.490095  |\n| FSDP     | 0.261191  | 0.390349  |\n| HENV     | 0.232173  | 0.406107  |\n| HTIQ     | 0.251693  | 0.384356  |\n| IAMI     | 0.129338  | 0.399676  |\n| IMRY     | 0.278553  | 0.362170  |\n\n## Results on Larger Systems\nWe recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.\n\nWe have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.\n\nImportantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).\n\nFor the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware.\nFor the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating *transferability* even in this setting.\n\nWe compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.\n\n# Response to Reviewer iqm5\n\n1. We believe our preliminary experiments on the fast-folding proteins and ATLAS where we keep all the training and sampling hyperparameters fixed as before, should answer your question. It does seem like these hyperparameters are robust.\n2. Yes, indeed the framework is general enough to support the sampling of collective variables (similar to how diffusion models can be learned for those data). Here, we wanted to utilize the connection between walk-jump sampling and MD which does operate in an all-atom space, to show that it might be possible to build transferable yet efficient samplers for molecular conformations."}, {"review_id": "1utjNpSKGt", "reviewer": "Reviewer_E7vk", "summary": "This paper addresses issues encountered when sampling the Boltzmann distribution using machine learning models, specifically limited sampling speed and transferability. Instead of performing Langevin dynamics directly in molecular conformation space, the authors propose performing Langevin dynamics in a noisy space, where atom coordinates are perturbed by Gaussian noise with a fixed variance. The authors demonstrate that only by knowing the score of the noisy data distribution, both Langevin Dynamics (“Walk”) and Projection to noise-free structures (“Jump”) can be performed.", "questions": "- Did the authors consider comparing in more detail to [1]? It seems like the approaches are closely related, since [1] also performs Langevin Dynamics using conformations with added noise, where they treat the noise level as a hyperparameter. However, [1] are missing the “Jump” step, projecting back to the data manifold. I would expect that the structures generated using JAMUN are of better fidelity and contain less noise in comparison. [2] show in Figure 4 that structures generated by [1] via Langevin Dynamics suffer from unphysical bond lengths due to the added noise. It would be great if the authors of JAMUN could show that their model doesn’t suffer from this problem.\n- The approach has been demonstrated to work for stationary observables only. It remains unclear if it would be possible to derive dynamic (time-dependent) observables from the generated trajectories as well?\n- Please consider adding captions to Table 1 for the different datasets\n- Jenson-Shannon distance could be briefly explained as one of the main metrics in the paper\n- Can you specify which settings in figure 2a correspond to which rows in table 1? E.g. how many samples did TBG (20x shorter) generate exactly?\n\n\n**References**\n\n[1] Arts, Marloes, et al. \"Two for one: Diffusion models and force fields for coarse-grained molecular dynamics.\" Journal of Chemical Theory and Computation 19.18 (2023): 6151-6159.\n\n[2] Plainer, M., Wu, H., Klein, L., Günnemann, S., & Noé, F. (2025). Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models. ArXiv. https://arxiv.org/abs/2506.17139", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Overall, I find this to be a nice paper where I like the idea of combining “Walks” via Langevin dynamics and projection via “Jumps” using the same score network. The performed experiments are clear and the paper is nicely written and easy to follow. In its current form, I recommend acceptance with a “borderline accept”. However, I am happy to raise my score if the authors add an ablation for the “Jump” step.\n\n**Strengths**\n\n- All-atom representation for point clouds -> not specific to proteins\n- Training the score network works only from samples of the Boltzmann distribution; physical trajectories or force labels are not needed\n- Only a single network must be trained for both “Walk” and “Jump” steps, where the model predicts the score of the noisy data distribution with fixed noise variance. The network is queried only twice to generate a new sample (once for “Walk” and once for the “Jump” step), which leads to high sampling speed.\n- The proposed method seems to be simple to implement, fast to train with modest amounts of parameters and effective in recovering Boltzmann distributions for unseen proteins.\n\n\n**Weaknesses**\n\n- From figure 2a it seems like JAMUN needs comparatively many sampling steps in comparison to other ML-based Boltzmann samplers, which in part diminishes the speed advantage of JAMUN.\n- Missing ablations: Especially ablating the “Jump” step (generating TICA plots from trajectories only using “Walk” and no “Jump”) would greatly enhance the quality of the paper to demonstrate that both “Walk” and “Jump” steps are indeed necessary. My intuition here is based on the following: Large structural dynamics analysed via TICA plot might be captured well even when local details like bond-distances or bond-angles are not accurately described.\n- The paper doesn’t analyze the physical validity of the generated structures; therefore, it doesn’t become clear if a single “Jump” step from a low noise level can already generate highly accurate structures compared to many consecutive sampling steps in standard generative models. An analysis about fine detail fidelity would allow better judging this (see “Questions”).", "was_revised": true, "final_justification": "During rebuttal, the authors successfully addressed the concerns raised in my original review. I am thus happy to increase my score and recommend acceptance.", "rebuttal": "# Common Response for All Reviewers\n\nWe would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our \"simple but powerful core idea\" and \"comprehensive experimentation\" demonstrating \"competitive performance\" in recovering Boltzmann distributions for unseen proteins.\n\nThe main concerns about the paper shared by the reviewers can be broadly placed in the following categories:\n- A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)\n- A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)\n- An ablation of the jump step (Reviewer E7vk)\n- A lack of results on larger systems (Reviewers iqm5 and tyev)\n\nWe have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\\approx$ 50 AA) and the shortest ATLAS proteins ($\\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.\n\n## Comparison with Diffusion Models\n\nTo highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):\n\n| Sampler  | Samples | NFE | JSD-Backbone Torsions   |\n|-----------|---------|----------|----------|\n| Walk-Jump | 3149    | 6298     | 0.150198 |\n| Diffusion | 3149    | 399923   | 0.136336 |\n| Walk-Jump | 200000  | 400000   | 0.049639 |\n| Diffusion | 200000  | 25400000 | 0.046087 |\n\nEssentially, we find that walk-jump sampling is $64\\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.\n\nWe would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.\n\n## Physical Validity and Energy Analysis of JAMUN Samples\n\nWe have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.\n \n| Posebusters Metric | Average Pass Rate |\n|--------|-------------------|\n| bond\\_lengths | 97.0% |\n| bond\\_angles | 99.4% |\n| internal\\_steric\\_clash | 100.0% |\n| internal\\_energy | 97.5% |\n| overall | 94.7% |\n\nWe see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the *overall quality of the JAMUN samples is high*.\nFor a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:\n| Posebusters Metric / Pass Rate | >90% | >92% | >94% | >96% | >98% | 100% |\n|-----------------------|------:|------:|------:|------:|------:|------:|\n| bond_lengths | 20/20 | 18/20 | 17/20 | 14/20 | 10/20 | 5/20 |\n| bond_angles | 20/20 | 20/20 | 20/20 | 20/20 | 18/20 | 14/20 |\n| internal_steric_clash | 20/20 | 20/20 | 20/20 | 20/20 | 20/20 | 20/20 |\n| internal_energy | 20/20 | 19/20 | 19/20 | 17/20 | 10/20 | 5/20 |\n| overall | 19/20 | 14/20 | 13/20 | 11/20 | 6/20 | 2/20 |\n\nNext, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:\n\n| Sequence | MDGen 4AA-Explicit Energy (kJ/mol)     | JAMUN Energy (kJ/mol)     | Difference (kJ/mol) |\n|----------|--------------------|--------------------|---------------------|\n| FHSE     | -675.6 ± 20.1      | -633.2 ± 113.3     | 42.4               |\n| FKKL     | -699.0 ± 24.0      | -562.5 ± 192.1     | 136.4              |\n| FLRH     | -1272.7 ± 18.9     | -1198.5 ± 129.3    | 74.2               |\n| FSDP     | -697.5 ± 23.7      | -687.6 ± 91.5      | 9.9                |\n| FSRK     | -1333.0 ± 21.9     | -1349.9 ± 0.0      | -16.9              |\n| GCIC     | -557.5 ± 21.4      | -538.8 ± 56.4      | 18.6               |\n| GGHN     | -905.3 ± 21.9      | -821.3 ± 129.6     | 84.0               |\n| GLIL     | -743.3 ± 20.7      | -711.6 ± 71.8      | 31.7               |\n| HELI     | -794.2 ± 25.5      | -780.6 ± 73.5      | 13.6               |\n| HENV     | -1156.9 ± 17.3     | -1123.4 ± 148.7    | 33.5               |\n| HTIQ     | -762.9 ± 17.0      | -726.8 ± 105.8     | 36.1               |\n| IAMI     | -428.0 ± 15.3      | -426.8 ± 74.8      | 1.2                |\n| IDRH     | -1416.8 ± 18.1     | -722.8 ± 2608.1    | 694.1              |\n| IHNV     | -845.4 ± 21.1      | -864.4 ± 48.8      | -19.1              |\n| IMRY     | -1230.7 ± 23.6     | -1100.6 ± 207.0    | 130.1              |\n| INVH     | -793.6 ± 22.9      | -745.3 ± 128.6     | 48.3               |\n| IPGD     | -611.5 ± 15.0      | -582.5 ± 55.2      | 29.1               |\n\n## Ablation of the Jump Step\nAs requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far *higher divergence metrics*. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.\n\n| Sequence | JSD-TICA of De-Noised X | JSD-TICA of Noisy Y     |\n|----------|-----------|-----------|\n| FGGW     | 0.149416  | 0.322338  |\n| FKKL     | 0.189984  | 0.490095  |\n| FSDP     | 0.261191  | 0.390349  |\n| HENV     | 0.232173  | 0.406107  |\n| HTIQ     | 0.251693  | 0.384356  |\n| IAMI     | 0.129338  | 0.399676  |\n| IMRY     | 0.278553  | 0.362170  |\n\n## Results on Larger Systems\nWe recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.\n\nWe have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.\n\nImportantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).\n\nFor the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware.\nFor the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating *transferability* even in this setting.\n\nWe compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.\n\n# Response to Reviewer E7vk\n\n1. Our detailed Posebusters analysis above shows that JAMUN captures bond lengths correctly; we hope you will find this sufficient, but we are happy to perform further analyses.\n2. In the current framework, no, we indeed lose kinetic information and cannot map the dynamics in $Y$ to those in $X$. \n3. and 4. We will fix the caption of Table 1 and also detail the JSD metric for clarity in the updated paper.\n5. According to Figure 2 a), TBG generates 5000 structures, so we used 5000/20 = 250 structures to represent TBG (20x shorter). JAMUN generated 100000 structures."}, {"review_id": "OGDD6uPwGL", "reviewer": "Reviewer_tyev", "summary": "This paper introduces JAMUN, a generative machine learning framework to accelerate molecular dynamics. JAMUN performs Langevin dynamics on a smoothed, noised space by leveraging walk-jump sampling. The model is fast and demonstrates transferability to unseen systems.", "questions": "-\tThe authors claim: This physical prior enables JAMUN to transfer well – just like force fields for molecular dynamics can – to unseen systems. I don’t think this claim is justified. For example, standard diffusion model also run Langevin dynamics over a “latent” space. What is exactly meant here and what in JAMUN should make it generalize better? Is it just that only partially corrupting samples is better than fully corrupting them to a standard gaussian? If so, why is that a physical prior?\n-\tThe authors claim: generation with JAMUN yields converged sampling of the conformational ensemble faster than MD with a standard force field, even outperforming several state-of-the-art baselines. I think this claim is misleading since the authors do not evaluate convergence to the Boltzmann distribution. They only evaluate convergence of torsional projections. In order to measure real speed-up, post-processing of these samples to energetically meaningful conformations should be accounted for in the total runtime.\n-\tCould the authors share more details about how the MSMs in Figure 1 were made? How do authors explain a substantial number of points with very high probabilities in Fig 1 b)?\n-\tHow do the authors explain that MDGen performs better in terms of Jensen-Shannon differences but worse in terms of mode coverage in sample test molecules?\n-\tHave the authors made sure that the MDGen model is sampled with long-enough lag times to be comparable with models of the Boltzmann distribution? Please add sampling details and mark initial structure in the Ramachandran/TICA plots.\n-\tDo the authors expect this strategy to scale seamlessly to larger systems? Results for larger systems would significantly strengthen the contribution.\n\nSuggestions:\n\n-\tAdd parameter count for all baselines models.\n-\tAdd energy evaluations of generated structures.\n\nMinor corrections:\n\n-\tLine 94: space after coma.\n-\tLine 167: implicit solvent.", "limitations": "-\tIt’s unclear to me how to disentangle the influence of the Walk-Jump Sampling to the one of the model. The advantage of using Walk-Jump Sampling would become clearer if the same architecture you use (with minimal modifications for noise conditioning) was used to train a baseline regular diffusion model.\n-\tThe manuscript would benefit from a thorough linguistic review to enhance clarity and readability.\n-\tIn Figure 1 b) and c) correlation coefficients of other baselines should be compared against.\n-\tPlease mark best (and second) performing models in bold in Table 1. \n-\tAuthors do not show examples for when JAMUN “hallucinates” metastable states. I suggest showing structure, computing energy and comparing with reference simulation.\n-\tI understand the lack of baselines for the extrapolation exercise, but I am not sure how relevant the comparison with Boltz-1 and BioEmu in Figure 5 are since the scope of these models is very different. Can the authors provide a better baseline?", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n-\tEmprical results suggest better generalization in terms of mode coverage with respect to baselines.\n-\tFast all-atom sampling and only need to learn one unconditional score model.\n\nWeakness:\n\n-\tComparable sampling time and worst performance in terms of Jenson-Shannon distances wrt to most modern (not all atom) baseline (MDGen).\n-\tUnclear how to disentangle the influence of the Walk-Jump Sampling because no comparison with respect to standard diffusion models is provided.", "was_revised": true, "final_justification": "JAMUN is a cheaper and performant method, however, improvement seems to be limited and entails some performance loss and I beleive some evaluations could be improved. I recommend acceptance.", "rebuttal": "# Common Response for All Reviewers\n\nWe would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our \"simple but powerful core idea\" and \"comprehensive experimentation\" demonstrating \"competitive performance\" in recovering Boltzmann distributions for unseen proteins.\n\nThe main concerns about the paper shared by the reviewers can be broadly placed in the following categories:\n- A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)\n- A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)\n- An ablation of the jump step (Reviewer E7vk)\n- A lack of results on larger systems (Reviewers iqm5 and tyev)\n\nWe have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\\approx$ 50 AA) and the shortest ATLAS proteins ($\\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.\n\n## Comparison with Diffusion Models\n\nTo highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):\n\n|Sampler|Samples|NFE|JSD-Backbone Torsions|\n|-------|-------|---|---------------------|\n|Walk-Jump|3149|6298|0.150198|\n|Diffusion|3149|399923|0.136336|\n|Walk-Jump|200000|400000|0.049639|\n|Diffusion|200000|25400000|0.046087|\n\nEssentially, we find that walk-jump sampling is $64\\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.\n\nWe would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.\n\n## Physical Validity and Energy Analysis of JAMUN Samples\n\nWe have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.\n\n|Posebusters Metric|Average Pass Rate|\n|------------------|-----------------|\n|bond\\_lengths|97.0%|\n|bond\\_angles|99.4%|\n|internal\\_steric\\_clash|100.0%|\n|internal\\_energy|97.5%|\n|overall|94.7%|\n\nWe see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the *overall quality of the JAMUN samples is high*.\nFor a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:\n\n|Posebusters Metric / Pass Rate|>90%|>92%|>94%|>96%|>98%|100%|\n|------------------------------|----|----|----|----|----|----|\n|bond_lengths|20/20|18/20|17/20|14/20|10/20|5/20|\n|bond_angles|20/20|20/20|20/20|20/20|18/20|14/20|\n|internal_steric_clash|20/20|20/20|20/20|20/20|20/20|20/20|\n|internal_energy|20/20|19/20|19/20|17/20|10/20|5/20|\n|overall|19/20|14/20|13/20|11/20|6/20|2/20|\n\nNext, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:\n\n|Sequence|MDGen 4AA-Explicit Energy (kJ/mol)|JAMUN Energy (kJ/mol)|Difference (kJ/mol)|\n|--------|----------------------------------|---------------------|-------------------|\n|FHSE|-675.6 ± 20.1|-633.2 ± 113.3|42.4|\n|FKKL|-699.0 ± 24.0|-562.5 ± 192.1|136.4|\n|FLRH|-1272.7 ± 18.9|-1198.5 ± 129.3|74.2|\n|FSDP|-697.5 ± 23.7|-687.6 ± 91.5|9.9|\n|FSRK|-1333.0 ± 21.9|-1349.9 ± 0.0|-16.9|\n|GCIC|-557.5 ± 21.4|-538.8 ± 56.4|18.6|\n|GGHN|-905.3 ± 21.9|-821.3 ± 129.6|84.0|\n|GLIL|-743.3 ± 20.7|-711.6 ± 71.8|31.7|\n|HELI|-794.2 ± 25.5|-780.6 ± 73.5|13.6|\n|HENV|-1156.9 ± 17.3|-1123.4 ± 148.7|33.5|\n|HTIQ|-762.9 ± 17.0|-726.8 ± 105.8|36.1|\n|IAMI|-428.0 ± 15.3|-426.8 ± 74.8|1.2|\n|IDRH|-1416.8 ± 18.1|-722.8 ± 2608.1|694.1|\n|IHNV|-845.4 ± 21.1|-864.4 ± 48.8|-19.1|\n|IMRY|-1230.7 ± 23.6|-1100.6 ± 207.0|130.1|\n|INVH|-793.6 ± 22.9|-745.3 ± 128.6|48.3|\n|IPGD|-611.5 ± 15.0|-582.5 ± 55.2|29.1|\n\n## Ablation of the Jump Step\nAs requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far *higher divergence metrics*. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.\n\n|Sequence|JSD-TICA of De-Noised X|JSD-TICA of Noisy Y|\n|--------|----------------------|-------------------|\n|FGGW|0.149416|0.322338|\n|FKKL|0.189984|0.490095|\n|FSDP|0.261191|0.390349|\n|HENV|0.232173|0.406107|\n|HTIQ|0.251693|0.384356|\n|IAMI|0.129338|0.399676|\n|IMRY|0.278553|0.362170|\n\n## Results on Larger Systems\nWe recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.\n\nWe have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.\n\nImportantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).\n\nFor the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware.\nFor the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating *transferability* even in this setting.\n\nWe compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.\n\n# Response to Reviewer tyev\n\n1. The idea is that walk-jump sampling shares many parallels with MD, but you are correct that we do not currently identify which specific aspects help the transferability of JAMUN. We would be happy to rephrase this in the paper. In a standard diffusion model, the noise level is continuously reducing during sampling which makes it hard to give a physical interpretation.\n2. Our analysis of the force field energy of JAMUN samples should address these concerns; we do not do any post-processing of the JAMUN samples. It is true that there is some bias in the final conformational distributions, which we cannot fully correct for, but our analysis does show that we are able to get quite close.\n3. The MSMs were made with the analysis code from MDGen (Jing, et al.) which we have modified and uploaded as part of our supplementary code. We utilize the same MSM hyperparameters (eg. TICA lag), but it is likely that those hyperparameters are not optimal for the Timewarp datasets.\n4. We ran the exact commands (and checkpoints) used by the MDGen authors for sampling, according to their README. It is possible that the recommended sampling parameters/model is not ideal for this comparison.\n5. We hope that you will find the results on the fast-folding proteins and ATLAS encouraging.\n\nWe would be very happy to improve the formatting and the clarity of the paper in the camera-ready version, if accepted."}]

Author rebuttal:

Meta-review: (4,5,5,4) This paper introduces JAMUN, a generative framework that accelerates molecular dynamics by performing Langevin dynamics in a smoothed, noised space using walk-jump sampling. Reviewers found the method simple, efficient, and transferable to unseen systems, with experiments showing substantial speedups over MD and competitive performance with ML baselines. While some concerns about baselines and validation remain, the overall assessment was positive, and I recommend acceptance.
Decision: Accept (poster)