CoMoFlow: Controllable Protein Conformational Ensemble Generation via Explicit Collective-Mode Modeling

Yusong Wang; Zheyuan Gu; Ailiang Lin; Minghao Shao; Yicheng Xu; Jiang Cai; Prayag Tiwari; Mingkun Xu

cs.AINov 6, 2026

CoMoFlow: Controllable Protein Conformational Ensemble Generation via Explicit Collective-Mode Modeling

Authors: Yusong Wang, Zheyuan Gu, Ailiang Lin, Minghao Shao, Yicheng Xu, Jiang Cai, Prayag Tiwari, Mingkun Xu

Organizations: Beijing University of Chinese Medicine · Peking University · New York University · Guangdong Institute of Intelligence Science and Technology ·

Abstract

Protein conformational ensemble generation aims to efficiently characterize the structural variations accessible to a protein beyond a single static structure. Recent flow-based methods efficiently generate conformational ensembles, but they model coordinated structural variation implicitly, which can limit ensemble fidelity and hinder direct control. To address this limitation, we propose CoMoFlow, a controllable framework for generating protein backbone conformational ensembles by explicitly modeling collective variation. CoMoFlow constructs protein-specific collective modes from the equilibrium structure, learns the distribution of their coefficients across the molecular dynamics (MD) ensemble, and uses sampled collective variations to guide SE(3) flow matching through feature conditioning and prior recentering. Experiments on ATLAS and de novo proteins show that CoMoFlow improves reconstruction of MD conformational distributions and coordinated structural variation relative to strong flow-based baselines, while enabling direct control over targeted conformational changes.

Figures & tables

Figure 1: Baseline flow matching and CoMoFlow on the ATLAS protein 7asgA. Baseline methods learn from individual MD conformations, whereas CoMoFlow learns collective modes from the equilibrium structure and their variation across the MD ensemble to guide generation. Manipulating selected modes enables controllable conformational changes. Å denotes inter-region distances.

Figure 2: Overview of the CoMoFlow framework, consisting of equilibrium encoding, collective-mode learning and coefficient distribution modeling, and coefficient-guided SE(3) flow generation.

Dataset	Method	$r$ $\uparrow$	MAE $\downarrow$	Med.	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$	Time [s] $\downarrow$
		RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$	Time [s] $\downarrow$
ATLAS	BioEmu	0.83	1.29	2.34	2.84	0.80	1.65	36	1.9
	AlphaFlow	0.86	0.59	1.51	1.35	0.86	1.47	41	32.0
	ConfDiff	0.88	0.62	2.00	1.45	0.86	1.41	39	20.2
	AlphaFlow-T	0.92	0.41	1.17	0.91	0.89	1.28	47	32.6
	ESMFlow-T	0.92	0.52	0.94	1.22	0.89	1.48	47	11.2

Table 1: Performance comparison on the ATLAS and de novo proteins datasets. The reference RMSF medians (Med.) are 1.48 for ATLAS and 0.91 for de novo proteins.

Variant	$r$ $\uparrow$	MAE $\downarrow$	Med. (1.48)	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
	RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
w/o collective guidance	0.89	0.47	1.50	0.77	0.87	1.38	29
Feature conditioning only	0.91	0.43	1.47	0.74	0.88	1.33	31
Prior recentering only	0.90	0.41	1.46	0.75	0.90	1.29	32
w/o $\mathcal{L}_{\mathrm{rec}}$	0.91	0.41	1.46	0.74	0.90	1.28	33
w/o $\mathcal{L}_{\mathrm{coeff}}$	0.90	0.42	1.44	0.76	0.89	1.31	31

Table 2: Ablation of collective-conditioning pathways and training objectives on the ATLAS dataset.

Figure 3: Visualization of coordinated conformational variation recovery for protein 7asgA. (a) Residue-wise DCCMs of the MD reference and generated ensembles. (b) Absolute errors.

Metric	ATLAS	de novo
Corr. $\uparrow$	0.983	0.971
Slope $\rightarrow 1$	0.864	0.797
NMAE $\downarrow$	0.345	0.401
Off-target $\downarrow$	0.32	0.22

Table 3: Controllable conformation generation.

Figure 4: Single-mode controllable generation on a representative protein.

Figure 5: Targeted exploration and conformational-basin recovery. (a) Cumulative opening/closing target success for BBFlow, passive CoMoFlow, and CoMoFlow with CEM. (b) Basin coverage and population JSD for BBFlow and CoMoFlow on ATLAS and de novo proteins.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	$K$	$r$ $\uparrow$	MAE $\downarrow$	Med.	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
		RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
ATLAS	4	0.90	0.43	1.45	0.75	0.90	1.28	34
	8	0.92	0.38	1.48	0.71	0.93	1.21	37
	16	0.92	0.36	1.49	0.72	0.92	1.19	38
de novo	4	0.84	0.27	0.89	0.30	0.86	0.64	39
	8	0.86	0.23	0.93	0.27	0.89	0.59	43

Appendix

