Behavioral Monitoring of JEPA World Models with Jacobian Centroids
cs.LG
Submitted: 2026-09-27
Updated: 2026-09-27
Comments: Presented at Workshop on World Models - Hosted by Chicago Booth
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
Abstract
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.
Sources
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
- Foundational World Models Accurately Detect Bimanual Manipulator Failures
- A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification
- The Linear Centroids Hypothesis: Features as Directions Learned by Local Experts
- DeepMind Control Suite
- Temporal Straightening for Latent Planning
- D4RL: Datasets for Deep Data-Driven Reinforcement Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks