Title: JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensemble Generation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Conformational Ensembles

Molecules

Protein Dynamics

Drug Discovery

Statistical Mechanics

Generative Models
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 Dynamics and Score-Based Learning for 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 .

1 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 is 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], or the beta-sheet transition to a disordered strand for insulin to dissociate and find and bind to its receptor [5]. Similarly, drug discovery on protein kinases depends on 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 to have similar impact. ML models for generating molecular ensembles are widely considered the 'next frontier' [11, 57, 91]. In this work, we present JAMUN (WalkJ ump A ccelerated M olecular ensembles with U niversal N oise), a generative ML model which advances this frontier by demonstrating improvements in both speed and transferability 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

a)

b)

c)

d)

Noise

+ ge(0)

x(0)

（0）

Smoothened

Energy Landscape

Walk

pr(y)

Walk

g(t)

Noise

Walk

(once)

Jump

Langevin Dynamics

with Vy log pr (g(c)

(t+1)

Molecular Dynamics

Energy Landscape

px(c)

2(0)

Jump

(t)

(t+1)

+ o²Vy log py(y(t)

c(t)

dynamics (MD), which will sample the Boltzmann distribution in the limit of infinite sampling time, but is 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 is '. . . equivalent to tracking the 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 is nothing fundamental about this small time-step limitation; it is an artifact of high-frequency motions, such as bond vibrations, that have little effect on protein ensembles [51]. Enhanced sampling methods have been developed in 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 4. 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, we focus on small peptides whose MD trajectories can be run to convergence in a reasonable amount of time, unlike those for much larger proteins [78].

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 - to unseen systems. The key idea is to run Langevin molecular dynamics in a noisy 'latent' space Y ∈ R N × 3 , instead of the original space X ∈ R N × 3 of all-atom 3D positions. Indeed, Y is constructed by adding a small amount of i.i.d Gaussian noise ε to X :

To run this Langevin MD over Y , the crucial component is the score function ∇ y log p Y ( y ) , which needs to be modelled. (This is identical to how a force field needs to be parametrized in classical MD). Once the MD trajectory over Y is run, the resulting samples need to be mapped back to clean data via a denoising procedure.

This framework is actually mathematically described by Walk-Jump Sampling (WJS), as first introduced by Saremi and Hyv¨ arinen [70]. WJS has been used in voxelized molecule generation [62, 63]

and protein sequence generation [24]. In particular, JAMUN corresponds to a SE (3) -equivariant Walk-Jump sampler of point clouds.

The WJS framework actually tells us that the score function ∇ y log p Y ( y ) can be used for denoising as well, removing the need to learn a 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 a 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 ∇ y log p Y ( y ) is learned by adding noise to clean data x , and a 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 to 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 as is commonly done in diffusion [32, 79] and flow-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 a high noise level.

We train JAMUN on a large dataset of MD simulations of small peptides. We demonstrate that this model can generalize to a holdout set of unseen peptides. In all of these cases, 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. 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 as Point Clouds

Each point cloud of N atoms can be represented by the tuple ( x, h ) where x ∈ R N × 3 represents the 3D coordinates of each of the N atoms and h ∈ R D represents atom and covalent bonding information. h can be easily computed from the amino acid sequence for each peptide, and hence is not learned or sampled. For clarity of presentation, we omit the conditioning on h in the distributions and models below. We discuss how our model uses h in Section 2.4.

At sampling time, we assume access to an initial sample x (0) ∈ R N × 3 sampled from the clean data distribution p X . 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 x (0) . In theory, x (0) could also be obtained from experimental data, such as crystallized structures from the Protein Data Bank [9]. We plan to explore this approach in the future to seed JAMUN simulations for larger proteins.

2.2 Walk-Jump Sampling

JAMUN operates by performing walk-jump sampling on molecular systems represented as 3D point clouds. A conceptual overview of the process is illustrated in Figure 1.

Given the initial sample x (0) ∼ p X , walk-jump sampling performs the following steps:

Noise the initial structure x (0) to create the initial sample y (0) from the noisy data distribution p Y (Figure 1a):

Walk to obtain samples y (1) , . . . , y ( N ) from p Y using Langevin dynamics which conists of numerically solving the following Stochastic Differential Equation (SDE) (Figure 1b):

where v y represents the particle velocity, ∇ y log p Y ( y ) is the gradient of the log of the probability density function (called the score function) of p Y , γ is friction, M is the mass, and B t is the standard Wiener process in N × 3 -dimensions: B t ∼ N (0 , t I N × 3 ) . In practice, we employ the BAOAB solver (Appendix I) to integrate Equation 3 numerically.

Jump back to p X to obtain samples ˆ x 1 , . . . , ˆ x N (Figure 1c) :

where ˆ x ( · ) ≡ E [ X | Y = · ] is called the denoiser. It corresponds to the minimizer (Section H.1) of the ℓ 2 -loss between clean samples X and samples denoised back from Y = X + σε .

where f : R N × 3 → R N × 3 . As shown by Robbins [66], Miyasawa [58] (and Section H.2), the denoiser ˆ x is closely linked to the score ∇ y log p Y :

Importantly, the score function ∇ y log p Y shows up in both the walk and jump steps, and we need to approximate this quantity.

2.3 Learning to Denoise

In order to run Walk-Jump Sampling as outlined above, we have the choice of modelling either the score ∇ y log p Y or the denoiser ˆ x as they are equivalent by Equation 7. Following trends in diffusion models [41, 42], we model the denoiser as a neural network ˆ x θ ( y, σ ) ≈ ˆ x ( y ) parameterized by model parameters θ .

Importantly, we only need to learn a model at a single, fixed noise level σ . This is unlike training diffusion or flow-matching models where a wide range of noise levels are required for sampling. In particular, the choice of noise level σ for WJS is important because mode-mixing becomes faster as σ is increased, but the task asked of the denoiser becomes harder.

The denoiser ˆ x θ thus takes in noisy point clouds y formed by adding noise (at a fixed noise level σ ) to clean point clouds x . The denoiser is tasked to reconstruct back x , given y . To be precise, training the denoiser ˆ x θ consists of solving the following optimization problem:

to obtain θ ∗ , the optimal model parameters. As is standard in the empirical risk minimization (ERM) [84] setting, we approximate the expectation in Equation 8 by sampling X ∼ p X and ε ∼ N (0 , I N × 3 ) . We minimize the loss as a function of model parameters θ using the first-order optimizer Adam [44] in PyTorch 2.0 [4, 22].

2.4 Parametrization of the Denoiser Network

We summarize the key features of the denoiser network ˆ x θ ( y, σ ) which will approximate ˆ x ( y ) in this section. Note that σ is fixed in our setting, but we explicitly mention it 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. [41, 42] (in the context of image generation):

where F θ represents a learned network parameterized by parameters θ . In particular, F θ (Figure 14) is a geometric graph neural network (GNN) model similar to NequIP [8, 81]. Importantly, F θ is chosen to be SE (3) -equivariant, in contrast to existing methods [33, 45, 46] that utilize the E (3) -equivariant EGNN model [71]. As rightly pointed out by Dumitrescu et al. [20] and Schreiner et al. [73], E (3) -equivariant models are equivariant under parity, which means that are forced to transform

CYS

CYS

ACE

ALA

ALA

NME

mirrored structures identically. When we experimented with E (3) -equivariant architectures, we found symmetric Ramachandran plots which arise from the unnecessary parity constraint of the denoising network. For this reason, TBG [45] and Timewarp [46] use a 'chirality checker' to post-hoc fix the generated structures from their model. For JAMUN, such post-processing is unnecessary because our model can distinguish between chiral structures.

The coefficients c skip ( σ ) , c out ( σ ) , c in ( σ ) , c noise ( σ ) in Equation 9 are normalization functions (from R + to R ) which adjust the effective inputs and outputs to F θ . They are chosen to encourage re-use of the input y at low noise levels, but the opposite at high noise levels. Importantly, based on the insight that F θ uses relative vectors in the message passing steps, we adjust the values of these coefficients instead of simply using the choices made in Karras et al. [41, 42], Wohlwend et al. [89], Abramson et al. [1], as discussed in Appendix G.

In F θ , edges between atoms are computed using a radial cutoff of 10 ˚ A 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 are 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 are obtained using the embedding of the residue code (eg. ALA, CYS) and concatenated to each atom in the residue. Importantly, we do not use the sequence index of the residues (eg. 0 , 1 , . . . ) as we found that it hurts generalization to longer peptide lengths.

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-LARGE and TIMEWARP 4AA-LARGE from Klein et al. [46], MDGEN 4AA EXPLICIT from Jing et al. [37], and our own MD data simulated with OpenMM [21]. A summary of these datasets is presented in Table 6. The differing simulation conditions across these datasets allows us to test the broad applicability of our approach. We describe these, as well as an internal dataset we use for more realistic setting analysis in C.

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.

4 Related Work and Baseline Models

The goal of building machine learning models that can generate conformational ensembles of molecular systems is not new. While a full overview of this field is beyond the scope of this work - see [6] 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 [90] learns a diffusion model over conformations of small organic molecules from MD17 [14], showing some level of transferability across C7O2H10 isomers [76]. Timewarp [46] uses a normalizing flow as a proposal distribution in MCMCsampling of the Boltzmann distribution to approximate the conditional distribution of future

conformational states x ( t +∆ t ) conditional on the present state x ( t ) . ITO [73] modelled this distribution using diffusion with a SE (3) -equivariant PaiNN architecture. Equijump [19] extended this idea with a protein-specific 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 flow-matching instead of maximum likelihood estimation and a more efficient continuous normalizing flow architecture.

We discuss a few more related methods and studies in E.

A related problem is that of protein structure prediction, as successfully tackled by the AlphaFold models [77, 38, 1]. Here, we compare to Boltz-1 [89], an open-source reproduction of AlphaFold3 [1]. Boltz-1 was trained exclusively on static crystalline structures of folded states, without any dynamics or conformational information, allowing us to 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 7 torsion angles for each residue. MDGen [37] cleverly builds on this parametrization by creating a SE (3) -invariant tokenization of the backbone torsion angles, relative to a known initial conformation (here, x (0) ). Then, they learn a stochastic interpolant [3] (a generalization of diffusion and flow-matching) 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 SE (3) -invariant tokenization is limited to single-chain proteins and peptides, but allows for more efficient architectures. On the other hand, JAMUN models all atoms explicitly with a SE (3) -equivariant network, which makes it far more flexible and easily extendable to molecules other than proteins, as discussed in the supplement.

Several models build upon AlphaFold2 to sample conformational ensembles, discussed further in E. One such method is BioEmu- technically only a backbone-only 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, 75]. As is standard practice, all TICA projections and MSMs are estimated using the reference MD data.

6 Experiments

Here, we compare to 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 10 × faster in practice. This allows for a fair comparison to JAMUN.

Due to the different simulation conditions across the datasets shown in Table 6, we train a different JAMUN model for each dataset. However, the same noise level of σ = 0 . 4 ˚ A 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 2 NVIDIA RTX A100 GPUs with 40 GB memory,

although competitive results can be obtained by training for only 1 day. Our denoiser model is built with the e3nn library [25], and contains approximately 10 . 5M parameters.

This scale of noise ( σ = 0 . 4 ˚ A) is 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, as shown in Figure 8a, we find that higher noise levels (such as σ = 0 . 8 ˚ A) result in samples with broken topologies, while lower noise levels (such as σ = 0 . 2 ˚ A) require many more sampling steps to explore the entire conformational landscape.

6.1 Results on TIMEWARP 2AA-LARGE and TIMEWARP 4AA-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-LARGE datasets.

Table 1 shows that on TIMEWARP 2AA-LARGE, JAMUN outperforms TBG when run for equal amounts of time (based on Figure 3), and is only slightly worse when TBG is run for 20 × 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.

a)

