PAPER KEY: 6Q8QMLF4
TITLE: P2DFlow: A Protein Ensemble Generative Model with SE(3) Flow Matching
AUTHORS: Jin, Yaowei; Huang, Qi; Song, Ziyang; Zheng, Mingyue; Teng, Dan; Shi, Qian

P2DFlow: A Protein Ensemble Generative
Model with SE(3) Flow Matching
Yaowei Jin1, Qi Huang5, Ziyang Song6, Mingyue Zheng2,3,4, Dan Teng2, *, Qian Shi1, *
1 Lingang Laboratory, Shanghai 200031, China.
2 Drug Discovery and Design Center, State Key Laboratory of Drug Research,
Shanghai Institute of Materia Medica, Chinese Academy of Sciences, 555 Zuchongzhi
Road, Shanghai 201203, China.
3 University of Chinese Academy of Sciences, No.19A Yuquan Road, Beijing 100049,
China.
4 School of Chinese Materia Medica, Nanjing University of Chinese Medicine, Nanjing
210023, China.
5 Institute for Electric Light Sources, School of Information Science and Technology,
Fudan University, Shanghai 200438, P. R. China.
6 Shanghai Key Lab of Chemical Assessment and Sustainability, School of Chemical
Science and Engineering, Tongji University, Shanghai 200092, P. R. China.
*To whom correspondence should be addressed.
*(Qian Shi) E-mail: shiqian@lglab.ac.cn
*(Dan Teng) E-mail: tengdan@simm.ac.cn


Abstract
Biological processes, functions, and properties are intricately linked to the ensemble of
protein conformations, rather than being solely determined by a single stable
conformation. In this study, we have developed P2DFlow, a generative model based on
SE(3) flow matching, to predict the structural ensembles of proteins. We specifically
designed a valuable prior for the flow process and enhanced the model’s ability to
distinguish each intermediate state by incorporating an additional dimension to describe
the ensemble data, which can reflect the physical laws governing the distribution of
ensembles, so that the prior knowledge can effectively guide the generation process.
When trained and evaluated on the MD datasets of ATLAS, P2DFlow outperforms
other baseline models on extensive experiments, successfully capturing the observable
dynamic fluctuations as evidenced in crystal structure and MD simulations. As a
potential proxy agent for protein molecular simulation, the high-quality ensembles
generated by P2DFlow could significantly aid in understanding protein functions across
various scenarios. Code is available at https://github.com/BLEACH366/P2DFlow.
Key words: protein ensembles, molecular dynamics, flow matching, equivariant graph
neural network
1. Introduction
Proteins exhibit a dynamic nature that leads to the generation of diverse conformations.
Crucial biological functions are executed by relying on the distinct states, collective
motions, and disordered fluctuations within the protein ensemble[1, 2]. Thus,


A Protein Ensemble Generative Model with SE(3) Flow Matching
understanding the distribution of these ensembles is essential for elucidating the
mechanism by which proteins function in different environments. While experimental
measurements, such as crystallographic B-factors and NMR spectroscopy, can provide
some insight into conformational changes, they are limited in spatial and temporal
scale[3, 4]. Although methods like AlphaFold[5], ESMFold[6], and other deep learning
approaches have demonstrated excellent performance in predicting the crystal structure
of proteins, it is still challenging to offer diverse predictions for protein conformational
ensembles [7-10].
Several computational methods are available for conformational sampling. Traditional
approaches include Monte Carlo (MC) and molecular dynamics (MD). By providing an
initial structure of a molecular, these methods explore its conformational space based
on the forces acting upon it, which can be calculated using molecular mechanics or
quantum mechanics[11]. However, these methods face several challenges: the
computational efficiency declines rapidly as the number of atoms and degrees of
freedom increase, and both MC and MD are reliant on force fields and energy, making
it difficult to overcome high energy barriers. Consequently, these methods often
become trapped in local minima. To address these issues, enhanced sampling methods,
such as umbrella sampling [12] and meta dynamics [13], are employed for broader
exploration,.
An alternative approach involves the use of generative models leveraging machine
learning and deep learning techniques. For instance, Boltzmann generator [14] represents
one of the earliest attempt to use normalizing flow to sample system-specific
conformational distributions from random noise. However, it requires pre-acquired