Table 4: Sensitivity to the number of collective conformational modes

K

on the ATLAS and de novo proteins. The reference RMSF medians (Med.) are 1.48 for ATLAS and 0.91 for de novo proteins.

Dataset	$C$	$r$ $\uparrow$	MAE $\downarrow$	Med.	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
		RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
ATLAS	1	0.91	0.42	1.46	0.74	0.90	1.27	34
	2	0.92	0.38	1.48	0.71	0.93	1.21	37
	4	0.92	0.37	1.50	0.71	0.91	1.19	37
de novo	1	0.85	0.26	0.90	0.29	0.87	0.62	40
	2	0.86	0.23	0.93	0.27	0.89	0.59	43

Appendix

Table 5: Sensitivity to the number of Gaussian mixture components

C

on the ATLAS and de novo proteins. The reference RMSF medians (Med.) are 1.48 for ATLAS and 0.91 for de novo proteins.

Figure 6: Joint sensitivity analysis of $\lambda_{\mathrm{rec}}$ , $\lambda_{\mathrm{coeff}}$ , and $\lambda_{\mathrm{ens}}$ on the validation set. Each panel fixes $\lambda_{\mathrm{ens}}$ , while varying $\lambda_{\mathrm{rec}}$ and $\lambda_{\mathrm{coeff}}$ . The white dashed box highlights the unit-weight setting, where $\lambda_{\mathrm{rec}}=\lambda_{\mathrm{coeff}}=\lambda_{\mathrm{ens}}=1$ .

Dataset	Mode basis	Coverage $\uparrow$	DCCM $r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
ATLAS	ANM	0.63	0.89	1.31	33
	Learned	0.79	0.93	1.21	37
de novo	ANM	0.60	0.85	0.66	39
	Learned	0.76	0.89	0.59	43

Appendix

Table 6: Comparison with classical ANM modes using

K=8

collective modes.

Dataset	MD fraction	Coverage $\uparrow$	$r$ $\uparrow$	MAE $\downarrow$	Med.	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
		Coverage $\uparrow$	RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
ATLAS	$10\%$	0.69	0.88	0.47	1.42	0.79	0.88	1.36	32
	$25\%$	0.75	0.91	0.42	1.46	0.75	0.91	1.26	36
	$50\%$	0.77	0.90	0.40	1.49	0.73	0.92	1.24	36
	$100\%$	0.79	0.92	0.38	1.48	0.71	0.93	1.21	37
de novo	$10\%$	0.67	0.82	0.30	0.88	0.34	0.84	0.70	38

Appendix

Table 7: Effect of MD supervision fraction on conformational ensemble generation. The reference RMSF medians (Med.) are 1.48 for ATLAS and 0.91 for de novo proteins.

Dataset	Reference	$r$ $\uparrow$	MAE $\downarrow$	Med.	MAE $\downarrow$	$r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
		RMSF			PwRMSD	DCCM	PCA $W_{2}$ $\downarrow$	$J_{\mathrm{tr}}$ (%) $\uparrow$
ATLAS	Original	0.92	0.38	1.48	0.71	0.93	1.21	37
	$Q_{25}$	0.90	0.42	1.43	0.76	0.90	1.29	34
	$Q_{50}$	0.87	0.49	1.59	0.86	0.88	1.43	35
	$Q_{75}$	0.82	0.57	1.64	1.01	0.82	1.58	28
de novo	Original	0.86	0.23	0.93	0.27	0.89	0.59	43

Appendix

Table 8: Effect of alternative reference structures on conformational ensemble generation. The reference RMSF medians (Med.) are 1.48 for ATLAS and 0.91 for de novo proteins.

Figure 7: Learned coefficient distributions. Joint distributions of MD-projected coefficients (blue) and samples from

q(c\mid x_{\mathrm{eq}})

(red) along the two highest-variance coefficient dimensions. Ellipses indicate the corresponding covariance structure.

Dataset	Covariance	Coeff. NLL $\downarrow$	CorrErr $\downarrow$	DCCM $r$ $\uparrow$	PCA $W_{2}$ $\downarrow$	Basin cov. $\uparrow$	Pop. JSD $\downarrow$
ATLAS	Diagonal	9.84	0.118	0.90	1.29	0.55	0.027
ATLAS	Full	9.21	0.074	0.93	1.21	0.62	0.018
de novo	Diagonal	9.47	0.126	0.86	0.65	0.51	0.043
de novo	Full	8.96	0.081	0.89	0.59	0.59	0.030

Appendix

Table 9: Ablation of covariance parameterization in the coefficient distribution.

Dataset	$\|\Delta c_{k}\|/\sigma_{k}$	C $\alpha$ clash (%) $\downarrow$	C $\alpha$ –C $\alpha$ bond MAE (Å) $\downarrow$
ATLAS	0	0.37	0.052
	1	0.39	0.055
	2	0.42	0.061
de novo	0	0.43	0.057
	1	0.42	0.060
	2	0.47	0.066

Appendix

Table 10: Physical plausibility along controlled single-mode traversals.

