ESMFlow: Generative Protein Ensemble Modeling

• ESMFlow is a conditional generative model that produces diverse protein structural ensembles by integrating flow-matching with rich PLM embeddings.
• It augments ESMFold by incorporating noisy backbone inputs and time tokens via Gaussian-Fourier features, modulating Evoformer-based pairwise representations.
• ESMFlow achieves enhanced ensemble diversity and improved ensemble observables, with flexible inference options including full, truncated, and distilled modes.

ESMFlow is a flow-matching generative model derived from ESMFold, constructed to sample protein structural ensembles conditional on sequence. The model transforms ESMFold—a deterministic, regression-based, sequence-to-structure predictor—into a conditional generative model that produces diverse, high-fidelity structural ensembles, aiming to directly approximate molecular dynamics (MD) distributions and ensemble observables with improved computational efficiency over traditional simulation or simple MSA subsampling (Jing et al., 2024).

1. Model Architecture and Sequence Conditioning

ESMFlow augments the original ESMFold pipeline by introducing a flow-matching paradigm. The architecture preserves the large ESM-2 protein LLM (PLM) for deriving rich residue-level embeddings, which are then processed through an Evoformer-like folding trunk and a structure module for all-atom coordinate prediction. To enable conditional flow generation, ESMFlow prepends an input embedding module that ingests a noisy C-β backbone $x_t$ and an explicit time token $t \in [0,1]$, producing a pairwise residue embedding that acts as a bias in the Evoformer.

Key architecture details:

• InputEmbedding module encodes the noisy structure and time coordinate via binned pairwise distances and Gaussian-Fourier features for $t$, with four layers of triangular self-attention/multiplication.
• Sequence embeddings are retained unaltered from the PLM and supplied to all folding trunk blocks.
• Structure module outputs both per-residue frames and all-atom coordinates, with cross-attention mixing MSA embeddings into the frame representations.

Sequence conditioning occurs throughout:

• Per-residue PLM embeddings $e_i$ are transformed to MSA-style arrays and injected across Evoformer blocks.
• Intermediate pairwise representations are modulated by flow-derived templates at every time step.
• Final structure prediction directly depends on the sequence, as enforced by cross-attention in the structure module.

2. Mathematical Formulation and Flow-Matching Objective

ESMFlow is trained as a flow-based generative model using the continuous flow-matching objective. For a given sequence $A$ and structural ensemble $\mathcal{X}$, the aim is to transport a simple prior $q(x_0)$ (typically a harmonically biased random walk for the C-βs) to the target data distribution $p(x_1|A)$ via a learned vector field $v_\theta(x, t)$ governed by the ODE

$$\frac{dx_t}{dt} = v_\theta(x_t, t),\quad t \in [0,1],\quad x_0 \sim q(x).$$

Model training minimizes

$$L_{FM}(\theta) = \frac{1}{2} \mathbb{E}_{x_0 \sim q, x_1 \sim p_{data}} \left[ \int_0^1 \left\| v_\theta(x_t, t) - u_t(x_t|x_1) \right\|^2 dt \right]$$

where $u_t(x|x_1) = \frac{x_1 - x}{1-t}$ is the oracle field along the path $x_t = (1-t)x_0 + t x_1$. The operational loss is implemented in the denoising form:

$$L(\theta) = \mathbb{E}_{t, x_0, x_1} \left[ \frac{\left\| \hat x_1(x_t, t; \theta) - x_1 \right\|^2}{(1-t)^2} \right]$$

with $\hat x_1(x_t, t; \theta)$ the model's predicted clean structure. To account for rigid-body symmetry, FAPE (Frame Aligned Point Error) replaces MSE as the principal metric, operating on $M = \mathbb{R}^{3N}/SE(3)$.

3. Training Process and Computational Considerations

Fine-tuning to protein ensemble prediction occurs in two primary stages:

• PDB ensemble stage: Training on $\sim 0.72$M PDB structures (pruned up to 2020) with MSAs from OpenProteinSet. Crop length is 256; batch size is 64.
• MD ensemble stage: Further supervised by 27k frames (82 proteins) from ATLAS all-atom MD trajectories.