A Protein Ensemble Generative Model with SE(3) Flow Matching
simulation data for specific protein systems. Multiple sequence alignment (MSA)
subsampling combined with AlphaFold [15-17] can predict different structures for each
MSA subset. While effective in certain contexts, it faces two limitations. First, the
results depend on the partition of the MSA, making it difficult to predict protein
ensembles lacking homologous proteins. Furthermore, it does not provide the approach
for training the model on protein ensemble data, which is critically important.
STR2STR [4] employs a diffusion model to learn the score matching of distributions
from random noise to protein ensemble distributions, combining stochastic perturbation
and score to guide the direction of conformational changes. It does not rely on
simulation data during either training or inference and is capable of performing zero
shot conformation sampling. However, the large variability in its predicted results
makes it difficult to focus on a particular low-energy conformation. This may be due to
inconsistencies in its prior distribution during training and sampling processes.
AlphaFlow [18] integrates AlphaFold and ESMFold with flow matching, fine-tune these
models to convert regression models to generative models. It can generate accurate and
diverse PDB structures, which can be utilized to compute additional dynamic properties.
Nonetheless, it encounters the challenges in generating non-existent intermediate states
between two or more minima, similar to the issues faced by AlphaFold in certain
proteins.
To address these limitations, we propose a new framework that samples protein
ensembles via a SE(3) equivariant flow matching model named P2DFlow (Fig. 1).
P2DFlow is trained on MD simulation data, which contains different macro states of
proteins (as seen in Fig. 4), allowing it to learn the conformational changes within
protein ensembles. To enhance the model's generalization ability while reducing the


A Protein Ensemble Generative Model with SE(3) Flow Matching
difficulty of training, we use a strong prior, ESMFold predictions with perturbations of
coordinates (detailed in Sec. 2.1), to obtain an estimated structure. Compared to
Gaussian prior and harmonic prior, the stronger prior that we use guides the model to
generate accurate structures by introducing more precise biases in bond lengths and
bond angles. This approach differs significantly from AlphaFlow, which uses harmonic
prior and relies on AlphaFold or ESMFold as the main component to control structure.
In contrast, P2DFlow employs ESMFold predictions with perturbations as a stronger
prior and utilizes a new SE(3) equivariant block to adjust the structure. To distinguish
between various conformations of the same protein, we introduce a new dimension
called ‘approximate energy’, which maps MD results onto a low-dimensional plane and
uses the density to represent probability (detailed in Sec. 2.2). It guides the model in
generating conformations with distinct energies, thereby avoiding the generation of
non-existent intermediate states.
We evaluate the performance of P2DFlow against AlphaFlow and STR2STR, two
representative and advanced models for protein ensemble generation. Based on
extensive experiments on the test dataset of ATLAS[19], it is evident that P2DFlow
outperforms other baseline models on the metrics which can reflect the fidelity and
dynamic properties of the generated ensembles compared to the ground truth.
Visualization results indicate that P2DFlow effectively captures important changes in
residue contact and more accurately recovers protein ensembles distributions.
Moreover, the ablation experiment shows that the ‘approximate energy’ significantly
aids in constructing the conformation distribution of proteins. Thus, P2DFlow can be
used to predict protein ensembles without the need for expensive MD simulations.


A Protein Ensemble Generative Model with SE(3) Flow Matching
Figure 1. Conceptual overview of P2DFlow. Protein structures are firstly sampled
from a prior based on perturbed ESMFold predictions. Subsequently, an ‘approximate
energy’ sampling serves as a condition to guide the generation process to produce
ensemble distribution.
2. Method
2.1 SE(3) Flow Matching
A flow-matching generative framework[20] aims to learn continuous normalizing flows
(CNFs) [21], generating target distributions by predicting a vector field and integrating
the ordinary differential equation (ODE).
For a flow φ(x, t), there is a related vector field v(x, t), defined by the following ODE:
dφ(x, t)
dt = v(x, t) , φ(x, 0) = x0 (1)
where t ∈ [0, 1] is a continuous parameter, x0 is a sample of prior distribution p0.
We can use the flow to transform the simple prior p0 towards the data distribution p1
by integrating the vector field v(x, t):


A Protein Ensemble Generative Model with SE(3) Flow Matching
x1 − x0 = ∫ v(x, t) dt
1
0
(2)
For the flow, since it only fixes the distribution at t = 1 and the prior at t = 0 , the
interpolation process between them can be manually defined. In our method, we choose
ESMFold prediction with perturbations as the prior. It means that, for a given protein
sequence requiring ensemble generation, we first use ESMFold to predict the static
structure. Then, Gaussian noise with specific variance is added to the coordinates of
each residue to generate the prior distribution. Previous studies has shown that the prior
distribution closer to data distribution which needs to be learned can lead to superior
performance[22], so we identify this as an especially valuable inductive bias since it will
offer a suitable bond length and torsion angle as the initial value of the iteration,
compared with Gaussian prior and harmonic prior [23, 24].
For the interpolation process, following previous work [25-27], the backbone atom
positions of each residue are parameterized by an orientation preserving rigid
transformation T ∈ SE(3), that maps from fixed coordinates of N∗ , Cα∗ , C∗ , O∗
centered at Cα∗ = (0,0,0). Each frame T = (r, x) consists of a rotation r ∈ SO(3) and
a translation vector x ∈ R3. As for the side chain of the residue, torsion angles θ are
used to define the twists. Since our focus lies on the backbone of protein, we only
interpolate for translation and rotation accompanied by the use of optimal transport (OT)
path [28]:
Translations (R3): xt = (1 − t)x0 + tx1 (3)
Rotations ( SO(3) ): rt = expr0 (t logr0(r1)) (4)