b)

6.2 Results on MDGEN 4AA-EXPLICIT

Table 1 shows that JAMUN is very competitive with MDGen on the JSD metrics. In fact, the TICA projections in Figure 6 shows that MDGen is missing some basins that JAMUN is able to sample, while the Ramachandran plots show an example where JAMUN hallucinates a 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 a rate 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 it to operate on molecules of larger sizes than it was originally trained on. Thus, we test whether JAMUN can generalize to peptides of lengths

-

π

-

π

/2

0

π

/2

π

ψ

Reference

-

π

-

π

/2

0

π

/2

π

ψ

Reference

(10x shorter)

-

π

-

π

/2

0

π

/2

π

ψ

JAMUN

-

π

-

π

/20

π

/2

π

ϕ

-

π

-

π

/2

0

π

/2

π

ψ

-

π

-

π

/20

π

/2

π

ϕ

-

π

-

π

/20

π

/2

π

ϕ

LIRH

10 -3

10 -2

10 -1

10 0

10 1

10 2

10 3

Speedup Factor

0

50

100

150

Frequency

10 -3

10 -2

10 -1

10 0

10 1

10 2

10 3

Speedup Factor

0

25

50

75

100

125

Frequency

beyond its training set. This is a challenging task, and one that we 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, we choose Boltz-1 and BioEmu which support sampling on UNCAPPED 5AA to compare against JAMUN trained on TIMEWARP 4AA-LARGE.

