UBio-MolFM-V1.5: Universal Bio-Molecular Foundation Model

UBio-MolFM is a foundation-model suite for molecular modeling, designed for bio-systems. This release, UBio-MolFM-V1.5 (Stage 3), is built on the E2Former-V2 linear-scaling equivariant transformer and is the checkpoint used for every simulation reported in UBio-MolFM: Enabling Biomolecular Dynamics at DFT Accuracy and 10⁵ Atoms with One Untuned Potential.

One potential, no per-system tuning: the same weights run pure water, aqueous electrolytes, an RNA pseudoknot with Mg²⁺, a cyclic-peptide free-energy landscape, and a 109,000-atom transmembrane K⁺ channel — all atoms on the ML potential, water included.


Technical report	`MolFM-1p5-Technical-Report.pdf` — in this repo
Code	https://github.com/IQuestLab/UBio-MolFM
Weights (this repo)	https://huggingface.co/IQuestLab/IQuest-UBio-MolFM-V1.5
Previous release	https://huggingface.co/IQuestLab/IQuest-UBio-MolFM-V1 (arXiv:2602.17709)
Training data (subset)	https://huggingface.co/datasets/IQuestLab/UBio-Protein26
Architecture	E2Former-V2 (arXiv:2601.16622), long-range message passing (arXiv:2601.03774)

What is new since V1

1. More training data — UBio-Mol26 grew from 17M to ≈19M labels. Structures are labelled at three DFT fidelities and paired with OMol25 for ≈160M quantum-chemical labels in total. The added configurations extend coverage of bio-relevant chemistry (nucleic acids, lipids, ions, protein environments) rather than simply scaling the existing distribution.

2. The condensed-phase density offset is fixed, by a periodic data tier. Earlier checkpoints were trained on non-periodic (cluster) data almost exclusively and over-bound bulk liquids: a 512-molecule NPT water box settles at 1.118 g/cm³ under the stage-2 checkpoint. UBio-Mol26's periodic tier enters at stage 3 and corrects the bias — the released model holds 0.987 ± 0.009 g/cm³ over 1 ns at 300 K / 1 bar, 1.0 % below the experimental 0.997, with the 0.15 mol/L NaCl and KCl boxes inside the same margin. No density is supplied in training, so the number is emergent — and discriminating: of the pretrained baselines, DPA-4 (its OMol head, which is what bounds this result rather than DPA-4 as a whole) collapses toward a gas-like density while MACE-OMol and UMA-S-1p2 over-densify. This is what makes unbiased condensed-phase MD usable at all; every simulation below depends on it.

3. Redesigned data mixture and stage-3 recipe. The curriculum runs three stages over one shared backbone: S1 OMol25 with separate energy and force heads; S2 still OMol25 but with the force head retired in favour of autograd, F = −∇ᵣE; S3 the first mixed continued pretraining on OMol25 and UBio-Mol26 together, over four re-balanced data branches (plus a low-weight auxiliary svp energy head). Effect on the shipped model: stage 3 cuts the best pretrained baseline's force error by 47 / 47 / 21 / 64 % on held-out TZVPD protein / DNA / RNA / lipid fragments and by 32–51 % across the extreme-size TZVP classes — and, the property that matters for MD, trajectories that stay stationary over nanoseconds instead of drifting.

4. Extensive MD validation, not just single-point error. V1 was evaluated primarily on held-out DFT energies and forces. V1.5 is evaluated by running the simulations, against experiment where experiment exists and against classical force fields where it does not. See Validation. The technical report is the full account; the model card quotes its headline numbers.

Model details