A Protein Ensemble Generative Model with SE(3) Flow Matching
We calculate expr0 and logr0 using Rodrigues’ formula [29]. The vector field for
translation and rotation are then computed as follow:
ẋt = x1 − xt
1 − t , ṙ = logrt( r1)
1 − t (5)
These vector fields are learned with SE(3) equivariant neural network (mentioned in
Section 2.2). For the torsion angle within residues, we predict it using the node
embedding of the corresponding residues with a Multi-Layer Perceptron (MLP) due to
its SE(3) invariant[30], which is similar to the approach utilized in AlphaFold.
With these considerations and calculations mentioned above, the training loss of
P2DFlow is defined as:
L = Εt,p1,p0{ 1
(1 − t)2 (‖ ̂x1 − x1‖2 + ‖r̂1 − r1‖2)
+ α(‖ ̂C1( ̂x1, r̂1, ̂θ) − C1‖2 + ‖ ̂D1( ̂x1, r̂1, ̂θ) − D1‖2)} (6)
Where ( ̂x1, r̂1, ̂θ) refer to predicted translation, rotation, and torsion at t = 1 .
( ̂C1, ̂D1) refer to the atom coordinates and distance matrix of all atoms, which is
recovered using ( ̂x1, r̂1, ̂θ) and standard frames proposed in AlphaFold at t = 1. α is
a hyperparameter to adjust the auxiliary loss with the vector field loss. Furthermore, we
subtract the Center of Mass (CoM) from the prior sample and each protein data to
achieve a zero CoM. This step is crucial for ensuring that the distribution of the sampled
results remains SE(3) equivariant.
2.2 Network Architecture
To learn the aforementioned vector field, we developed P2DFlow, a flow-matching
model that is SE(3) equivariant. The model's equivariance is achieved by the stacking
of Invariant Point Attention (IPA) module [5] and E(n) equivariant Graph Neural


A Protein Ensemble Generative Model with SE(3) Flow Matching
Network (EGNN) module [31]. These components are utilized to encode the spatial
features of the protein. In the molecular graph representation, residues are considered
as nodes, and edges are defined by two criteria: the cutoff coordinate distance between
adjacent residues and the cutoff sequence distance. We utilize the union of these criteria
to establish the edges.
Figure 2. Workflow of P2DFlow. Protein sequence and ‘approximate energy’ are
inputs of P2DFlow. We utilize ESM-2 to get embeddings from sequence, and use
ESMFold + Perturb to get the noised structure as the prior distribution, then apply a
SE(3) equivariant block composed with IPA and EGNN to predict the vector filed for
the flow process. The differences between training and sampling are marked with
different border colors.
The inputs of the model include sequence and ‘approximate energy’ (Figure 2). To
improve the generalization ability of the model, we utilize the ESM-2 protein language
model to generate the initial Node Representation (Node Repr.) and Pair Representation
(Pair Repr.) for residues. After successive aggregations and updates of the Node Repr.


A Protein Ensemble Generative Model with SE(3) Flow Matching
and Pair Repr. using the SE(3) Equivariant Block, we obtain a compressed, full-graph
representation of the protein structure, which is sampled from prior of ESMFold
prediction with perturbation. Subsequently, we employ the representation to predict the
vector fields associated with translation and rotation.
To distinguish structures with different energies in the conformation set, we introduce
the concept of ‘approximate energy’. We project the molecular dynamics (MD)
simulation ensembles onto a two-dimensional plane defined by the radius of gyration
(RG) and the root-mean-square deviation (RMSD) relative to the crystal structure. We
then compute the Gaussian kernel density of this 2D map and apply the Boltzmann
Equation to convert the density values into ‘approximate energy’ values after
normalization.
The sampling process is illustrated in Figure 2. Given a protein sequence as input, the
process begins with the prediction of a stable structure using ESMFold. Subsequently,
we add Gaussian noise with the same variance used during training to perturb the
structure, so that we can get the initial perturbed structure from the specialized prior
distribution. We then sample the 'approximate energy' from Boltzmann Distribution,
and keep it unchanged throughout the entire autoregressive sampling process to guide
the generation. Following the prediction of vector fields, a Euler integration step is
applied to solve the ODE. Ultimately, we reconstruct the coordinates of each atom
based on the translation, rotation, and torsion angles, utilizing the frame representation
provided by AlphaFold, to derive the protein conformation structure. After conducting
a sufficient number of samples, we obtain the predicted protein ensembles.


