MotionJEPA

Preventing Temporal Feature Collapse by Capturing Visual Changes in Latent Space

Markus Karmann^1,*, Shile Li^2,*, Christian Internò^3,1, Bruno Andreis^1,4, David Klindt⁵, Randall Balestriero^6,7, Jindong Gu¹, Philip Torr^1,4, Qi Zhang⁸, Peng-Tao Jiang⁸, Hao Zhang⁸, Bo Li⁸, Onay Urfalioglu²

¹University of Oxford, ²vivo Tech Research GmbH, ³Bielefeld University, ⁴Slater Labs, ⁵Cold Spring Harbor Laboratory, ⁶Brown University, ⁷AMI Labs, ⁸vivo BlueImage Lab, vivo Mobile Communication Co., Ltd., China

^*Equal contribution

Paper arXiv Code

Feature suppression. Visual probes of JEPA models trained on Pong. LeWM* focuses on slow features, reconstructing the scoreboard and bats but suppressing the fast-moving ball. SMWM* focuses on action-labeled features and encodes mainly the bats. MotionJEPA learns a more balanced representation, correctly predicting the score change and simulating the ball movement.

Abstract

Joint Embedding Predictive Architectures (JEPAs) are a promising paradigm for learning task-agnostic latent world models without visual reconstruction. However, standard JEPA training exhibits a strong inductive bias towards slow features, causing feature suppression and the collapse of latent representation. While inverse dynamics provides temporal anti-collapse, it relies on action labels and offers little incentive to embed general, unlabeled dynamics. We introduce Difference Image and Single image embedding Regularization (DISReg), a novel regularizer that builds on an inverse-dynamics-style module that predicts temporal difference image embeddings without any pixel reconstruction loss, encouraging balanced static and dynamic feature learning. DISReg consists of a static term that shapes the distribution of the image embedding and encourages slow features, and a dynamic term, which, unlike direct regularization on the embedding, imposes no constraint on the image embedding's shape or distribution and instead only incentivizes that dynamic features be present. By integrating this regularizer into a standard JEPA, we establish our new architecture, MotionJEPA. Latent probing demonstrates that MotionJEPA produces more complete representations than other methods, and our trajectory analysis shows it maintains geometrically simple latent embeddings with low curvature. We further show that MotionJEPA improves downstream planning success under static-background distractors across four environments.

TL;DR: MotionJEPA is a JEPA-based world model that avoids temporal feature collapse using DISReg, a difference-image regularizer that enables more stable, balanced learning of static and dynamic features. DISReg's dynamic term encourages dynamic features to be present without constraining the shape or distribution of the image latent, resulting in simple latent trajectories for downstream planning in the presence of static distractors.

MotionJEPA architecture: a JEPA encoder–predictor in blue, DISReg in green, and losses in red.

MotionJEPA overview. MotionJEPA uses our novel DISReg, whose temporal difference-image embedding $d_t$ provides visually grounded collapse prevention for changing features.

Approach

DISReg. DISReg adds two components to a standard JEPA: a difference-image encoder $\mathrm{DiffEnc}_{\alpha}$ and an inverse-dynamics-style predictor $\mathrm{DiffPred}_{\beta}$. $\mathrm{DiffPred}_{\beta}$ predicts the difference latent $d_t$ from consecutive image embeddings $z_t$ and $z_{t+1}$, with no pixel reconstruction and no action or state labels as targets:

$$\begin{aligned} d_t &= \mathrm{DiffEnc}_{\alpha}(o_{t+1} - o_t) \\ \hat{d}_{t} &= \mathrm{DiffPred}_{\beta}(z_t, z_{t+1}) \end{aligned}$$

Training objective. We use SIGReg to prevent $z$ and $d$ from collapsing. The MSE of the dynamic term, $\mathrm{MSE}(d_t,\hat{d}_t)$, only asks that visual changes be present in $z_t$ and $z_{t+1}$, without constraining the shape or distribution of the image latent. The full loss is the typical next-embedding prediction plus DISReg ($\lambda_z{=}0.25$, $\lambda_{\mathrm{pred}}{=}0.5$, $\lambda_d{=}2$):

$$\mathcal{L}_{\mathrm{DISReg}} = \lambda_z\,\mathrm{SIGReg}(z) + \lambda_d\,\mathrm{SIGReg}(d) + \lambda_{\mathrm{pred}}\,\mathrm{MSE}(d_t,\hat{d}_t)$$

$$\mathcal{L}_{\mathrm{MotionJEPA}} = \mathcal{L}_{\mathrm{DISReg}} + \mathrm{MSE}(z_{t+1},\hat{z}_{t+1})$$

