UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

2026-08-07

AI Summary

This work addresses the challenge of jointly modeling photometric invariance in images and temporal dynamics in videos within a unified framework. The authors propose UniJEPA, the first architecture that learns both image-level photometric prediction and video-level temporal state prediction end-to-end in a shared latent space, without relying on exponential moving averages (EMA), stop-gradient operations, or pretrained encoders. By combining next-embedding prediction loss with Gaussian regularization, UniJEPA achieves controllable abstraction: its photometric branch captures structural invariance, while its temporal branch learns dynamic equivariance. Experiments demonstrate that UniJEPA matches or exceeds the performance of specialized models across image, video, and control tasks, using only a single loss hyperparameter. Moreover, it enables zero-shot planning that is tens of times faster than generative world models while maintaining comparable accuracy.

Abstract

Joint-Embedding Predictive Architectures (JEPAs) have emerged as a principled framework for self-supervised learning of world models in compact latent spaces, yet existing methods are fragmented: some predict masked parts of a single image in latent space (I-JEPA), others learn to predict global photometric transformations (Image World Models), while video-scale JEPAs predict future temporal states and are post-trained for action-conditioned planning (V-JEPA~2, DINO-World, DINO-WM). These objectives are treated as distinct recipes with separate encoders, predictors, and anti-collapse regularizers, hindering a single model from unifying image-level and video-level world modeling. We present UniJEPA, a unified JEPA that jointly learns photometric prediction (image-level transformations) and temporal prediction (video-level next-state dynamics) in one shared latent space. A single end-to-end objective, composed of a next-embedding prediction loss and a Gaussian regularizer, yields a provably anti-collapse encoder-predictor pair trainable from raw pixels without EMA, stop-gradient, or pre-trained encoders. We show that the same latent space supports controllable abstraction: photometric prediction learns invariant structure while temporal prediction learns equivariant dynamics. After action-conditioned post-training on offline trajectories, UniJEPA enables zero-shot planning by treating goal features as prediction targets. On image, video, and control benchmarks, UniJEPA matches or surpasses task-specific JEPAs while requiring a single loss hyperparameter, and plans up to tens of times faster than generative world models at comparable accuracy.

Problem

• Joint-Embedding Predictive Architecture
• self-supervised learning
• world modeling
• latent space
• task-agnostic

Innovation

• Joint-Embedding Predictive Architecture
• Unified World Modeling
• Anti-Collapse Regularization
• Zero-Shot Planning
• Latent Space Abstraction