The optimizer is AdamW; learning rates decline from $1\mathrm{e}{-3}$ to $1\mathrm{e}{-4}$; weight decay is $1\mathrm{e}{-2}$. Self-conditioning is used in 50% of the PDB-stage minibatches. Curriculum on time $t$ injects unnoised (deterministic) examples for stability.

Computationally, inference time scales linearly with the number of flow steps $N$:

Variant	Time/sample (s)	Description
ESMFold (baseline)	3.2	Single-point, deterministic
MSA subsampling (48 passes)	3.5	48-fold diversity via MSAs
ESMFlow (10 steps)	30.4	Full generative chain
ESMFlow (2 steps)	9.2	Truncated chain
ESMFlow (distilled)	3.1	Single-pass, distilled

Distillation enables single-pass inference with modest loss in diversity and ensemble fidelity.

4. Sampling Algorithm and Diversity Control

Sampling proceeds by initializing a random backbone $x_0 \sim$ HarmonicPrior, then iterating:

1. At each step $t$, denoise $x_t$ via ESMFold given $(A, x_t, t)$, producing $\hat x_1$.
2. RMSD-align $x_t$ and $\hat x_1$.
3. Interpolate to $x_{s}$ for the next step.
4. Repeat for $N$ steps (default $N=10$).

Diversity-precision tradeoff can be managed by shortening the chain (e.g., $N=2$), truncating initial steps, or using the distilled model.

5. Ensemble Evaluation Metrics and Empirical Results

ESMFlow ensembles are benchmarked against PDB and MD diversity baselines using:

• PDB metrics: Precision, recall, and diversity computed from $lDDT_{C\alpha}$ comparisons.
• MD metrics: Pairwise Cα-RMSD, RMSF (root mean square fluctuation), root-mean Wasserstein distance (RMWD), and various ensemble observables (e.g., weak contacts, transient SASA exposures, mutual information on state transitions).

Empirical results for median performance across 100 PDB and 82 MD targets:

Method	Precision	Recall	Diversity	Pairwise RMSD (Å)	RMWD (Å)	Weak Contacts J	Time/sample (s)
ESMFold	0.809	0.761	0.000	--	--	--	3.2
MSA subsampling (48)	0.757	0.760	0.125	1.67	4.28	0.37	3.5
ESMFlow (10 steps)	0.777	0.777	0.210	3.25	3.60	0.55	30.4
ESMFlow (2 steps)	0.795	0.774	0.100	--	--	--	9.2
ESMFlow (distilled)	0.775	0.752	0.152	2.76	4.23	0.48	3.1

ESMFlow (full or distilled) achieves greater diversity and improved weak/transient contact statistics and exposure behavior compared to MSA subsampling, at a modest computational cost (Jing et al., 2024).

6. Significance and Broader Context

ESMFlow demonstrates that coupling high-fidelity regression models (ESMFold, PLMs) with flow-matching frameworks produces generative models superior to naïve ensemble generators in terms of accuracy, conformational diversity, and statistical fidelity to MD. The design supports fine-grained diversity-precision tradeoff and efficient distillation to single-pass inference, yielding practical advantages for rapid ensemble generation and protein design pipelines. A plausible implication is that flow-matching, combined with expressive PLM-based structure networks, could replace computationally demanding MD sampling in some contexts, accelerating downstream biophysical or design analyses.

7. Limitations and Future Directions

ESMFlow's runtime increases linearly with the number of flow steps, leading to a ~10-fold penalty versus deterministic predictors at maximum diversity. Distillation partially alleviates this at some cost to ensemble quality. The method inherits sequence–structure biases from ESMFold and is limited by the expressive capacity of the denoiser and flow parameterization. Possible research directions include scaling flow depth, hybridizing with explicit physical priors, integrating side-chain or solvent modeling, and further leveraging large-scale MD or evolutionary data for training. Extending ESMFlow to non-protein (e.g., RNA, complex assemblies) remains an open problem (Jing et al., 2024).