Atomistic representation learning / Molecular dynamics
MACE JEPA: representations for molecular dynamics
An ongoing study of pretrained MACE features, explicit atomic velocities, and JEPA-style temporal learning, using atomistic water to test what representations retain about molecular dynamics.
My role. Encoder comparison, representation and adapter development, controlled training, and physical-readout evaluation.
Research question
Can temporal representation learning make pretrained atomistic features more useful for physical prediction, beyond the effects of the input interface and physical supervision?
Approach
- Audited encoder inputs and intermediate features to understand which geometric and dynamical information each representation exposes.
- Built an adapter that combines MACE-MH-1 scalar and vector features with atomic velocities and local geometry. Compared frozen and fine-tuned encoders under direct physical supervision, same-state representation pretraining, and future-state representation pretraining.
- Used common physical readouts and true-versus-shuffled temporal-pairing controls to examine the contribution of temporal learning, alongside short-horizon prediction and recursive-rollout diagnostics.
- Extended the investigation to structure-dependent hydrogen-bond exchange propensity: predict the event frequency averaged over repeated velocity initializations, comparing geometric inputs with the same inputs augmented by frozen MACE features.
Work to date
- Completed controlled training and evaluation pilots on atomistic water. In the initial short-horizon pilot, frozen MACE with direct physical supervision gave the lowest final physical errors; temporal pretraining showed different effects for oxygen and hydrogen.
- Follow-up temporal-pairing controls found task-dependent value in temporal learning. The benefit depended on the physical readout and supervision protocol, and reliable long-horizon dynamics remain unestablished.
- A separate frozen-feature study reduced hydrogen-bond exchange propensity error relative to its specified geometric baseline on two evaluation trajectory sources. This tests the usefulness of MACE features; it is not a result of JEPA training.
This line grew out of the encoder comparison within my CG JEPA investigation. It now asks which molecular targets benefit from pretrained representations and temporal learning, with hydrogen-bond exchange providing a target beyond immediate displacement.
Scope & limitations
- These are exploratory atomistic-water studies. They do not establish a production molecular-dynamics surrogate, broad materials transfer, or a general advantage for JEPA.
- The initial six-condition training pilot used one shared seed. Later technical-seed repetitions do not create independent physical systems.
- The propensity comparison adds feature dimensions as well as MACE information. Its evaluation sources share simulation ancestry and their labels had been inspected previously; the comparison is not a fully blind test or a capacity-matched causal attribution.