UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
2026-08-07 • Computer Vision and Pattern Recognition
Computer Vision and Pattern Recognition
AI summaryⓘ
The authors introduce UniJEPA, a new approach that combines different ways of predicting image changes and video changes into one model. Unlike past methods that handled image and video predictions separately, UniJEPA learns both from raw pixels using a single training step without complex tricks. This unified model can understand both how images change in style and how scenes change over time, supporting planning tasks without extra training. Their results show UniJEPA performs as well or better than specialized models and plans faster than other world models.
Joint-Embedding Predictive Architecturesself-supervised learninglatent spacephotometric predictiontemporal predictionencoder-predictoranti-collapse regularizerzero-shot planningoffline trajectoriesworld models
Authors
An Lanji, Dawei Liu, Jin Li, Haoran Xu, Mei Chen, Yu Tian
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.