geom2vec

Chignolin looks like it folds into a single native structure, but its folded ensemble splits into four states that differ only in the χ1 rotamers of two threonine side chains. An analysis built on backbone dihedrals or distances between Cα atoms — the features molecular dynamics practitioners reach for first — cannot see that distinction at all. Picking features by hand is the standing bottleneck in interpreting simulations, and the choice quietly decides which mechanism the analysis reports.

geom2vec, from the University of Chicago with a collaborator at New York University's Courant Institute, takes the choice away. Equivariant graph neural networks are pretrained once on a general corpus of molecular conformations with a self-supervised denoising objective — random Gaussian displacements are added to Cartesian coordinates and the network learns to predict them — then frozen and used as universal geometric featurizers. A simulation frame goes in as raw coordinates; per-atom scalar (rotation-invariant) and vector (rotation-equivariant) features come out.

The consequence is structural, not merely convenient. Earlier work coupling GNNs to variational dynamics objectives had to backpropagate through the graph network on every batch, which forced small graphs — Cα-only representations, or very small molecules — because the batch sizes those objectives need do not fit in GPU memory alongside a trainable encoder. Decoupling the encoder from the downstream head means the expensive part runs once, in inference mode, and only a lightweight head is trained per simulation system. Encoding every non-hydrogen atom of a small protein becomes routine.

#Key Features

• Frozen transferable encoder: the pretrained GNN parameters are never fine-tuned for a downstream analysis, so one set of checkpoints serves any trajectory and any task built on top of the features.
• Residue-level tokens: per-atom representations are summed within user-defined structural units — amino acids for proteins — cutting node count by roughly an order of magnitude while keeping all-atom information inside each token.
• Expressive token mixers: instead of pooling tokens into one vector, a transformer (SubFormer) or MLP-Mixer (SubMixer) head learns interactions between tokens, optionally with geometric vector perceptron layers and a global token encoding pairwise Cα distances.
• Interpretable attention: SubFormer attention maps show which residues a learned state depends on, and unlike graph-attention baselines confined to one-hop neighbours they resolve long-range, blockwise patterns across the chain.
• Feasible on one GPU: precomputing features turns an analysis that would otherwise need hundreds of gigabytes of GPU memory into one that trains in minutes on a single card.

#Technical Details

The encoders are ViSNet and TorchMD-ET equivariant graph transformers, trained with 6 message-passing layers, hidden dimensions of 64 to 384, 64 radial basis functions, cutoffs of 5 or 7.5 Å, and a 0.2 Å noise level. Pretraining uses OrbNet Denali — 215,000 molecules and complexes averaging 45 atoms, with 2.3 million conformations — holding out 10,000 conformations for validation; a single pretraining run takes hours to days depending on architecture and card. The released checkpoints span ViSNet, TorchMD-ET, and TensorNet backbones.

Downstream evaluation covers VAMPnets and the state predictive information bottleneck on three fast-folding proteins from long unbiased simulations: chignolin (10 residues, 106 µs), trp-cage (20 residues, 208 µs), and villin HP35 (35 residues, 125 µs), split by time rather than at random to avoid the inflated validation scores that correlated frames produce. Token mixers beat simple pooling on trp-cage and villin, and adding geometric vector perceptrons with a global token lets them beat a distance-based MLP baseline. SPIB recovers 13 metastable trp-cage states at a 20 ns lag time. On cost: training a comparable VAMPnet with an unfrozen TorchMD-ET encoder is estimated at roughly 106 GB of GPU memory for trp-cage and 219 GB for villin, against hundreds of seconds on a single A40 with cached features.

#Applications

The immediate users are simulation groups who need collective variables, metastable state decompositions, or Markov state models from long trajectories without hand-curating input features per system. Because the featurizer is task-agnostic, the same cached embeddings support committor estimation, dimensionality reduction, and other learnable analyses, and the released package wraps trajectory inference through MDAnalysis so features can be computed over a folder of trajectories and reused.

#Impact

geom2vec's contribution is the demonstration that generic denoising pretraining on small organic molecules yields representations good enough for protein conformational analysis, with no protein-specific pretraining and no fine-tuning. That makes equivariant all-atom analysis affordable on hardware most groups already have, and it separates a slow, reusable encoding step from a fast, disposable analysis step. The limits are stated plainly by the authors: the demonstrations are three small fast-folding proteins from one simulation source, generalization to other classes of molecules and other tasks is expected but not yet quantified, and the downstream head still has to be trained afresh for each system studied.