Model type	Equivariant transformer (E2Former-V2), linear-scaling tensor products
Training stage	Stage 3 (final stage of the curriculum)
Parameters	≈23.5 M (the `.pt` also carries optimizer state, hence its size)
Prediction heads	`omol25` (default, used for every reported result) and `svp`
Receptive field	Hierarchical hybrid cutoff: three short-range all-atom blocks (5.0 Å) feed a medium-range block (8.0 Å) over heavy-atom neighbourhoods through a residual skip, both merging at a fusion node. Node-level SO(2) transformations replace dense tensor products, so equivariance stays linear-scaling.
Node irreps	`256x0e+256x1e+256x2e`
Outputs	Total energy and per-atom forces (forces via autograd, so they are conservative)
Scale	High fidelity up to ~1,500 atoms benchmarked against DFT; simulations up to ~10⁵ atoms on a single GPU (memory-dependent)
Boundary conditions	Periodic and non-periodic

Files

File	Description
`MolFM-1p5-Technical-Report.pdf`	Technical report — methods, benchmarks and all four simulation campaigns in full.
`molfm-v1p5-stage-3.pt`	Pretrained checkpoint.
`config.yaml`	Flattened inference configuration. Must accompany the checkpoint — see the note below.
`img/`	Figures used in this card.
`reproduce-data/`	4.5 GB companion deposit: source data behind every figure and table, starting structures and build scripts, per-frame observables, analysis scripts, and 4.4 GB of stride-decimated trajectories with forces. Start with its `REPRODUCIBILITY_CHECKLIST.md`.

On the config file. Training goes through Hydra and composes config_file/config_molfm.yaml from several fragments; inference loads a single flattened yaml. config.yaml here is the flattened form that matches these weights. Passing the Hydra root to the inference interface instead fails with TypeError: SmallMolConfig.__init__() got an unexpected keyword argument 'defaults'. Keep the two files together: load_state_dict runs with strict=False, so a mismatched backbone block silently leaves parameters randomly initialized rather than raising.

Usage

Install the codebase first — see the repository for environment setup (Python 3.12, PyTorch 2.8.0, CUDA 12.8).

# Weights + config only (~265 MB). Drop the file list to also pull reproduce-data (4.5 GB).
hf download IQuestLab/IQuest-UBio-MolFM-V1.5 molfm-v1p5-stage-3.pt config.yaml \
    --local-dir ckpt/molfm-v1p5

By default the inference interface resolves the config from the checkpoint's directory, so pointing at the downloaded folder is enough.

Single-point energy and forces

from ase.build import molecule
from molfm.interface.ase.calculator.e2former_calculator import E2FormerCalculator

atoms = molecule("H2O")
atoms.set_cell([10, 10, 10])
atoms.pbc = [True, True, True]

calc = E2FormerCalculator(
    checkpoint_path="ckpt/molfm-v1p5/molfm-v1p5-stage-3.pt",
    config_name="ckpt/molfm-v1p5/config.yaml",
    head_name="omol25",
    device="cuda",
    use_tf32=True,
    use_compile=True,
)

atoms.calc = calc
print(f"Energy: {atoms.get_potential_energy()} eV")
print(f"Forces:\n{atoms.get_forces()}")

Molecular dynamics with ASE

from ase import units
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution

MaxwellBoltzmannDistribution(atoms, temperature_K=300)
dyn = Langevin(atoms, 1 * units.fs, temperature_K=300, friction=0.01)
dyn.run(100)

Production MD from the command line (recommended)

For anything longer than a smoke test, use the GPU-native simulation stack in the repository (src/molfm/interface/ase/torch_ext/), driven by interface/ase/cli.py. Positions, velocities, forces and the integrator all stay on the GPU for the whole trajectory — no per-step host round-trip — trajectories are streamed to HDF5 by a background writer, and runs resume from the last saved frame automatically. This is the stack every result below was produced with.

# 216-water box shipped with the repo: 100 ps NVT (200,000 x 0.5 fs) at 300 K
python src/molfm/interface/ase/cli.py \
    --input_path tools/samples/water216_nvt.xyz \
    --checkpoint ckpt/molfm-v1p5/molfm-v1p5-stage-3.pt \
    --config ckpt/molfm-v1p5/config.yaml \
    --head_name omol25 \
    --task md --steps 200000 --temp 300 --dt 0.5 --ensemble nvt \
    --full_sync_interval 10 --log_interval 500 --seed 42 \
    --use_tf32 True --use_compile True --save_in_fp32 True \
    --device cuda:0 --work_dir runs --name water216_nvt_100ps

