Representation Learning for Spatiotemporal Physical Systems

Abstract: Machine learning approaches to spatiotemporal physical systems have primarily focused on next-frame prediction, with the goal of learning an accurate emulator for the system's evolution in time. However, these emulators are computationally expensive to train and are subject to performance pitfalls, such as compounding errors during autoregressive rollout. In this work, we take a different perspective and look at scientific tasks further downstream of predicting the next frame, such as estimation of a system's governing physical parameters. Accuracy on these tasks offers a uniquely quantifiable glimpse into the physical relevance of the representations of these models. We evaluate the effectiveness of general-purpose self-supervised methods in learning physics-grounded representations that are useful for downstream scientific tasks. Surprisingly, we find that not all methods designed for physical modeling outperform generic self-supervised learning methods on these tasks, and methods that learn in the latent space (e.g., joint embedding predictive architectures, or JEPAs) outperform those optimizing pixel-level prediction objectives. Code is available at https://github.com/helenqu/physical-representation-learning.

• The paper shows that JEPA outperforms MAE and MPP in physical parameter inference, reducing error by 28-51% across various PDE systems.
• The paper benchmarks self-supervised strategies on systems like active matter, Rayleigh-Bénard convection, and shear flow to emphasize robust latent representations.
• The paper highlights that latent prediction objectives yield improved data efficiency and capture physically meaningful invariants for downstream scientific tasks.

Representation Learning for Spatiotemporal Physical Systems

Introduction

The investigation presented in "Representation Learning for Spatiotemporal Physical Systems" (2603.13227) systematically evaluates self-supervised representation learning methods for partial differential equation (PDE)-governed physical systems. The focus is on moving beyond next-frame prediction—an approach prevalent in scientific machine learning—and towards learning representations that are robustly informative for downstream scientific tasks such as physical parameter inference. The study benchmarks general-purpose self-supervised strategies, notably joint embedding predictive architectures (JEPA), masked autoencoding (MAE), and operator meta-learning approaches on canonical dynamical systems including active matter, Rayleigh-Bénard convection, and shear flow, intending to quantify and compare their physical relevance through parameter estimation accuracy.

Problem Framing and Physical Systems

Traditional surrogates in scientific ML typically aim to emulate costly numerical simulations at the pixel or field level. However, these approaches can suffer from compounding rollout errors and may not preserve global or physically meaningful invariants essential for inverse tasks such as system identification or scientific inference. By instead evaluating representations through the lens of physical parameter recovery, the study reliably quantifies whether the model's latent space encodes physical informativity.

The evaluated physical testbeds—active matter, Rayleigh-Bénard convection, and shear flow—are all governed by nontrivial PDEs and parameterized by physically interpretable quantities. For instance, Rayleigh-Bénard convection is governed by the Rayleigh ($\mathrm{Ra}$) and Prandtl ($\mathrm{Pr}$) numbers, which drive diverse qualitative dynamics as visualized below.

Figure 1: Example trajectories from the physical systems in our evaluation.

The chosen testbeds expose a range of dynamical features, including order–disorder transitions and coherent structure generation, that challenge the ability of learned representations to capture global system behavior, rather than surficial visual/textural detail.

Figure 2: Different values for the physical parameters, e.g., Rayleigh ($\mathrm{Ra}$) and Prandtl ($\mathrm{Pr}$) numbers for Rayleigh-Bénard, lead each system to evolve very differently.

Methodological Framework

The empirical analysis targets four model classes:

• Joint Embedding Predictive Architectures (JEPA): These models predict future representations in the latent space by aligning embeddings of contiguous frame blocks. The JEPA objective incorporates invariance, variance, and covariance regularization as per VICReg [bardes2021vicreg], and eschews pixel-level constraints.
• Masked Autoencoding (MAE): MAE is implemented as VideoMAE, learning spatial-temporal feature representations by reconstructing masked image tubes in pixel space.
• Autoregressive Surrogate Modeling (MPP): MPP is a next-frame pixel predictor, pre-trained as a foundation model across multiple PDE domains.
• Operator Meta-learning (DISCO): DISCO infers an operator per-trajectory via data prompts, merging neural operator priors and in-context learning [morel2025discolearningdiscoverevolution].

