🤖 machine learning

Phys-JEPA: Physics-Informed Latent World Models for Multivariate Time-Series Forecasting

Phys-JEPA introduces a physics-informed joint-embedding predictive architecture that enforces physical consistency directly within latent state transitions rather than just at the output level, thereby improving multivariate time-series forecasting accuracy and interpretability on real-world datasets like Jena Climate, Traffic, and Electricity.

Original authors: Weizhi Nie, Weichao Liu, Honglin Guo, Yuting Su

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Weizhi Nie, Weichao Liu, Honglin Guo, Yuting Su

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to predict the future weather or traffic patterns.

The Old Way: The "Guess-and-Check" Student
Most current AI models are like students who are only graded on their final test answers. They look at the past (yesterday's weather), try to guess the future (tomorrow's weather), and if the guess is wrong, they just tweak their answer until it looks right.

The Problem: The robot might get the final number right, but it doesn't actually understand the physics. It's just memorizing patterns. If you ask it to predict a week from now, it might start making wild guesses because its internal "brain" (the hidden state) isn't organized by real-world rules like temperature or pressure. It's like a student who memorized the answer key but doesn't know the math.

The New Way: Phys-JEPA (The "Physics-Savvy" Student)
The authors of this paper, from Tianjin University, built a new model called Phys-JEPA. Instead of just checking the final answer, they force the robot to organize its internal brain using the actual laws of physics before it even makes a prediction.

Here is how it works, using a simple analogy:

1. The Two-Part Brain

Imagine the robot's brain is split into two drawers:

Drawer A (The Physics Drawer): This holds the "real" stuff we understand, like temperature, pressure, and wind speed. The robot is forced to put these specific, known facts here.
Drawer B (The Mystery Drawer): This holds the "residual" stuff—the messy, complicated parts that are hard to explain with simple physics (like sudden traffic jams or weird humidity spikes).

2. The "Internal Check" (The Magic Trick)

In the old models, the robot only checked if its final prediction matched reality.
In Phys-JEPA, the robot has to pass two tests inside its brain before it ever writes down a prediction:

