Geometry Is Not Enough: Symplectic Structure and Noise Robustness in JEPA World Models for Beam Dynamics

Joint-Embedding Predictive Architectures (JEPAs) learn dynamics in latent space and promise cheaper, noise-robust learning-based surrogates for physical simulation. Without a reconstruction loss anchoring the representation, however, the encoder and predictor are free to settle on latents that are mutually consistent with little knowledge about the system dynamics. But whether enforcing that structure in latent space improves performance remains poorly understood. We address this by studying JEPAs on a system whose latent dynamics have a known target: longitudinal beam dynamics, whose one-turn map is a symplectic Hamiltonian flow, and use that structure as a reference for the learned latent. Our contributions are threefold: we examine whether enforcing symplectic structure on the latent predictor restores the geometry of the true one-turn map and improves accuracy, benchmark JEPA against reconstruction-based and pixel-space surrogates in error and cost, and assess its robustness to noise and distribution shift. On $23,101$ BLonD simulations of the FNAL Recycler, we find that symplectic predictors restore the geometry of the one-turn map but only direct map parameterisations improve accuracy, and that JEPA trains $2.1\times$–$2.5\times$ faster, and filters input noise, though reconstruction retains the lowest error. Our findings align with prior work on transfer-operator and latent Hamiltonian learning and point to encoder objectives and training regimes that would make JEPA a reliable surrogate for complex physical dynamics.