Self-supervised pretraining is performed individually for each physical dataset, except for models packaged as “foundation models.” Fine-tuning is performed only on the predictor head for non-autoregressive models to directly measure latent informativity and avoid collapse to the pretraining objective.

Results and Key Findings

JEPA Outperforms Pixel-based Methods

A central empirical conclusion is that JEPA substantially outperforms both masked autoencoder (VideoMAE) and autoregressive surrogate (MPP) approaches on physical parameter inference tasks across all evaluated systems. MSE is reduced by 51% on active matter, 43% on shear flow, and 28% on Rayleigh-Bénard compared to pixel-level MAE. Notably, JEPA's error approaches that of DISCO on two out of three systems, despite JEPA's general-purpose nature.

Furthermore, the performance gap is systematic: the two latent-space predictive methods (JEPA and DISCO) are consistently the top performers in each domain, whereas the pixel-level generative methods exhibit both lower accuracy and higher sensitivity to task/domain shift.

Sample Efficiency and Scaling Properties

JEPA also demonstrates improved data efficiency during downstream fine-tuning. With only 10% of available finetuning data, JEPA already matches the best case result achievable by VideoMAE with a full dataset. This scaling property suggests the inductive biases inherent in latent-predictive objectives preferentially encode globally informative, physically-meaningful features, offering clear practical advantages for scenarios with limited labeled data.

Latent Prediction vs. Pixel Reconstruction

The finding that generic self-supervised learning approaches (e.g., JEPA) can easily outperform heavily-engineered pixel-level physics surrogates (e.g., MPP) on parameter inference directly contradicts the assumption that pixel-fidelity in surrogate modeling leads to optimal representations for scientific tasks. The results parallel analogous findings in NLP, where encoder-only representation learning (e.g., BERT) often supersedes autoregressive models on discriminative tasks [devlin2018bert].

Theoretical and Practical Implications

These results have significant consequences for scientific machine learning and surrogate modeling:

• Foundation Model Development: The demonstrated outperformance of JEPA on physical inference indicates that the next generation of ML-based scientific surrogates should prioritize the optimization of physically resonant representations over mere field-level generative fidelity. Models designed solely for pixel-accurate rollouts may not encode the preferred structure for scientific inference.
• Representation Disentanglement: The distinction between generative rollouts and useful, physically T-informative embeddings formalizes a pathway to disentangle emulation fidelity from informativeness for downstream tasks. This has meaningful implications for both interpretability and scientific transfer learning.
• Sample Efficiency: Increased sample efficiency in JEPA points towards a reduced dependency on large-scale labeled data for fine-tuning, which is a dominant bottleneck in scientific domains.
• Generalization and Robustness: Encoding physical invariants and global system features in latent spaces is likely to improve robustness to domain shifts, parameter extrapolation, and compositional generalization, all of which are critical for scientific deployment.

Future Directions

The superior scaling and cross-task transfer of latent-space self-supervised approaches like JEPA motivate continued research in several directions:

• Hybrid Objectives: Combining latent and pixel-level predictive losses or embedding explicit physical inductive biases within JEPA-like frameworks.
• Scalable Pretraining and Transfer: Investigating the limits of transfer across unrelated PDE-governed systems, and into more complex multi-physics environments as found in The Well [ohana2025welllargescalecollectiondiverse].
• Limitations of Pixel Emulation: Further characterizing the regimes and scientific questions for which emulation fidelity is necessary versus cases where latent informativity suffices.
• Foundation Operator Models: Development of operator foundation models that leverage latent self-supervision to encode physically actionable representations, amenable to adaptation with minimal downstream data.

Conclusion

The study rigorously demonstrates that latent prediction objectives in self-supervised learning yield representations with superior physical informativeness and data efficiency for downstream scientific inference, compared to pixel reconstruction and autoregressive modeling. These empirical findings underscore the necessity of aligning representation learning paradigms with the demands of scientific parameter inference, rather than solely optimizing for surrogate rollout fidelity. This work introduces a robust, quantifiable, and transferable benchmark for physical representation learning and is expected to inform the design and evaluation of the next generation of AI models for scientific discovery and simulation (2603.13227).