--task relax runs L-BFGS instead; --ensemble npt adds an energy-only Monte-Carlo barostat; --thermostat qtb switches to a Quantum Thermal Bath for nuclear quantum effects. python src/molfm/interface/ase/cli.py --help lists everything.

Performance notes

use_tf32=True enables TensorFloat-32 on supported NVIDIA GPUs — higher throughput, a small precision cost.
use_compile=True enables torch.compile; it usually speeds things up and can reduce memory.
recompute_budget (0.0–1.0, torch.compile only) caps the fraction of activations kept in memory and recomputes the rest. This is how the largest systems are made to fit — e.g. --recompute_budget 0.3 for the 10⁵-atom range.
The first run is slower: kernels and compilation artifacts have to be built.

Validation: what the model was actually run on

Every number below comes from the technical report and is recomputable from reproduce-data/; the file backing each claim is named. All simulations use this single checkpoint with no system-specific tuning, and put all atoms — solvent included — on the ML potential.

Force accuracy holds past a thousand atoms

Force MAE (meV/Å) against the best pretrained baseline, on held-out UBio-Mol26 TZVPD fragments (390–909 atoms). Table 1 of the report has all three tiers, five models and the energy metrics; source_data/table1_accuracy/ has the per-class files.

Model	Protein	DNA	RNA	Lipid
UMA-S-1p2 (best baseline)	40.8	38.9	59.6	33.0
UBio-MolFM S2 (OMol25 only)	44.4	41.7	65.8	34.0
UBio-MolFM S3 (released)	21.8	20.7	47.3	11.9

The ranking inverts at biomolecular size and holds there: S3 leads every baseline on every class of the TZVPD fragments and on all five categories of the extreme-size TZVP trajectories (1,215–1,555 atoms). On OMol-Bio-10k — the baselines' own training distribution — UMA-S-1p2 is still strongest, and the report does not contest that tier.

Water and electrolytes — agreement with experiment

512-molecule water box and 0.15 mol/L NaCl / KCl boxes (12,121 atoms), 1 ns each at dt 0.5 fs, NVT / NPT / NVE. source_data/fig2_thermo/

Density 0.987 ± 0.009 g/cm³ (water, 1.0 % below the experimental 0.997), 0.995 (NaCl), 0.996 (KCl) — against 1.118 g/cm³ for the stage-2 checkpoint, i.e. the density offset is gone.
Self-diffusion D_PBC = 1.60 × 10⁻⁵ cm² s⁻¹; after the Yeh–Hummer finite-size correction D_∞ = 1.88 × 10⁻⁵ cm² s⁻¹, ~18 % below experiment (and that 18 % is a lower bound — the correction uses the experimental viscosity; see REPRODUCIBILITY_CHECKLIST.md §G.3).
Structure O–O, ion–water RDFs and coordination numbers consistent with neutron/X-ray reference data.
NVE total energy stationary over 1 ns at 0.5 fs, which is the conservativeness check the autograd forces are supposed to pass.

RNA pseudoknot with Mg²⁺ — the ion behaves like an ion

BWYV pseudoknot (PDB 1L2X), 22,804 atoms, 5 × 1 ns unbiased NPT, against AMBER OL3 + Li–Merz 12-6-4. per_frame_observables/rna/comparison/

The Mg²⁺ inner-sphere phosphate contact holds in 100 % of frames across all ten trajectories, and the complete six-ligand octahedral shell in 99.79 % of ours — obtained without any ion-specific correction term, which the classical baseline achieves only with an explicit C4 polarization term.

Cyclosporine A — a free-energy landscape a fixed-charge field cannot produce

80 umbrella windows × 400.8 ps, one global 2D histogram-MBAR, versus an automatically-parameterised GAFF2 / AM1-BCC baseline on the identical box. source_data/fig4_csa/, inputs/csa/us_experiment_results/cv_samples/