Protein Conformational Ensemble Generation

See all 8 papers

Jun 3, 2026

ProHiFlo: Hierarchical Flow Matching with Functional Guidance for De Novo Protein Generation

De novo protein generation has transformative potential in therapeutic design, enzyme engineering, and synthetic biology. While diffusion-based and flow matching approaches have achieved progress, they typically operate at single resolution and lack mechanisms for incorporating functional constraints. We introduce ProHiFlo, a hierarchical flow matching framework with three innovations: (1) coarse-to-fine generation that models backbone geometry before refining to all-atom coordinates, reducing computational cost while maintaining accuracy; (2) functional guidance leveraging pretrained predictors to steer generation toward desired properties without retraining; (3) adaptive SE(3)-equivariant architecture for efficient multi-scale processing. Experiments on unconditional generation, motif scaffolding, and functional design demonstrate state-ofthe-art performance while requiring 4 fewer sampling steps. On enzyme active site scaffolding, ProHiFlo achieves 58.9% success rate compared to 41.2% for RFDiffusion.

Chuanzhen Wang, Meade Cleti, Pete Jano

3Tongji University · 1Arizona State University · University of Wisconsin-Madison

Flow Matching 3D Molecular Generation Protein Design Conditional Generative Modeling Protein Structure Generation

Ask PDF

Jun 7, 2026

Few-step Cofolding with All-Atom Flow Maps

All-atom generative modeling of 3D biomolecular complexes has emerged as the dominant paradigm for predicting the structure of proteins and protein-ligand systems. Generating structures at the atomic level of fidelity, however, typically requires expensive iterative diffusion rollouts, making both conventional deployment and inference-time search techniques computationally costly. In this paper, we introduce the Denoiser Cofolding All-Atom Flowmap (DeCAF) framework for distilling state-of-the-art all-atom cofolding models into all-atom flow maps that produce high-quality samples in only a few inference steps. We build DeCAF on a denoiser-based formulation of flow maps with endpoint losses that naturally support SE(3) rigid alignment, which we show is critical for training accurate models. We further derive a simple change of variables that lets DeCAF operate in the σ-space noise schedule of EDM-style architectures, enabling direct distillation from pretrained cofolding diffusion models. Equipped with DeCAF's flowmap lookahead, we introduce a purpose-built inference-time framework that improves sampling through reward-guided search. Empirically, DeCAF-Boltz statistically improves over Boltz-1x in both accuracy (RMSD) and physical validity scores of protein-ligand poses at strict NFE budgets on the challenging Runs N' Poses, while also showing a more optimal Pareto frontier across all inference compute budgets on PoseBusters. Distilling the state-of-the-art Pearl cofolding model, DeCAF-Pearl outperforms diffusion-based cofolding models and matches its teacher on success rate while using 5x fewer NFEs. We release our code at https://github.com/genesistherapeutics/decaf.

Gianluca Scarpellini, Ron Shprints, Peter Holderrieth +7

1Genesis Molecular AI · Massachusetts Institute of Technology · Carnegie Mellon University +2

Protein Structure Prediction Diffusion Distillation Few-Step Diffusion Sampling 3D Molecular Generation

Ask PDF

Jul 29, 2026

SE(3)-MeanFlow: Few-Step Protein Backbone Generation on Lie Groups

Generative modeling of protein backbones promises the de novo design of proteins with prescribed structural and functional properties. Existing diffusion and flow-matching models produce high-quality backbones on SE(3)^N, but inference requires numerically integrating an ODE over hundreds of network evaluations, each involving a Lie group exponential map - a bottleneck for high-throughput design campaigns. We introduce SE(3)-MeanFlow, a few-step generative framework that extends MeanFlow from Euclidean space to the Lie group geometry of protein frames. Working natively in the Lie algebra so(3) and in R^3, we derive closed-form average-velocity identities for rotations and translations, giving simulation-free training targets. We further introduce an SE(3) alpha-Flow objective that removes the Jacobian-vector product from the rotation branch and serves as a warm-up stage, after which training switches to a small-t stabilized MeanFlow loss that is used for the remainder of pretraining and for rectification-based post-training. In protein backbone generation, SE(3)-MeanFlow matches or exceeds flow-matching baselines that use several times more sampling steps, and its advantage widens in the few-step regime, where rectification lets it lead at every matched budget - at a modest cost in diversity.

Yikun Bai, Binghang Lu, Yikai Liu +7

Department of Mathematics, Purdue University · School of Electrical and Computer Engineering, Purdue University · Department of Biochemistry, Purdue University +5

Flow Matching Lie Group Methods Generative Modeling 3D Molecular Generation Mean Flow Molecular Generation+2

Ask PDF

Abstract

Protein Conformational Ensemble Generation

Explore similar work

ProHiFlo: Hierarchical Flow Matching with Functional Guidance for De Novo Protein Generation

Few-step Cofolding with All-Atom Flow Maps

SE(3)-MeanFlow: Few-Step Protein Backbone Generation on Lie Groups