The Snapshot Test: Does the "Physics Drawer" look like the real weather right now? (e.g., If the real temperature is 20°C, the robot's internal physics number must also represent 20°C).
The Movement Test: If the robot predicts the weather will change, does the change in its Physics Drawer match the real change? (e.g., If the real temperature drops by 5 degrees, the robot's internal number must drop by 5 degrees, not just randomly change).

3. Why This Matters

The paper tested this on three real-world datasets:

Jena Climate (Weather): This is the "purest" test because weather follows strict physical laws. Here, Phys-JEPA worked great. It predicted temperature and pressure more accurately than the old models because its internal brain was actually tracking the physics.
Traffic & Electricity: These are "messier" systems. Traffic isn't governed by a single physics equation like gravity; it's a mix of human behavior and road rules.
- On Traffic, the new model was the clear winner. By forcing the brain to respect the "physics" of traffic flow (even if that physics is a bit fuzzy), it predicted future congestion much better.
- On Electricity, it was a mix. For short-term predictions, just checking the "Snapshot" (current state) was best. For very long-term predictions (192 hours ahead), the full "Movement" check (tracking how things change over time) became the most helpful.

The Big Takeaway

The paper claims that you get better predictions if you force the AI to understand the rules of the world inside its thinking process, rather than just fixing its mistakes at the very end.

Think of it like building a house:

Old Model: Build a wobbly tower, then tape a sign on it that says "This is a stable house."
Phys-JEPA: Build the foundation and walls using the actual laws of physics first, so the house is stable by design.

The authors admit this is just the beginning. They found it works best when the "rules" (the physics) are clear, like in weather. When the rules are vague (like traffic), it still helps, but you have to be careful not to force the rules too hard. They suggest the next step is to test this on even stricter physical systems to see how far this "physics-first" thinking can go.

Technical Summary: Phys-JEPA

Problem Statement

Multivariate time-series forecasting in physical systems (e.g., energy grids, transportation, meteorology) requires models that predict coupled temporal variables while preserving meaningful state evolution. While deep forecasting models excel at capturing complex statistical correlations, their hidden predictive states are often optimized solely for accuracy, resulting in latent dynamics that are difficult to interpret and weakly constrained by physical laws.

Existing physics-informed approaches typically impose constraints only at the decoded-output level (e.g., penalizing predictions that violate conservation laws). The authors argue this is insufficient for long-horizon forecasting because it leaves the internal latent predictive state—the source from which future trajectories are generated—physically unstructured. Consequently, a model may produce plausible short-term forecasts while learning unstable or non-physical internal transitions that accumulate error over time.

Methodology: Phys-JEPA

The paper proposes Phys-JEPA (Physics-Informed Joint-Embedding Predictive Architecture), a framework that moves physical regularization from the output space into the latent predictive state space.

Core Architecture

Phys-JEPA extends the Joint-Embedding Predictive Architecture (JEPA) paradigm, which learns to predict future representations rather than reconstructing raw observations. The model consists of:

Context Encoder ( $E_c$ ): Maps a historical time window to a current latent state $z_t$ .
Latent Predictor ( $P$ ): Forecasts the future latent state $\hat{z}_{t+H}$ directly in the representation space.
Decoder ( $D$ ): Maps the predicted latent state to the future trajectory $\hat{x}_{t+1:t+H}$ .
Target Encoder ( $E_y$ ): Maps the actual future window to a target representation $z^*_{t+H}$ (used with stop-gradient for JEPA alignment).

Physics-Informed Latent Decomposition

The central innovation is the decomposition of the latent state $z_t$ into two distinct components:

Physical Component ( $z^{phy}_t$ ): Encodes variables with direct physical meaning.
Residual Component ( $z^{res}_t$ ): Captures unresolved factors and statistical variations not explicitly covered by the physical projection.

A lightweight Physics Projector ( $G$ ) maps the physical latent component to the space of observed physical variables ( $q_t$ ).

Loss Functions

The training objective combines standard forecasting losses with specific constraints on the latent physical branch:

JEPA Objective ( $L_{JEPA}$ ): Aligns the predicted latent state with the target latent representation to learn predictive features without pixel-level reconstruction.
Forecasting Loss ( $L_{forecast}$ ): Standard MSE between decoded predictions and ground truth.
Static Latent Physical State Loss ( $L_{state}$ ): Ensures the physical latent state corresponds to observed physical variables at both the current and future time steps:
$L_{state} = \|G(z^{phy}_t) - q_t\|^2_2 + \|G(\hat{z}^{phy}_{t+H}) - q_{t+H}\|^2_2$
Dynamic Latent Physical Transition Loss ( $L_{dyn}$ ): Ensures the change in the physical latent state matches the observed physical change:
$L_{dyn} = \|[G(\hat{z}^{phy}_{t+H}) - G(z^{phy}_t)] - [q_{t+H} - q_t]\|^2_2$
Residual Regularization ( $L_{SIGReg}$ ): Applied only to the residual component to prevent representation collapse (ensuring variance and covariance structure) without imposing isotropic priors on the structured physical component.

The total loss is a weighted sum: $L = L_{forecast} + \alpha L_{JEPA} + \beta L_{SIGReg} + \lambda_s L_{state} + \lambda_d L_{dyn}$ .

Key Contributions

Latent World Model for Physical Forecasting: The paper introduces a method where physical consistency is enforced on the predictive latent state and its transitions, rather than solely on decoded outputs. This organizes the internal world model to be physically interpretable.
Physical–Residual Decomposition: A novel architecture that separates latent states into physical and residual components, allowing for explicit physical supervision while preserving flexibility for non-explicit dynamics.
Dual Consistency Constraints: The formulation utilizes both static state alignment (anchoring states to variables) and dynamic transition alignment (anchoring state changes to physical laws), addressing the limitations of using either constraint in isolation.

Experimental Results

The authors evaluated Phys-JEPA on three benchmarks: Jena Climate (meteorological), Electricity (grid load), and Traffic (road occupancy).

Jena Climate (Strong Physical Descriptors):
- Phys-JEPA reduced aggregate MSE from 0.12482 to 0.12273 and temperature MSE from 0.01892 to 0.01831 at horizon $H=24$ .
- Improvements were observed across most thermodynamic and vapor-pressure variables, validating the approach when variables have direct meteorological meaning.
Traffic (Weak-Domain Descriptors):
- Full Phys-JEPA (Static + Dynamic) consistently outperformed the supervised baseline across all horizons ( $H=24$ to $H=192$ ).
- At $H=192$ , aggregate MSE improved from 0.800784 to 0.773873.
Electricity (Weak-Domain Descriptors):
- Results were horizon-dependent. Static-only consistency performed best at short horizons ( $H=24, 48$ ), while full Phys-JEPA achieved the best results at the longest horizon ( $H=192$ ).
- This suggests that dynamic transition consistency is beneficial for long-range evolution but can be restrictive if the physical descriptor is too coarse for short-term dynamics.
Ablation Insights:
- Output-level physics (penalizing decoded forecasts) was found to be weaker than latent regularization in these settings.
- Latent diagnostics (effective rank and covariance rank) confirmed that the physical losses reshape the latent representation geometry rather than acting merely as output penalties.

Significance and Claims

The paper claims that moving physics-informed learning from the output space to the latent predictive state space is a promising route toward interpretable and physically consistent temporal world models.

Interpretability: By decomposing latent states, the model learns representations where physical variables and their transitions are explicitly structured.
Stability: Constraining the latent transition helps prevent error accumulation in long-horizon forecasting, addressing a key weakness of purely statistical deep forecasters.
Limitations and Future Work: The authors modestly note that current results rely on single-seed experiments and "weak-domain" descriptors for Electricity and Traffic (which are not exact governing equations). They emphasize that the method's full potential should be evaluated on systems with stronger, explicit physical constraints where transition laws can be directly verified. The work is presented as an initial validation of the latent-physics formulation rather than a final benchmark against all modern forecasting architectures.