F(CV₁, HB₁) from the same window layout under each potential, contoured every 1 kcal/mol; hatched cells carry no estimate. Left, UBio-MolFM: one connected basin with ordered contours — a funnel, and a landscape spanning 4.97 kcal/mol. Right, classical fixed charge: a mottled surface of shallow local minima spanning 1.14 kcal/mol, with no organised basin — the latch is flattened away. The 5.51 / 1.51 kcal/mol below are the funnel depth along the well-sampled window; the numbers on the panels are the span of the full 2D landscape.

UBio-MolFM resolves a continuous, connected closed → open path with a well-defined closed basin and a real barrier: a 5.51 kcal/mol funnel across the well-sampled window, putting a 3.54 kcal/mol toll (window-bootstrap range 2.69–4.41) on the membrane-permeable conformer. The classical baseline is essentially flat — its entire surface over the same window fits inside 1.51 kcal/mol, with no funnel at all. The switch is gated by a single kinetically asymmetric hydrogen bond (unlatch / relatch 0.43 / 3.73 kcal/mol, against 0.70 / 1.20 classically): it is driven by intramolecular H-bond rearrangement and the charge redistribution that comes with it, which a fixed-charge model has no mechanism to represent. The model does, and it shows up as structure in the landscape rather than as noise.

The classical campaign was run with concessions in its own favour (2 fs with H-bond constraints, first 30 % of each window discarded); a 0.5 fs unconstrained 80-state re-run is also deposited, and it is not uniformly favourable to us — the closed-basin depth moves against us under it.

KcsA channel — 10⁵ atoms, 1 ns, and the selectivity filter holds

Full KcsA tetramer in a POPE:POPG 3:1 Lipid21 bilayer, 108,964 atoms, 5 K⁺ loaded in the pore, 5 × 1 ns at dt 0.5 fs / 310 K, against two classical potentials × 5 replicas each (Li–Merz 12-6-4 and 12-6 + Joung–Cheatham). source_data/fig5_kcsa/, per_frame_observables/kcsa/

	UBio-MolFM (5)	12-6-4 (5)	12-6 (5)
Contiguous direct-contact K⁺ column	78–89 % of late frames in 4 of 5 replicas	0 % in all 5	0 % in all 5
Closest K⁺–K⁺ spacing (5-replica mean)	3.32 ± 0.13 Å (direct-contact guide: 3.4 Å)	3.74 ± 0.10 Å	3.71 ± 0.18 Å
Frames with water in the ion column	0 % in all 5	0–86 %	0–100 %
Cavity K⁺ → nearest carbonyl O	2.81 Å (coord. 3.4)	4.11 Å (coord. 0.8)	4.43 Å (coord. 0.4)
S4 K⁺ → carbonyl O / coordination	2.56 Å / 8.0	2.62 Å / 7.9	2.59 Å / 7.9

The classical replicas relax to water-mediated spacings and let water into the filter; ours keep the ions anhydrous in all five and in direct carbonyl contact in four of five, and in the fifth until 700 ps — the knock-on geometry proposed for permeation, and the coordination chemistry the crystallography implies. The 12-6-4 replicas also drift outward significantly (+0.29 to +0.54 Å/ns on the closest K⁺–K⁺ pair), while the four UBio-MolFM replicas that hold the column carry no significant trend (+0.001 to +0.046 Å/ns).

One replica per potential, same equilibrated 108,964-atom system, 1 ns each. Per-site K⁺→nearest-carbonyl-O distances (top block), adjacent K⁺–K⁺ spacings against the 3.4 Å direct / 5.4 Å water-mediated guides (middle), and on-axis occupancy (bottom). Top, UBio-MolFM — deliberately seed 126, our least favourable replica: it holds four ions on axis for 716 ps, every occupied site sits 2.53–2.60 Å from its nearest carbonyl O (the cavity ion included, at 2.60 Å), and the adjacent spacings track the direct-contact guide (S₄–S₃ 3.28 Å, S_cav–S₄ 3.32 Å). Bottom, Amber Lipid21 + 12-6 ions: down to three on-axis ions within 22 ps, its cavity ion parked 5.73 Å from the carbonyls — a water layer below them — and S_cav–S₄ at 7.30 Å, past even the water-mediated guide. The two deepest sites agree with ours (2.59 / 2.68 Å); the divergence is at S2 and the cavity.