Surprisingly, we find that the JAMUN model trained only on 4AA peptides can accurately predict ensembles for 5AA peptides . Figure 8b and Figure 11 show that JAMUN is able to recover most states and even reproduce relative probabilities. Interestingly, the same experiment does not work if we train on TIMEWARP 2AA-LARGE instead, suggesting that the 2AA reference MD data may not be informative enough to generalize from.

On the other hand, we find that Boltz-1 is unable to sample the diversity of peptide conformations. This is not entirely surprising as 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-wise distance matrix,

and thus encodes a significant portion of the geometry. Keeping this representation fixed possibly prevents the sampling of large conformational changes.

Surprisingly, BioEmu also seems to 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 1 shows that JAMUN significantly outperforms Boltz-1 and BioEmu on the metrics from Section 5. As seen in Figure 3, JAMUN is roughly 5 × faster than Boltz-1, and is roughly 60 × faster than BioEmu when we perform side-chain reconstruction with H-Packer.

7 Conclusion

We 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 represents an important step toward the ultimate goal of a transferable generative model for protein conformational ensembles. Performing MDin the noised space gives the model a clear physics interpretation, and allows faster decorrelation (and hence, sampling) than classical MD.

The model has some limitations that motivate future work. While it is highly transferable in the space of two to five amino acid peptides, scaling up is likely to require more exploration and intensive data generation in future work. Additionally, while the current SE (3) -equivariant denoiser architecture works well, further development of the denoising network could speed up sampling. Alternative jump methods, such as multiple denoising steps (` a la diffusion), could also serve to sharpen generation. Lastly, a promising direction that has not yet been explored is the application of classical enhanced sampling methods, such as metadynamics, for traversing the noisy space.

References

J. Abramson, J. Adler, J. Dunger, R. Evans, T. Green, A. Pritzel, O. Ronneberger, L. Willmore, A. 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, E. Arvaniti, C. Beattie, O. Bertolli, A. Bridgland, A. Cherepanov, M. Congreve, A. I. Cowen-Rivers, A. Cowie, M. Figurnov, F. B. Fuchs, H. Gladman, R. Jain, 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: 10.1038/s41586-024-07487-w. URL https://doi.org/10.1038/s41586-024-07487-w .

V. Agarwal and A. C. McShan. The power and pitfalls of alphafold2 for structure prediction beyond rigid globular proteins. Nature Chemical Biology , 20(8):950-959, Aug 2024. ISSN 1552-4469. doi: 10.1038/s41589-024-01638-w. URL https://doi.org/10.1038/ s41589-024-01638-w .

M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions, 2023. URL https://arxiv.org/abs/2303.08797 .

J. 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. DeVito, 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, S. 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 '24) . ACM, Apr. 2024. doi: 10.1145/3620665. 3640366. URL https://pytorch.org/assets/pytorch2-2.pdf .

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):5571-5587, 2020. doi: 10.1021/acs.jpcb.0c03521. URL https://doi.org/10.1021/acs.jpcb.0c03521 . PMID: 32515958.

A. Aranganathan, X. Gu, D. Wang, B. Vani, and P. Tiwary. Modeling Boltzmann weighted structural ensembles of proteins using AI based methods. 2024.

M. Arts, V. Garcia Satorras, C.-W. Huang, D. Zugner, M. Federici, C. Clementi, F. No´ e, R. Pinsler, and R. van den Berg. Two for one: Diffusion models and force fields for coarsegrained molecular dynamics. Journal of Chemical Theory and Computation , 19(18):6151-6159, 2023.

S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications , 13(1):2453, May 2022. ISSN 2041-1723. doi: 10.1038/s41467-022-29939-5. URL https://doi.org/10.1038/s41467-022-29939-5 .

H. M. Berman, J. Westbrook, Z. Feng, G. Gilliland, T. N. Bhat, H. Weissig, I. N. Shindyalov, and P. E. Bourne. The Protein Data Bank. Nucleic Acids Research , 28(1):235-242, 01 2000. ISSN 0305-1048. doi: 10.1093/nar/28.1.235. URL https://doi.org/10.1093/nar/28.1.235 .

D. W. Borhani and D. E. Shaw. The future of molecular dynamics simulations in drug discovery. Journal of computer-aided molecular design , 26:15-26, 2012.

G. R. Bowman. AlphaFold and Protein Folding: Not Dead Yet! The Frontier Is Conformational Ensembles. Annual Review of Biomedical Data Science , 7, 2024.

M. Buttenschoen, G. M. Morris, and C. M. Deane. Posebusters: Ai-based docking methods fail to generate physically valid poses or generalise to novel sequences. Chemical Science , 15(9): 3130-3139, 2024. ISSN 2041-6539. doi: 10.1039/d3sc04185a. URL http://dx.doi.org/ 10.1039/D3SC04185A .

D. Chakravarty and L. L. Porter. AlphaFold2 fails to predict protein fold switching. Protein Sci , 31(6):e4353, June 2022.

S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. Sch¨ utt, and K.-R. M¨ uller. Machine learning of accurate energy-conserving molecular force fields. Science Advances , 3 (5):e1603015, 2017. doi: 10.1126/sciadv.1603015. URL https://www.science.org/doi/ abs/10.1126/sciadv.1603015 .

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.

T. Darden, D. York, and L. Pedersen. Particle Mesh Ewald: An N log(N) method for Ewald sums in large systems. The Journal of Chemical Physics , 98(12):10089-10092, June 1993. doi: 10.1063/1.464397. URL https://doi.org/10.1063/1.464397 .

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.7554/eLife.75751. URL https://doi.org/10.7554/eLife.75751 .

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-938, Nov. 2018.

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. URL https://arxiv.org/abs/2410.09667 .

A. Dumitrescu, D. Korpela, M. Heinonen, Y. Verma, V. Iakovlev, V. Garg, and H. L¨ ahdesm¨ aki. Field-based Molecule Generation, 2024. URL https://arxiv.org/abs/2402.15864 .

P. Eastman, J. Swails, J. D. Chodera, R. T. McGibbon, Y. Zhao, K. A. Beauchamp, L.-P. Wang, A. C. Simmonett, M. P. Harrigan, C. D. Stern, R. P. Wiewiora, B. R. Brooks, and V. S. Pande. OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. PLOS Computational Biology , 13(7):e1005659, July 2017. doi: 10.1371/journal.pcbi.1005659. URL https://doi.org/10.1371/journal.pcbi.1005659 .

W. Falcon and The PyTorch Lightning team. PyTorch Lightning, Mar. 2019. URL https: //github.com/Lightning-AI/lightning .

M. L. Fern´ andez-Quintero, 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 therapeutics and vaccines. Antibodies , 12(4):67, 2023.

N. C. Frey, D. Berenberg, J. Kleinhenz, I. Hotzel, J. Lafrance-Vanasse, R. L. Kelly, Y . Wu, A. Rajpal, 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.

M. Geiger and T. Smidt. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.09453 , 2022.

N. R. Gough and C. G. Kalodimos. Exploring the conformational landscape of protein kinases. Current Opinion in Structural Biology , 88:102890, 2024.

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://arxiv.org/abs/2310.01234 .

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.1038/s41597-024-03698-y. URL https://doi.org/10.1038/ s41597-024-03698-y .

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(35):e2410662121, 2024. doi: 10.1073/pnas.2410662121. URL https://www.pnas.org/doi/abs/10.1073/pnas.2410662121 .

K. Henzler-Wildman and D. Kern. Dynamic personalities of proteins. Nature , 450(7172): 964-972, Dec. 2007.

B. Hess, H. Bekker, H. J. Berendsen, and J. G. Fraaije. Lincs: a linear constraint solver for molecular simulations. Journal of computational chemistry , 18(12):1463-1472, 1997.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models, 2020. URL https: //arxiv.org/abs/2006.11239 .

E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant Diffusion for Molecule Generation in 3D, 2022. URL https://arxiv.org/abs/2203.17003 .

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 , 20(6):2335-2348, 2024.

X. Huang, R. Pearce, and Y. Zhang. FASPR: an open-source tool for fast and accurate protein side-chain packing. Bioinformatics , 36(12):3758-3765, June 2020.

B. Jing, B. Berger, and T. Jaakkola. AlphaFold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845 , 2024.

B. Jing, H. St¨ ark, T. Jaakkola, and B. Berger. Generative modeling of molecular dynamics trajectories. arXiv preprint arXiv:2409.17808 , 2024.

J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. ˇ Z´ ıdek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, A. J. Ballard, A. Cowie, B. Romera-Paredes, S. 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 , 596(7873):583-589, Aug 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-03819-2. URL https://doi.org/10.1038/s41586-021-03819-2 .

A. Jussupow and V. R. I. Kaila. Effective molecular dynamics from neural network-based structure prediction models. Journal of Chemical Theory and Computation , 19(7):19651975, 2023. doi: 10.1021/acs.jctc.2c01027. URL https://doi.org/10.1021/acs.jctc. 2c01027 . PMID: 36961997.

W. Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A , 32(5):922-923, Sep 1976. doi: 10.1107/S0567739476001873. URL https://doi. org/10.1107/S0567739476001873 .

T. Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the Design Space of Diffusion-Based Generative Models. In Proc. NeurIPS , 2022.

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.

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.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2017. URL https: //arxiv.org/abs/1412.6980 .

L. Klein and F. No´ e. Transferable Boltzmann Generators. arXiv preprint arXiv:2406.14426 , 2024.

L. Klein, A. Foong, T. Fjelde, B. Mlodozeniec, M. Brockschmidt, S. Nowozin, F. No´ e, and R. Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning timecoarsened dynamics. Advances in Neural Information Processing Systems , 36, 2024.

L. Klein, A. Kr¨ amer, and F. No´ e. Equivariant flow matching. Advances in Neural Information Processing Systems , 36, 2024.

M. Lazou, O. Khan, T. Nguyen, D. Padhorny, D. Kozakov, D. Joseph-McCarthy, and S. Vajda. Predicting multiple conformations of ligand binding sites in proteins suggests that alphafold2 may remember too much. Proceedings of the National Academy of Sciences , 121(48): e2412719121, 2024. doi: 10.1073/pnas.2412719121. URL https://www.pnas.org/doi/ abs/10.1073/pnas.2412719121 .

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.1093/amrx/abs010. URL https://doi.org/10.1093/amrx/abs010 . eprint: https://academic.oup.com/amrx/article-pdf/2013/1/34/397230/abs010.pdf.

B. Leimkuhler and C. Matthews. Robust and efficient configurational molecular sampling via langevin dynamics. The Journal of Chemical Physics , 138(17):174102, May 2013. doi: 10.1063/1.4802990. URL https://doi.org/10.1063/1.4802990 .

B. Leimkuhler and C. Matthews. Molecular Dynamics: With Deterministic and Stochastic Numerical Methods . Number 39 in Interdisciplinary Applied Mathematics. Springer International Publishing : Imprint: Springer, Cham, 1st ed. 2015 edition, 2015. ISBN 978-3-319-16375-8.

S. Lewis, T. Hempel, J. Jim´ enez-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, J. Nigam, F. Barbero, V. Stimper, A. Campbell, J. Yim, M. Lienen, Y. Shi, S. Zheng, H. Schulz, U. Munir, C. Clementi, and F. No´ e. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv , 2024. doi: 10.1101/2024.12.05.626885.

Z. Lin, H. Akin, R. Rao, B. Hie, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. dos Santos Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, and A. Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science , 379(6637):11231130, 2023. doi: 10.1126/science.ade2574. URL https://www.science.org/doi/abs/10. 1126/science.ade2574 .

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://openreview.net/forum?id=PqvMRDCJT9t .

J. Lu, B. Zhong, Z. Zhang, and J. Tang. Str2str: A score-based framework for zero-shot protein conformation sampling. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=C4BikKsgmK .

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 B , 111(27):7812-7824, Jul 2007. ISSN 1520-6106. doi: 10.1021/jp071097f. URL https://doi.org/10.1021/jp071097f .

M. D. Miller and G. N. Phillips. Moving beyond static snapshots: Protein dynamics and the Protein Data Bank. Journal of Biological Chemistry , 296, 2021.

K. Miyasawa. An Empirial Bayes Estimator of the Mean of a Normal Population. Bulletin de l'Institut international de statistique. , 38(4):181-188, 1960.

L. Molgedey and H. G. Schuster. Separation of a mixture of independent signals using time delayed correlations. Phys. Rev. Lett. , 72:3634-3637, Jun 1994. doi: 10.1103/PhysRevLett.72. 3634. URL https://link.aps.org/doi/10.1103/PhysRevLett.72.3634 .

F. No´ e, S. Olsson, J. K¨ ohler, and H. Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science , 365(6457):eaaw1147, 2019.

G. P´ erez-Hern´ andez, F. Paul, T. Giorgino, G. De Fabritiis, and F. No´ e. Identification of slow molecular order parameters for markov model construction. J Chem Phys , 139(1):015102, July 2013.

P. O. Pinheiro, A. Jamasb, O. Mahmood, V. Sresht, and S. Saremi. Structure-based Drug Design by Denoising Voxel Grids. arXiv preprint arXiv:2405.03961 , 2024.

P. O. Pinheiro, J. Rackers, J. Kleinhenz, M. Maser, O. Mahmood, A. Watkins, S. Ra, V. Sresht, and S. Saremi. 3D molecule generation by denoising voxel grids. Advances in Neural Information Processing Systems , 36, 2024.

J. W. Ponder and D. A. Case. Force fields for protein simulations. 66:27-85, 2003. ISSN 0065-3233. doi: https://doi.org/10.1016/S0065-3233(03)66002-X. URL https: //www.sciencedirect.com/science/article/pii/S006532330366002X .

P. Pracht and S. Grimme. Automated exploration of the low-energy chemical space with fast quantum chemical methods. Physical Chemistry Chemical Physics , 22(14):71697192, 2020. doi: 10.1039/C9CP06869D. URL https://pubs.rsc.org/en/content/ articlelanding/2020/cp/c9cp06869d .

H. Robbins. An Empirical Bayes Approach to Statistics. In Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, Volume 1: Contributions to the Theory of Statistics , volume 3.1, 1956.

J. P. Roney and S. Ovchinnikov. State-of-the-art estimation of protein model accuracy using alphafold. Phys. Rev. Lett. , 129:238101, Nov 2022. doi: 10.1103/PhysRevLett.129.238101. URL https://link.aps.org/doi/10.1103/PhysRevLett.129.238101 .

M. Sachs, B. Leimkuhler, and V. Danos. Langevin dynamics with variable coefficients and nonconservative forces: from stationary states to numerical methods. Entropy , 19(12):647, 2017.

T. Salda˜ no, N. Escobedo, J. Marchetti, D. J. Zea, J. Mac Donagh, A. J. Velez Rueda, E. Gonik, A. Garc´ ıa Melani, J. Novomisky Nechcoff, M. N. Salas, T. Peters, N. Demitroff, S. Fernandez Alberti, N. Palopoli, M. S. Fornasari, and G. Parisi. Impact of protein conformational diversity on alphafold predictions. Bioinformatics , 38(10):2742-2748, 04 2022. ISSN 1367-4803. doi: 10.1093/bioinformatics/btac202. URL https://doi.org/10.1093/bioinformatics/ btac202 .

S. Saremi and A. Hyv¨ arinen. Neural Empirical Bayes. Journal of Machine Learning Research , 20(181):1-23, 2019.

V. G. Satorras, E. Hoogeboom, and M. Welling. E(n) Equivariant Graph Neural Networks, 2022. URL https://arxiv.org/abs/2102.09844 .

M. K. Scherer, B. Trendelkamp-Schroer, F. Paul, G. P´ erez-Hern´ andez, M. Hoffmann, N. Plattner, C. Wehmeyer, J.-H. Prinz, and F. No´ e. Pyemma 2: A software package for estimation, validation, and analysis of markov models. Journal of Chemical Theory and Computation , 11(11):55255542, 2015. doi: 10.1021/acs.jctc.5b00743. URL https://doi.org/10.1021/acs.jctc. 5b00743 . PMID: 26574340.

M. Schreiner, O. Winther, and S. Olsson. Implicit transfer operator learning: Multiple timeresolution surrogates for molecular dynamics, 2023. URL https://arxiv.org/abs/2305. 18046 .

M. Schreiner, O. Winther, and S. Olsson. Implicit transfer operator learning: multiple timeresolution surrogates for molecular dynamics. arXiv preprint arXiv:2305.18046 , 2023.

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 , 9(4):2000-2009, Apr. 2013.

K. T. Sch¨ utt, F. Arbabzadah, S. Chmiela, K. R. M¨ uller, and A. Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Communications , 8(1), Jan. 2017. ISSN 20411723. doi: 10.1038/ncomms13890. URL http://dx.doi.org/10.1038/ncomms13890 .

A. W. Senior, R. Evans, J. Jumper, J. Kirkpatrick, L. Sifre, T. Green, C. Qin, A. ˇ Z´ ıdek, A. W. R. Nelson, A. Bridgland, H. Penedones, S. Petersen, K. Simonyan, S. Crossan, P. Kohli, D. T. Jones, D. Silver, K. Kavukcuoglu, and D. Hassabis. Improved protein structure prediction using potentials from deep learning. Nature , 577(7792):706-710, Jan 2020. ISSN 1476-4687. doi: 10.1038/s41586-019-1923-7. URL https://doi.org/10.1038/s41586-019-1923-7 .

D. E. Shaw, P. Maragakis, K. Lindorff-Larsen, 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 , 330(6002):341-346, 2010. doi: 10.1126/science. 1187409. URL https://www.science.org/doi/abs/10.1126/science.1187409 .

J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models, 2022. URL https: //arxiv.org/abs/2010.02502 .

M. Steinegger and J. S¨ oding. 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/nbt.3988. URL https://doi.org/10.1038/nbt.3988 .

N. Thomas, T. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

S. Umeyama. Least-squares estimation of transformation parameters between two point patterns. IEEE Transactions on Pattern Analysis and Machine Intelligence , 13(4):376-380, 1991. doi: 10.1109/34.88573.

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.

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. MorganKaufmann, 1991. URL https://proceedings.neurips.cc/paper_files/paper/1991/ file/ff4d5fbbafdf976cfdc032e3bde78de5-Paper.pdf .

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. ˇ Z´ ıdek, 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 , 52(D1):D368-D375, 11 2023. ISSN 0305-1048. doi: 10.1093/nar/gkad1011. URL https://doi.org/10.1093/ nar/gkad1011 .

G. M. Visani, W. Galvin, M. Pun, and A. Nourmohammad. H-packer: Holographic rotationally equivariant convolutional neural network for protein side-chain packing. In D. A. Knowles and S. Mostafavi, editors, Proceedings of the 18th Machine Learning in Computational Biology meeting , volume 240 of Proceedings of Machine Learning Research , pages 230-249. PMLR, 30 Nov-01 Dec 2024. URL https://proceedings.mlr.press/v240/visani24a.html .

A. Vitalis and R. V. Pappu. Methods for monte carlo simulations of biomacromolecules. Annual reports in computational chemistry , 5:49-76, 2009.

H. K. Wayment-Steele, A. Ojoawo, R. Otten, J. M. Apitz, W. Pitsawong, M. H¨ omberger, S. Ovchinnikov, L. Colwell, and D. Kern. Predicting multiple conformations via sequence clustering and alphafold2. Nature , 625(7996):832-839, Jan 2024. ISSN 1476-4687. doi: 10.1038/s41586-023-06832-9. URL https://doi.org/10.1038/s41586-023-06832-9 .

J. 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.11.19.624167. URL https://www.biorxiv. org/content/early/2024/11/20/2024.11.19.624167 .

F. Wu and S. Z. Li. Diffmd: A geometric diffusion model for molecular dynamics simulations, 2023. URL https://arxiv.org/abs/2204.08672 .

L.-E. Zheng, S. Barethiya, E. Nordquist, and J. Chen. Machine learning generation of dynamic protein conformational ensembles. Molecules , 28(10):4047, 2023.

S. Zheng, J. He, C. Liu, Y. Shi, Z. Lu, W. Feng, F. Ju, J. Wang, J. Zhu, Y . Min, et al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence , pages 1-10, 2024.

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 are that JAMUN is able to 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 1 and Table 1 show 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 1 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 to 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 of 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 SE(3)-equivariant 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 a complete (and correct) proof?

Answer: [Yes] The main contribution of our work is not theoretical. We largely build upon the theoretical framework of Walk-Jump Sampling as developed by [70]. For completeness, we re-derive all key theoretical results in Appendix H and re-derive the Langevin dynamics update in Appendix J.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of 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 F. 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes] The code was submitted in accordance with the ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) which includes scripts to reproduce the results in the paper. Section 3 includes the description of the data used within the paper. The TIMEWARP 2AA-LARGE, TIMEWARP 4AA-LARGE, and MDGEN 4AA-EXPLICIT datasets are available to the public under an MIT licence. We are in the legal process of sharing the CAPPED 2AA dataset for public use.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes] The data splits are described in Section 3, and were chosen to allow a comparison to existing baseline models. The parametrization and the hyperparameters of our model are described in Section 2.4 and Appendix F. We also provide code to reproduce our results.

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: [Yes] Our metrics are averaged over the test dataset, where we report the mean and 1-sigma errors for all JSD metrics.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes] We provide the compute details for training JAMUN in Section 6 and sampling JAMUN and all baseline models in Figure 3.

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 to the niche field of study, we do not believe that there are 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 or models that have a 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

Question: 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?

Answer: [Yes] Yes, we credit all authors for creating the benchmark datasets: Klein et al. [46] for TIMEWARP 2AA-LARGE and TIMEWARP 4AA-LARGE, 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 3 and Appendix C. We are in the legal process of releasing this dataset for public use.

14. Crowdsourcing 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 as 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 research 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?

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 an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or 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 in the core research methods.

A Comparison of Walk-Jump Sampling Against Full Diffusion

Section 6 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 a single noise level) to full diffusion (which samples a 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 a noise schedule of 64 steps from 0 . 01 ˚ A to 10 ˚ A using the Heun second order method as recommended by EDM [41]) and walk-jump sampling. We 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):

MD ref

MD ref

diffusion

diffusion

walkjump

walkjump

JS Divergence

JS Divergence

10-1

10-1

10-2

10-2

102

103

104

105

102

103

104

105

106

107

N samples

NFE

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 it works in a 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 4AAEXPLICIT 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 a 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 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:

C Data Details

The TIMEWARP and MDGEN datasets consist of 'uncapped' peptides, whose termini are zwitterionic amino and carboxyl groups, as shown in the left panel of Figure 2. These are not ideal analogues of amino acids in proteins due to local charge interactions as 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, a 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 of conformations. We choose the same splits as in TIMEWARP 2AA-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 against biased molecular dynamics runs using Non-Equilibrium Umbrella Sampling (NEUS) [18, 83], a trajectory stratification based enhanced sampling algorithm. The protein is represented by the amber99sbildn force field [64]. The simulations are performed at 300 K with the BAOAB integrator [50] in OpenMM [21]; LINCS is used to constrain the lengths of bonds to hydrogen atoms [31]; Particle Mesh Ewald is used to calculate electrostatics [16]; the step size was 2 fs. The systems are 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.1 Further Results on internally simulated data

Here, we compute the ratio of the decorrelation times for the backbone and sidechain torsions in JAMUN and the reference MD data. Figure 10 highlights how sampling in the smoothed space Y compared to the original space X enables much faster decorrelation.

In Table 7, we compare JAMUN to the reference MD, shortened by a factor of 10 along the variables mentioned in Section 5, using the Jensen-Shannon distance to the full reference MD data. JAMUN

outperforms this shortened MD trajectory across all metrics, even though it takes approximately 2 × longer to sample than JAMUN, from Figure 3.

E Related Work

Protein structure prediction only requires prediction of a few folded states, and is usually trained on 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 to 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) [9]. Another issue is that the quality of these models' predictions can be dependent on multiple sequence alignment (MSA) information, a preprocessing step where sequence databases are queried for similar protein sequences which indicate evolutionary conservation patterns as a supplementary input to the model. In fact, some of the first attempts to sample alternative conformations of proteins with AlphaFold2 were performed by subsampling [17], clustering [88] or manipulating MSA information [29]. Sequence-to-structure

10 -3

10 -2

10 -1

10 0

10 1

10 2

10 3

Speedup Factor

0

20

40

60

Frequency

10 -3

10 -2

10 -1

10 0

10 1

10 2

10 3

Speedup Factor

0

20

40

60

80

Frequency

models such as ESMFold [53] which do not require MSA information can be 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 of 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 AlpahFold2 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 to parametrize an additional harmonic potential for coarse-grained MD with the MARTINI [56] force field. AlphaFlow [36] develops flow-matching over the quotient space of 3D positions modulo rotations to learn a distribution over 3D positions of the β -carbon atoms, outperforming MSA subsampling with AlphaFold2 over the Protein Data Bank (PDB) [9]. Distributional Graphormer [92] parametrizes a diffusion model over α -carbon positions, also trained on the PDB. Str2Str [55] proposes a noising-denoising process along a 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 a 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 4, '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 Y as JAMUN does, which again limits the effective timestep of simulation.

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 a Boltzmann Emulator.

JAMUN is a walk-jump sampling method which uses an SE (3) -equivariant 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 is the first to our knowledge to apply walk-jump sampling to point clouds.

F Overview of Denoiser

The denoiser is a SE (3) -equivariant graph neural network. The graph is defined by a radial cutoff of 10 ˚ A in y . The overall computation performed by the denoiser is shown in Figure 12, with the initial embedding, message-passing and output head blocks shown in Figure 13, Figure 14 and Figure 15 respectively.

For all datasets, we train and sample with σ = 0 . 4 ˚ A. For the Langevin dynamics (Equation 68), we set M = 1 , friction of γ = 1 . 0 and a step size of ∆ t = σ .

F

Initial

Atom

SE(3)-Equivariant

Geometric

ndano

Cskip

Embedder

Message-Passing

Head

Module

Cin(o)

Cout(o)

(y,0)

y

(4blocks)

Cnoise(0)

Noise Level

The hidden features h ( n ) for n = 0 , . . . , 4 contain 120 scalar and 32 vector features per atom. We use spherical harmonics up to l = 1 for the tensor product.

We use the Adam optimizer with learning rate . 002 . Models are trained with a batch size of 32 over 2 NVIDIA RTX A100 GPUs.

0

Compute

Compute

Compute

Edges with

Compute

Edge-Level

Features

Atom-Level

Features

Residue-Level

Features

Atom Type

Residue Type

Atom Distances

Atom PDB Code

Residue Index

y

h(n+1)

Noise

SE(3)-Equivariant

SE(3)-Equivariant

Noise

Conditional

Tensor Product

Linear Layer and

Conditional

Scaling

with Spherical

Gate Activation

Skip

Harmonics

Connection

y

Cnoise(α)

Cnoise(o)

Noise Level

h(4)

SE(3)-Equivariant

Linear Layer

Coordinates by

1 vector per atom

Vectors

y

y

G Normalization

As the noise level σ is increased, y = x + σε where ε ∼ N (0 , I N × 3 ) expands in space. Let ˜ y represent the 'normalized' input y , as seen by the network F θ :

To control the expansion of y , c in ( σ ) is chosen such that the following property holds:

Note that here, we define E such that Note that this is distinct from the normalization chosen by [41, 42], which normalizes ∥ y ∥ directly. The intuition behind this normalization is that the GNN model F θ does not operate on atom positions y directly, but instead uses the relative vectors y i -y j 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:

where C = E ( i,j ) ∼ Uniform ( E ) ∥ x i -x j ∥ 2 can be easily estimated from the true data distribution. The full derivation can be found in Section G.1.

As the input is now appropriately normalized, the target output of the network F θ should also be appropriately normalized. A full derivation, found in Section G.2, leads to:

The noise normalization is a scaled version of the recommendation of 1 4 ln σ for images in Karras et al. [41, 42].

G.1 Input Normalization

Fix an ( i, j ) ∈ E from Equation 11. As ε i , ε j iid ∼ N (0 , I 3 ) , we have ε i -ε j ∼ N (0 , 2 I 3 ) from the closure of the multivariate Gaussian under linear combinations. Thus, for each component d = 1 , 2 and 3 , we have: ( ε i -ε j ) ( d ) ∼ N (0 , 2) and hence:

We can now compute:

Now, taking the expectation over all ( i, j ) ∈ E uniformly:

Let C = E ( i,j ) ∼ Uniform ( E ) ∥ x i -x j ∥ 2 , which we estimate from the true data distribution. Then, from Equation 19 and our intended normalization given by Equation 11:

G.2 Output Normalization

The derivation here is identical that of [41, 42], but with our normalization. The denoising loss at a single noise level is:

which gets weighted across a distribution p σ of noise levels by (unnormalized) weights λ ( σ ) :

where:

is the effective training target for the network F θ . We want to normalize F similarly as the network input:

Again, for a fixed ( i, j ) ∈ E , we have:

and hence:

where C was defined above. Now, to minimize c out ( σ ) to maximize reuse and avoid amplifying network errors, as recommended by Karras et al. [41, 42]:

Substituting into Equation 26, we get after some routine simplification:

The noise normalization is chosen as c noise ( σ ) = log 10 σ , a scaled version of the recommendation of 1 4 ln σ for images in Karras et al. [41, 42].

From Equation 22, we set λ ( σ ) = 1 c out ( σ ) 2 to normalize the loss at at all noise levels, as in Karras et al. [41, 42].

G.3 Rotational Alignment

As described in Algorithm 1, we use the Kabsch-Umeyama algorithm [40, 82] to rotationally align y to x before calling the denoiser.

Algorithm 1 Rotational Alignment with the Kabsch-Umeyama Algorithm

Note that both y and x are mean-centered to respect translational equivariance:

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¨ arinen [70].

H.1 The Denoiser Minimizes the Expected Loss

Here, we prove Equation 6, rewritten here for clarity:

First, we can decompose the loss over the domain R N × 3 of Y :

where l ( f, y ) ≥ 0 for all functions f and inputs y . Hence, any minimizer f ∗ must minimize the local denoising loss l ( f ∗ , y ) at each point y ∈ R N × 3 . For a fixed y ∈ R N × 3 , the loss l ( f, y ) is convex as a function of f ( y ) . Hence, the global minimizer can be found by finding the critical points of l ( f, y ) as a function of f ( y ) :

by Bayes' rule. Hence, the denoiser as defined by Equation 5 is indeed the minimzer of the denoising loss:

as claimed.

H.2 Relating the Score and the Denoiser

Here, we rederive Equation 7, relating the score function ∇ log p Y and the denoiser ˆ x . Let X ∼ p X defined over R N × 3 and η ∼ N (0 , I N × 3 ) . Let Y = X + ση , which means:

Then:

Rearranging:

To prove this:

By Bayes' rule:

and, by definition of the marginals:

For the left-hand side, we have:

For the right-hand side, we have:

Thus, 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 particular, 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 O below:

where both y, v y ∈ R d . This leads to the following update operators:

where B ∼ N (0 , I d ) is resampled every iteration. As highlighted by Kieninger and Keller [43], the A and B updates are obtained by simply discretizing the updates highlighted in Equation 63 by the Euler method. The O update refers to a explicit solution of the Ornstein-Uhlenbeck process, which we rederive for completeness in Appendix J.

Finally, the iterates of the BAOAB algorithm are given by a composition of these update steps, matching the name of the method:

J The Ornstein-Uhlenbeck Process

For completeness, we discuss the distributional solution of the Ornstein-Uhlenbeck process, taken directly from the excellent Leimkuhler and Matthews [51]. In one dimension, the Ornstein-Uhlenbeck Process corresponds to the following Stochastic Differential Equation (SDE):

Multiplying both sides by the integrating factor e γt :

and identifying:

We get after integrating from t 1 to t 2 , two adjacent time steps of our integration grid:

Now, for a Wiener process B t , if g ( t ) is a deterministic function, ∫ t 2 t 1 g ( t ) dB t is distributed as N ( 0 , ∫ t 2 t 1 g ( t ) 2 dt ) by Itˆ o's integral. Thus, applying this result to g ( t ) = e γt , we get:

In the N × 3 dimensional case, as the Wiener processes are all independent of each other, we directly get:

Setting ∆ t = t 2 -t 1 , we get the form of the O operator (Equation 65) of the BAOAB integrator in Appendix I.

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: y (0) 1 , . . . , y (0) N ch , where:

where ε (0) ch iid ∼ N (0 , I N × 3 ) for ch = 1 , . . . , N ch 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 y ( t ) ch over all chains ch at each iteration t .

Note that at t = 0 , the chains are correlated as they are all initialized from the same x (0) . However, if the number of samples per chain is large enough, the chains are no longer correlated, as they have now mixed into the stationary distribution.

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 to 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 are unavailable and it is necessary to use 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 CREMPdataset [28], generate ensembles with the CREST protocol [65], and benchmark the resulting conformers against the RINGER model [27]. Figure 16 illustrates 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 are new basins uncovered. We note that our outputs look significantly more diffusive than the ground truth. The main reason for this is 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 is designed for. It is impressive that in 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 to run inference for 5 and 6-mers with the same model to test generalizability.

While this is 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 a low-data regime.

(a)

(b)

MeS.MeS.V.L

(c)

n/2

CREMP

Z/u

JAMUN

Z/u

RINGER

(d)

TT

π/2

CREMP

T/2

π/2

JAMUN

0

π/2

(e)

nπ/2

CREMP

0

T/2

nπ/2

JAMUN

0·

-π/2

-T

-n/2

0

rt/2

-π/2

0

nπ/2

-π/2

0

rt/2

n-n

-π/2

0

nπ/2

-nt/2

0

nπ/2

-π/2

0

nπ/2

中

中

中

中

中

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n-\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.\n\nLIMITATIONS:\n-\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?", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, 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”).\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n1. **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?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: 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.

STRENGTHS AND WEAKNESSES:
### Strength

1. The idea of connecting MD and Score-based Learning ( more specifically, Langevin Dynamics) is interesting
2. The experimental results presented in the paper are comprehensive and demonstrate competitive performance.  

### Weakness

1. 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.
2. 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.

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.
2. Can JAMUN benefit from MD trajectory data, or is it solely designed for sampling from equilibrium states? 
3. If so, how does MD connects with Langevin Dynamics in JAMUN in physical sense?
4. 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.
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-04