A Protein Ensemble Generative Model with SE(3) Flow Matching
Other operational details are in line with those of FrameFlow [27]. We modify the loss
function by substituting 1/(1 − t)2 with 1/(1 − min{ t, 0.9})2 to prevent the loss
from exploding. Furthermore, we conduct a pre-alignment[32, 33] to align x0 and x1
using the Kabsch algorithm. Specifically, we solve r∗ = argminr∈SO(3)‖rx0 − x1‖2
and use the aligned position r∗x0 during training.
3. Experiments
3.1 Setup
Training To learn the dynamic distribution of protein ensembles, we utilize ATLAS
dataset[19], which contains ~1300 MD simulation results for each kind of protein. To
select representative structures from the conformations, we compute the ‘approximate
energy’ (mentioned in section 2.2) for each of them, then choose 11 structures at equal
intervals of ‘approximate energy’. For evaluation, we randomly choose ~100 ensembles
from ATLAS dataset excluding the training set. Our model is trained for 2 days using
4 NVIDIA A100-80G GPUs.
Metrics To assess the performance of P2DFlow, we use evaluation metrics which can
be categorized into: (a) Fidelity. Reflects the distributional gap between sampled
ensemble and reference MD simulation, which contains Jensen-Shannon divergence of
pairwise distance[4] (PWD J) and radius of gyration (RG J), root mean Wasserstein
distance (RMWD W2) and root mean square fluctuation (RMSF), intersection over
union of Ramachandran plots (RP IoU); (b) Dynamics. Shows the observable structural
change which is often associated with thermal fluctuations, such as the Weak Contacts
Ja and Transient Contacts Ja [18].


A Protein Ensemble Generative Model with SE(3) Flow Matching
Baseline We compare our results with Alphaflow and STR2STR. Alphaflow is a flow
matching model based on fine-tuning of AlphaFold and ESMFold, and STR2STR is a
diffusion model which changes inference stage by adding AlphaFold prediction as the
initial structure. We use their pretrained weights offered at GitHub to sample ensembles
on the test set of ~100 proteins from ATLAS dataset.
3.2 Results
Table 1 shows the evaluation results for P2DFlow, AlphaFlow and STR2STR. For each
method, we compare the predicted ensemble with the ground truth MD ensembles and
evaluate them on various metrics. We can see that P2DFlow performs better than
AlphaFlow and STR2STR in PWD J , RMWD W2 , RMSF, RP IoU and Weak
Contacts Ja. This indicates that P2DFlow more realistically reflect the physical space
distribution of structures, resulting in generated ensembles that closely resemble the
ground truth derived from MD simulation. As for the dynamics of protein residue
contacts, P2DFlow can predict the residue pairs with a higher likelihood of dissociation,
aiding the identification of potential protein pockets.
Table 1. Evaluation on MD ensembles. J : Jensen-Shannon divergence; W2: 2
Wasserstein distance; IoU: Intersection over Union; Ja: Jaccard similarity. Among
these metrics, RP IoU, Weak Contacts Ja and Transient Contacts Ja are the higher
the better (); while PWD J, RG J, RMWD W2, RMSF are the lower the better ().
The best result from generative models is bolded.


A Protein Ensemble Generative Model with SE(3) Flow Matching
STR2STR AlphaFlow P2DFlow
Fidelity
PWD J () 0.446 0.452 0.412
RG J () 0.461 0.519 0.493
RMWD W2 () 78.246 33.628 23.020
RMSF ()
RP IoU ()
3.997
0.240
0.576
0.276
0.523
0.301
Dynamics Weak Contacts Ja () 0.329 0.458 0.478
Transient Contacts Ja () 0.148 0.288 0.230
Figure 3. Distribution of evaluation results for some metrics. a) RMWD W2. b)
Weak Contacts Ja. c) Ramachandran plots for 4wyw_A. There are ~100 types of
protein ensembles in the test set, and we set bins for the metrics to generate the
histogram. After that, we use a Gaussian distribution to fit the histogram. For the


A Protein Ensemble Generative Model with SE(3) Flow Matching
Ramachandran plots of 4wyw_A, we use scatter points of different colors to represent
the model sampling and MD results.
According to Figure 3, it illustrates the distribution of RMWD W2 and Weak Contacts
Ja results on the ATLAS test set. We can see that there are a few bad outliers predicted
by AlphaFlow and STR2STR, whereas P2DFlow does not exhibit such issues.
Additionally, the Gaussian distribution of P2DFlow generated results displays narrower
peaks, indicating more stable predictions with a lower likelihood of extreme outcomes
across different protein ensembles. As for the Ramachandran plots, we 