That last point is the shape of the whole result: the potentials agree on the deep sites S3/S4 and split at S2, the Gly77–Tyr78 cage — at 10⁵ atoms the difference traces to one backbone carbonyl's orientation. Seed 126, the replica in the figure, is the exception among ours and is reported as one: 28.3 % contiguous contact and occupancy ≥ 4 in 72.9 % of frames, against 78–89 % and 100 % for the other four. It is included in every statistic above, the 3.32 Å mean included.

On cost: at this system size the pretrained baselines either exhaust memory or fall an order of magnitude short on throughput.

Known limitations

Read reproduce-data/REPRODUCIBILITY_CHECKLIST.md §G before drawing conclusions from any of this — it lists eleven limitations that bound what a reproduction can establish, including the water-viscosity caveat above, the isotropic-barostat constraint on the KcsA runs (with its released-box control), and the fact that the released 5M UBio-Protein26 subset is not sufficient to retrain this model.

Training data

Three-stage curriculum learning over ≈160M quantum-chemical labels:

UBio-Mol26 — ≈19M labels built for this work, at three DFT fidelities (the third periodic). It pairs bottom-up enumeration of explicitly solvated building blocks (amino acids, nucleotides, lipid segments) with top-down, environment-aware sampling of local regions excised from AlphaFold proteins together with their surrounding water, reaching 1,370 atoms at a mean of ~414. A 5M protein-focused subset is released as UBio-Protein26 — enough to study the distribution, not enough to retrain this model.
OMol25 — ≈140M labels; the pretraining corpus for stages 1–2 and one branch of the stage-3 mixture.

Citation

If you use UBio-MolFM-V1.5, please cite the technical report (MolFM-1p5-Technical-Report.pdf in this repository):

@techreport{huang2026ubiomolfm1p5,
      title={UBio-MolFM: Enabling Biomolecular Dynamics at DFT Accuracy and 10^5 Atoms with One Untuned Potential},
      author={Lin Huang and Frank Peng and JiaJun Cheng and Zion Wang and Hao Yin and Hao Li and Ji Zhang and Jack Jia and Junping Zhao and Arthur Jiang and Jia Zhang},
      institution={IQuest Research, UBio Team},
      year={2026},
      note={Technical report},
      url={https://huggingface.co/IQuestLab/IQuest-UBio-MolFM-V1.5},
}

@misc{huang2026ubiomolfm,
      title={UBio-MolFM: A Universal Molecular Foundation Model for Bio-Systems},
      author={Lin Huang and Arthur Jiang and XiaoLi Liu and Zion Wang and Jason Zhao and Chu Wang and HaoCheng Lu and ChengXiang Huang and JiaJun Cheng and YiYue Du and Jia Zhang},
      year={2026},
      eprint={2602.17709},
      archivePrefix={arXiv},
      primaryClass={physics.chem-ph}
}

@misc{huang2026e2formerv2,
      title={E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory},
      author={Lin Huang and Chengxiang Huang and Ziang Wang and Yiyue Du and Chu Wang and Haocheng Lu and Yunyang Li and Xiaoli Liu and Arthur Jiang and Jia Zhang},
      year={2026},
      eprint={2601.16622},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@misc{wang2026scalable,
      title={Scalable Machine Learning Force Fields for Macromolecular Systems Through Long-Range Aware Message Passing},
      author={Chu Wang and Lin Huang and Xinran Wei and Tao Qin and Arthur Jiang and Lixue Cheng and Jia Zhang},
      year={2026},
      eprint={2601.03774},
      archivePrefix={arXiv},
      primaryClass={physics.chem-ph}
}