UniFlow

Molecular dynamics (MD) is the principled route to a protein's equilibrium conformational landscape, but its computational cost puts long timescales and larger systems out of reach. Two families of machine-learning methods have grown up to relieve that cost, and they have grown up apart: generative ensemble models, which draw conformations directly from a learned distribution, and machine-learned coarse-grained force fields, which supply the energies and forces that drive accelerated simulation. Both describe the same equilibrium distribution, yet they are typically built and evaluated as separate tools.

UniFlow, developed at Purdue University by a group led by Guang Lin, collapses that split into one model: an internal-coordinate normalizing flow that acts as both ensemble sampler and coarse-grained force field, supporting efficient i.i.d. sampling, exact likelihood evaluation, and differentiable energy and force computation from a single learned density. Structure-conditioned generators such as AlphaFlow and BioEmu-1 also sample ensembles from one trained checkpoint, but a diffusion or flow-matching sampler does not hand back an exact, inexpensive likelihood, so the same network cannot be reused as an energy function. A normalizing flow can: it is an invertible map with a tractable Jacobian, so its density is available in closed form and differentiable with respect to coordinates, making the generator an energy model whose gradients are forces.

Across diverse protein systems the model produces ensembles that closely match reference MD, generalizes to proteins outside its training set, and samples substantially faster than diffusion-based ensemble-generation baselines, with the same density supporting stable long-timescale MD. The work is a preprint awaiting peer review, and the released code ships training and sampling scripts with configuration files but no pretrained checkpoint.

#Key Features

• One density, two jobs: A single learned distribution both draws conformational ensembles and scores the energies and forces that drive simulation, instead of training a sampler and a force field separately against the same target.
• Exact likelihood evaluation: The invertible architecture yields closed-form probabilities for arbitrary conformations, computing internal-coordinate density and the Cartesian-space likelihood with the appropriate Jacobian correction, so existing trajectories can be scored rather than only generated.
• Differentiable forces for Langevin dynamics: Gradients of the learned energy drive underdamped Langevin simulation, with an optional generalized hybrid Monte Carlo correction for unbiased sampling.
• Internal-coordinate representation: Operating on bonds, angles, and torsions rather than raw Cartesian coordinates builds in local backbone geometry and gives the flow a well-conditioned space to model.
• Generalization beyond training proteins: One checkpoint, conditioned on a target native structure, produces ensembles for proteins absent from the training dataset.

#Technical Details

The default configuration stacks 12 flow layers at a hidden width of 256, built from six graph-neural-network blocks and six transformer blocks with eight attention heads each; training defaults cap systems at 256 residues and run in bfloat16 mixed precision. Training data is supplied as per-protein directories pairing a native PDB structure with MD trajectory files, and the objective combines a negative log-likelihood term with geometric penalties for steric clashes and contact violations, so sampled conformations are physically plausible as well as probable.

Three inference modes are exposed. One-shot sampling draws an ensemble conditioned on a native structure and reports free energy surfaces, Ramachandran distributions, and secondary structure. Likelihood mode scores frames from an existing trajectory. Langevin mode treats the flow as a force field, propagating parallel walkers at 300 K with a 0.004 ps timestep under a potential combining the learned energy with bonded and repulsive priors. The implementation is PyTorch with PyTorch Lightning and MDTraj. The repository README states an MIT license but carries no license file, leaving reuse terms unsettled; the preprint is CC-BY.

#Applications

The unification matters most where sampling and simulation are both needed. A structural biologist can generate an equilibrium ensemble for a flexible protein in one pass, then continue the most interesting states with the same model as a coarse-grained force field rather than handing them to a separate simulation stack. Because the density is exact and cheap to evaluate, existing MD trajectories can be scored and reweighted, supporting free-energy estimation and the detection of undersampled regions. Teams hunting cryptic pockets, and biophysicists studying conformational transitions or allostery, gain a fast ensemble prior for proteins beyond the training set.

#Impact

UniFlow's contribution is architectural: it shows that the tractable-density property of normalizing flows, often treated as a theoretical convenience, is what lets one protein model do the work of two otherwise separate research programs. That argues for normalizing flows as a design choice in a subfield that has largely converged on diffusion and flow matching, where cheap likelihood access is harder to obtain. The caveats are real: the preprint has not been peer reviewed, the reported comparisons are in silico against reference MD and generative baselines, and no pretrained weights accompany the code, so adopters must train on their own trajectory data first.