henhaobio
Sign in

bioRxivBiochemistryPreprint

UniFlow: Unifying protein conformational ensemble generation and machine-learned force fields with a scalable normalizing Flow

Yikai LiuiD, Ming CheniD, Guang LiniD

PDFDOICode

In-depth summary

Claude Opus 5.5 · xhigh effort · 8 Nov 2026 · from the full text of version 1 (20 Jul 2026) · CC BY 4.0

A normalizing flow over backbone internal-coordinate angles generates protein ensembles comparable to BioEmu in one step and doubles as a differentiable coarse-grained potential for Langevin dynamics.

Generative ensemble models and machine-learned coarse-grained force fields both aim to capture equilibrium protein conformational distributions yet developed separately, one prioritizing scalable sampling and the other physical consistency for dynamics. UniFlow merges them in one transferable flow with exact likelihoods (unlike diffusion-based approaches), i.i.d. sampling, differentiable energies, and distillation into fast system-specific potentials.

Key results

  1. SuggestedUniFlow ensembles agree with reference MD about as well as BioEmu on seen CATH1 proteins and on the fully held-out ATLAS set, with similar or lower divergence on held-out octapeptides, indicating transfer to unseen sequences.

    Jensen-Shannon divergence (JSD) to MD over four collective variables (end-to-end distance, radius of gyration, RMSD to native, top two TICA components); CATH1 benchmark has 50 proteins; ATLAS never used in training; per-metric JSD values sit in Table 1 and are not reproduced in the extracted text. · Data; 4.2 Results; Tables

  2. SuggestedThe exact angular likelihood, extended with stiff harmonic bond-length terms, a Jacobian volume correction and short-range repulsion, acts as a Cartesian potential of mean force whose Langevin dynamics recover the main metastable basins seen in reference trajectories.

    BAOAB-integrated underdamped Langevin runs on one octapeptide and one CATH1 protein, 64 parallel walkers x 50 ns each (3.2 µs aggregate per system); two-dimensional TICA free-energy surfaces and one-dimensional profiles (end-to-end distance, Rg, RMSD) compared visually to reference MD; no numeric divergence given. · 4.4 UniFlow as a Machine-Learned Force Field for…; Figure legends

  3. ShownDistilling the pretrained model into a per-system student flow with lightweight multilayer perceptron (MLP) conditioners, trained on UniFlow's own samples instead of external simulation data, sharply lowers the cost of force evaluation during MD.

    2 ns run (10^6 steps, dt = 2 fs) with 64 parallel walkers: distillation plus torch.compile cut wall-clock time 134.2 to 5.3 hours (61-residue CATH1 protein) and 128.77 to 5.2 hours (eight-residue peptide) relative to the full pretrained model. · 4.4 UniFlow as a Machine-Learned Force Field for…; Tables

  4. SuggestedAblations indicate that both the geometry-aware conditioner and the geometric losses matter: Cartesian and internal-coordinate flow baselines perform poorly, and a likelihood-only UniFlow does well on octapeptides but degrades on larger proteins.

    Variants: Cartesian TARFlow (PROSE baseline), internal-coordinate Transformer spline flow (SplitFlow baseline), UniFlow with negative log-likelihood (NLL) only, full UniFlow with NLL plus clash and contact losses; SplitFlow weakest on global observables, even below Cartesian TARFlow; full model best overall (Table 2; values not in extracted text). · 3.3 Training Objective; 4.3 Ablation Studies; Tables

  5. SuggestedSince sampling needs only one inverse pass through the flow, UniFlow generates ensembles with much lower wall-clock cost than BioEmu while keeping competitive fidelity.

    Timing protocol: 2,000 i.i.d. samples generated on a single NVIDIA L40S GPU, UniFlow batch size 64 vs BioEmu default protein-length-dependent batch size; authors report a substantial wall-time reduction in Table 1, but exact times are absent from the extracted text. · 4.2 Results; Tables

Methods, data and reproducibility

