JEPA · 2026-09-30
Behavioral Monitoring of JEPA World Models with Jacobian Centroids
Thomas Walker, Randall Balestriero, Richard Baraniuk
Auto-summarized: this summary was generated by a language model from the paper’s text and has not been reviewed by an editor. Check the paper before relying on it. Benchmark numbers appear only after a human has verified them.
TL;DR
The authors propose behavioral monitoring signals for JEPA world models using centroids defined as sub-component Jacobian row-sums. Centroids can be computed via Jacobian vector products, characterize internal input-space geometry, and yield saliency maps. On Push-T and TwoRoom under distribution shift, a structural dissociation (encoder goal represented, predictor unresponsive) predicts planning failure before actions; centroid-based methods outperform activation- and reconstruction-based shift detectors, enabling pre-execution goal resampling.
Why it matters
Runtime monitoring for world-model planning needs signals that reflect behavioral alignment, not just representational mismatch. The authors show that centroid-based behavioral signals (encoder/predictor alignment and goal-code cosine) identify a specific dissociation that predicts failure, and can gate planning before any action is taken. Under distribution shift, centroid methods outperform activation and reconstruction baselines for failure prediction, and centroid-based gating improves out-of-distribution success via goal resampling.
Method
- Compute per-sub-component centroids as row-sums of input-output Jacobians (centroid c = (Jg[x])⊤ 1), efficiently via Jacobian vector products; compare with layerwise activations extracted from MLP blocks.
- Construct failure/distribution-shift monitoring signals from centroid alignment (e.g., post-execution calibration ccen and goal-code cosines via OMP or top-k sparse autoencoder codes).
- Use a pre-execution gate based on absolute predictor-centroid cosine thresholds to flag dissociated episodes and apply adaptive replanning via goal resampling.
Limitation
Generalization to other JEPA families, such as DINO-WM, and to non-JEPA architectures remains open.
Abstract (from arXiv)
Detecting failures in World Model (WM)-based planning requires monitoring whether the model is behaviorally aligned with the current task, which in turn requires studying its internal representations. Here, we show that centroids---sub-component Jacobian row-sums---effectively identify the behavioral properties of WMs, complementing traditional activation-based knowledge signals. The centroids of a model are easily computed through Jacobian vector products and characterize how the model organizes the geometry of its input space, yielding an efficient perspective on internal representations, including the generation of task-relevant saliency maps. Evaluated on continuous control tasks using JEPA WMs, this behavioral view reveals a structural dissociation, where the encoder correctly represents the goal while the predictor remains behaviorally unresponsive. This failure mode directly predicts planning failure before any action is taken, allowing for goal resampling to recapture out-of-distribution success. Moreover, centroid-based methods outperform baseline methods as distribution-shift detectors. Together, these tools yield a behavioral monitoring stack that is operational and consequential under distribution shifts.
Related papers
- LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
- Anisotropic Representations Improve Planning in JEPA World Models
- Bilinear World Models: Learning Representations with Structured Dynamics for Efficient Control
- FlexiWorld: Learning and Planning via Flexible Action Chunks Across Multiple Time Scales
- Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State