Latent Probing

We probe frozen encoder embeddings $z$ with residual MLPs on three synthetic games with separable feature groups. An NMSE of $1.0$ corresponds to predicting the evaluation-target mean (the probe extracts no useful information from $z$). Besides LeWM and SMWM, we include LeWM variants with common anti-collapse modifications: LeWM-Time applies SIGReg along the time dimension, LeWM-Flat applies SIGReg along the flattened batch and time dimension, and LeWM-Detached adds a stop-gradient on the JEPA predictor target $\mathrm{MSE}(\mathrm{StopGradient}(z_{t+1}), \hat{z}_{t+1})$. Starred baselines use per-environment hyperparameter search; MotionJEPA uses one setting ($\lambda_z{=}0.25$, $\lambda_{\mathrm{pred}}{=}0.5$, $\lambda_d{=}2$) on all three games.

Offline residual-MLP probe NMSE $\downarrow$ of the image embedding $z$. MotionJEPA is the only method with NMSE $\leq 0.1$ on every feature group. Feature-group values $\leq 0.1$ are green; values above are red. Bold / underline mark best / second-best among the rows shown.

Method	Pong				Dino						Golf
Method	avg.	ball	bats	score	avg.	cactus	clouds	dino	hearts	sky	avg.	bars	cross	ball	wood
SMWM*	0.651	0.983	0.002	1.003	0.606	0.398	1.005	0.003	0.775	0.069	0.727	1.000	0.986	0.002	1.003
LeWM*	0.061	0.163	0.018	0.012	0.282	0.990	0.059	0.751	0.011	0.012	0.760	1.001	1.002	0.909	0.110
LeWM-Detached*	0.048	0.103	0.038	0.008	0.332	1.005	0.073	0.959	0.011	0.017	0.745	1.002	1.002	0.908	0.048
LeWM-Time*	0.687	0.022	0.972	1.002	0.697	0.005	1.004	0.003	0.971	1.004	0.985	0.972	0.992	0.984	1.002
LeWM-Flat*	0.008	0.012	0.004	0.008	0.301	0.650	0.498	0.002	0.007	0.016	0.299	0.317	0.059	0.312	0.377
MotionJEPA (Ours)	0.004	0.005	0.003	0.004	0.021	0.011	0.037	0.006	0.006	0.004	0.061	0.095	0.083	0.033	0.033

Latent Straightness

We visualize how individual state features of the Golf game map into the image embedding $z$. Each panel varies one feature over a 1D range while holding the rest of the state fixed. We show the first two principal components of this latent trajectory (with no variance rescaling). LeWM* collapses the crosses and bars nearly to a single point, and the LeWM-Flat* variant yields trajectories with substantially more folds and scrambles, while MotionJEPA keeps these paths geometrically simple.

Animated PCA trajectories of latent embeddings while sweeping individual Golf state features.

Planning under Static Distractors

We evaluate CEM planning on the four LeWorldModel control tasks after replacing the original background with a wood texture, using a different crop per episode and keeping that crop fixed within the episode. Initial states are sampled from the test split; horizons are 25 and 50 environment steps. Averaged across all environments, MotionJEPA outperforms IDM by +4.6 at 25-step planning and +11.0 at 50-step planning.

Planning success rate (%) under static-background distractors. Bold / underline: best / second-best among the rows shown.

Method	Cube	PushT	Reach.	2Room	Mean
25-step planning
LeWM	43.6	4.0	5.2	28.8	20.4
IDM	76.0	83.6	49.2	100.0	77.2
MotionJEPA (Ours)	78.0	78.4	71.2	99.6	81.8
50-step planning
LeWM	24.4	2.4	2.0	18.4	11.8
IDM	53.2	26.0	62.0	96.0	59.3
MotionJEPA (Ours)	63.2	28.0	98.0	92.0	70.3

Qualitative 25-step rollouts

Side-by-side CEM planning with a wood-texture background. Left: MotionJEPA. Right: original LeWM. All clips below are cases where MotionJEPA succeeds and LeWM fails.

Cube

PushT

Reacher

Two-Room

BibTeX

@misc{motionjepa2026,
      title={MotionJEPA: Preventing Temporal Feature Collapse by Capturing Visual Changes in Latent Space},
      author={Markus Karmann and Shile Li and Christian Internò and Bruno Andreis and David Klindt and Randall Balestriero and Jindong Gu and Philip Torr and Qi Zhang and Peng-Tao Jiang and Hao Zhang and Bo Li and Onay Urfalioglu},
      year={2026},
      eprint={2609.23881},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2609.23881},
}