A 12-layer circular coupling flow with 8-bin rational-quadratic splines models angle offsets from native torsions and ideal valence angles; each conditioner rebuilds the backbone via NeRF and feeds an SE(3)-invariant geometric Transformer plus a sequence-structure Transformer, trained with negative log-likelihood plus clash and contact losses.

  • mdCATH equilibrium MD trajectories (training)
  • BioEmu simulation corpus subsets CATH1, CATH2, Octapeptides and MEGAsim (training) · n = tens of milliseconds of MD (combined with mdCATH)
  • CATH1 seen-protein evaluation benchmark · n = 50 proteins
  • Held-out octapeptides testing length-matched generalization to unseen sequences
  • ATLAS MD dataset, excluded from training and used only as an external test set of unseen proteins
  • UniFlow-driven Langevin MD test systems: one octapeptide and one CATH1 protein · n = 64 parallel walkers x 50 ns per system
Methods & tools
UniFlow (this work), BioEmu, PROSE (TARFlow-based Cartesian normalizing flow), SplitFlow (internal-coordinate Transformer flow), Circular rational-quadratic neural spline flows, NeRF backbone reconstruction, ESM3 geometric Transformer (IPA-style attention), TICA on Cα pseudo-torsion angles, Jensen-Shannon divergence, BAOAB underdamped Langevin integrator, torch.compile
Assay conditions & compute
Sampling was timed by drawing 2,000 i.i.d. conformations with one NVIDIA L40S GPU (UniFlow batch size 64). 2 ns Langevin MD with 64 walkers: 134.2 h full model vs 5.3 h distilled plus torch.compile (61-residue protein). Training hardware and time not stated.
Reproducibility gaps
Only a code repository is linked; no data download URLs, model-weight availability, training compute or software versions are given, and the Tables 1 and 2 numeric results are absent from the extracted text.
Limitations (5)
  • Generation is conditioned on a reference backbone torsion set that must be supplied at inference, so the sampled ensemble may depend on that reference; reference-free conditioning is left for future work. · stated by the authors
  • Running MD directly with the full transferable model is expensive because each force call passes through the geometry-aware Transformer and sequential NeRF conditioner, so practical simulation requires training a separate system-specific distilled student. · stated by the authors
  • Likelihood-only training in internal coordinates is biased toward local angular statistics and degrades on larger proteins whose ensembles depend on global organization, so auxiliary geometric losses are needed. · stated by the authors
  • Force-field validation covers only two systems (an octapeptide and one CATH1 protein) using visual free-energy comparisons; the extracted text reports no quantitative divergence, kinetics, or comparison against other ML coarse-grained force fields. · noted in this summary
  • The model is backbone-only with ideal bond lengths during generation, and BioEmu is the sole external ensemble baseline; JSD values from Tables 1 and 2 are not in the extracted text, so effect sizes cannot be checked here. · noted in this summary

Read this if you build ML models of protein conformational ensembles or coarse-grained force fields and want one exact-likelihood flow that both samples backbones and drives Langevin dynamics.

Citations
0

Abstract

AO_SCPLOWBSTRACTC_SCPLOWMolecular dynamics (MD) provides a principled method for modeling equilibrium protein conformational energy landscapes, but its computational cost limits access to long timescales and larger protein systems. Recently, generative protein ensemble models and machine-learned coarse-grained force fields have emerged as complementary approaches for accelerating conformational sampling. However, they are typically developed separately despite modeling the same underlying equilibrium distribution. We introduce UniFlow, the first scalable generative model that unifies protein ensemble generation and machine-learned coarse-grained force fields for molecular dynamics simulation within a single framework. UniFlow employs an internal-coordinate normalizing flow that supports efficient i.i.d. sampling, exact likelihood evaluation, and differentiable energy and force computation. Across diverse protein systems, UniFlow generates ensembles that closely match reference MD simulations, generalizes to proteins beyond its training dataset, and samples substantially faster than diffusion-based ensemble-generation baselines. The same learned density further enables stable long-timescale molecular dynamics simulations. Together, UniFlow paves the way for a unified class of models that bridges generative ensemble modeling with physics-based molecular simulation. Codehttps://github.com/Harrydirk41/UniFlow.git

Reproducibility ledger

Nobody has reported trying to reproduce this paper yet. Did you replicate an experiment or run the code? Your report saves the next lab a week.

Sign in to file a reproducibility report.

Related papers

Similarity Map
SelectedSame fieldCross-field
THISLiu '26Xiang '26Liu '25Liu '25Li '26Yang '26
87% semantic matchBiochemistry
Yikai Liu, Guang Lin, Ming Chen·bioRxiv·Jul 2, 2026