UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
UniJEPA unifies image and video predictive modeling in a single latent space, using a single end-to-end loss with Gaussian regularization, enabling task-agnostic world modeling and zero-shot planning.
Key Findings
Methodology
UniJEPA introduces a unified joint-embedding predictive architecture that employs a shared encoder and predictor operating within a common latent space. It simultaneously models photometric transformations (such as brightness, contrast, hue adjustments) and temporal dynamics (future state prediction) through a single end-to-end loss function comprising a next-embedding prediction loss and a Gaussian regularizer. The regularizer theoretically guarantees avoidance of representational collapse, ensuring diverse and informative embeddings. The architecture is based on Vision Transformer (ViT), enabling shared feature extraction for both tasks. During training, the model directly learns from raw pixels without requiring EMA, stop-gradient, or pre-trained encoders. Post-training on offline trajectories allows the model to perform zero-shot action-conditioned planning by predicting future latent states conditioned on actions, facilitating efficient path planning in visual environments.
Key Results
- On ImageNet, UniJEPA achieves 74.9% linear probe accuracy, surpassing I-WM (73.5%) and approaching DINOv2 (81.3%), with a single hyperparameter for training, demonstrating high efficiency and simplicity.
- In video understanding benchmarks, it reaches 78.1% top-1 accuracy on Something-Something-v2, outperforming V-JEPA-2 (77.3%), and achieves 40.6% recall@5 on Epic-Kitchens, indicating strong generalization across modalities.
- For control tasks, after offline trajectory post-training, UniJEPA attains 75.8% success rate in zero-shot planning, with speeds up to 44× faster than pixel-based generative models, validating its practical applicability.
Significance
This work bridges the gap between image and video world modeling by unifying them into a single latent space, which enhances transferability, sample efficiency, and robustness. The theoretical guarantee against collapse ensures stable training, reducing reliance on heuristics. Its ability to perform zero-shot planning directly from raw observations opens new avenues for autonomous agents in robotics and simulation, enabling faster, more reliable decision-making in complex environments. The approach significantly advances the state-of-the-art in self-supervised, task-agnostic world modeling, with broad implications for AI research and deployment.
Technical Contribution
UniJEPA's key technical innovation lies in formulating a single, unified loss function that combines photometric and temporal prediction tasks, supported by a Gaussian regularizer with provable anti-collapse properties. The architecture employs a shared Vision Transformer encoder and a single predictor, simplifying multi-task learning and reducing model complexity. The theoretical analysis guarantees that the regularizer prevents collapse, ensuring diverse and meaningful representations. Additionally, the model supports post-training for action-conditioned planning, enabling zero-shot control without additional supervision. This integration of multi-task learning, theoretical rigor, and practical efficiency marks a significant step forward in self-supervised world modeling.
Novelty
UniJEPA is the first to unify image-level photometric prediction and video-level temporal prediction within a single, end-to-end trainable latent space, using a single regularizer to prevent collapse. Unlike prior methods that treat these tasks separately with distinct models and objectives, UniJEPA couples them into one framework, enabling simultaneous learning of invariant structure and dynamic behavior. Its theoretical proof of anti-collapse guarantees, combined with the ability to perform zero-shot planning, distinguishes it from existing approaches that rely on multiple models, heuristics, or pretraining. This holistic approach advances the field by demonstrating that a single, unified latent space can support diverse, complex tasks efficiently.
Limitations
- The model's capacity to handle highly complex, real-world environments with extreme variability in lighting, textures, or long-term dependencies remains limited. Its performance may degrade in scenarios with abrupt or unpredictable changes, requiring further robustness enhancements.
- While the Gaussian regularizer prevents collapse, it may restrict the richness of the latent representations, especially at larger scales or in multi-modal settings, potentially limiting the model's expressiveness.
- Current experiments focus on static datasets and controlled environments; deploying UniJEPA in real-world, dynamic scenarios with noisy observations, occlusions, or multimodal inputs will require additional validation and adaptation.
Future Work
Future research will explore extending UniJEPA to multi-modal data, such as combining visual, auditory, and proprioceptive inputs, to build more comprehensive world models. Incorporating hierarchical or multi-scale architectures could improve the handling of long-term dependencies and complex scenes. Additionally, integrating reinforcement learning algorithms with the current framework may enhance autonomous decision-making and planning in real-world robotics. Further theoretical work on regularization techniques and representation capacity will aim to balance stability and expressiveness, pushing the boundaries of task-agnostic, scalable world modeling.
AI Executive Summary
Building intelligent agents capable of understanding and predicting their environment without task-specific supervision has long been a central goal in AI research. Traditional generative models, such as VAEs and GANs, excel at reconstructing detailed observations but are computationally intensive and often lack robustness in complex scenarios. Recent advances in self-supervised learning have shifted focus toward latent space prediction methods, notably the Joint-Embedding Predictive Architectures (JEPA), which encode observations into compact representations and forecast their evolution over time.
Despite their promise, existing JEPA variants tend to specialize in either image-level tasks, such as masked region prediction, or video-level tasks, like future state forecasting. This specialization results in fragmented models with incompatible latent spaces, complicating multi-task learning and limiting their applicability in real-world autonomous systems. Moreover, these models often require multiple training objectives, heuristics, or pretraining steps, which increase complexity and reduce efficiency.
The paper introduces UniJEPA, a unified framework that addresses these limitations by integrating image and video prediction tasks into a single, end-to-end trainable model. The core innovation is a shared encoder based on Vision Transformer (ViT), coupled with a predictor that operates in the same latent space for both photometric and temporal predictions. A single loss function, combining a next-embedding prediction term and a Gaussian regularizer, ensures the model learns diverse, non-collapsed representations without auxiliary techniques like EMA or stop-gradient. Theoretical analysis confirms that this regularizer guarantees non-degenerate embeddings, providing stability and robustness.
Experimental results demonstrate that UniJEPA achieves state-of-the-art performance across multiple benchmarks. On ImageNet, it attains 74.9% linear probe accuracy, surpassing prior light-weight models, while on video datasets like Something-Something-v2, it reaches 78.1%. In control tasks, the model supports zero-shot planning with success rates up to 75.8%, with inference speeds significantly faster than pixel-based generative models. These results validate the model’s ability to unify diverse prediction tasks within a single, efficient architecture.
This work marks a significant step toward scalable, task-agnostic world modeling. By unifying image and video prediction in a single latent space, UniJEPA enhances transferability, reduces training complexity, and opens new avenues for autonomous decision-making. Its theoretical guarantees and practical efficiency make it a promising foundation for future research in self-supervised learning, robotics, and embodied AI. The authors also outline future directions, including multi-modal extensions, hierarchical modeling, and reinforcement learning integration, aiming to bring autonomous agents closer to human-like understanding and reasoning capabilities.
Deep Dive
Glossary
Joint-Embedding Predictive Architecture (JEPA)
一种自监督学习框架，将观察映射到紧凑潜空间，并在该空间中预测未来或遮挡内容，避免像素重建，提升效率。
论文中提出的核心模型架构，用于多任务世界建模。
潜空间 (latent space)
模型内部的抽象特征空间，用于表示观察的压缩信息，便于进行预测和推理。
UniJEPA在共享潜空间中同时实现图像和视频预测。
正则化 (regularizer)
一种机制，用于限制模型参数或表示，防止崩溃或过拟合，确保潜空间多样性。
UniJEPA采用高斯正则保证潜空间非崩溃。
非崩溃 (anti-collapse)
确保模型的潜表示不退化为单一或无信息的状态，保持多样性和表达能力。
理论证明正则化实现了模型的非崩溃性。
Vision Transformer (ViT)
一种基于Transformer的视觉模型，将图像划分为块，进行序列建模，具有强大的表达能力。
UniJEPA的编码器基础架构。
零样本规划 (zero-shot planning)
在没有额外训练或示范的情况下，利用已有模型预测目标路径或动作。
模型在离线轨迹上后训练，实现目标导向的路径规划。
正向预测 (forward prediction)
预测未来状态或内容的任务，在潜空间中进行，避免像素重建。
模型同时进行光度和时间预测。
模型预测控制 (Model Predictive Control, MPC)
一种基于模型的规划方法，通过多步预测优化动作序列。
UniJEPA用于零样本路径规划。
正态正则 (Gaussian regularizer)
引入高斯分布约束潜空间，保证表示的多样性和非崩溃性。
理论分析支持其防止崩溃的效果。
多任务学习 (multi-task learning)
同时训练模型完成多个相关任务，提高泛化能力。
UniJEPA在单一潜空间中实现多任务预测。
自监督学习 (self-supervised learning)
无需标注数据，通过预测部分信息实现模型训练。
UniJEPA利用自监督目标进行训练。
路径规划 (path planning)
在环境中找到从起点到目标的最优路径，常用于机器人控制。
模型支持零样本路径规划。
信息压缩 (information compression)
将高维观察压缩到低维潜空间，保留关键信息。
UniJEPA在潜空间中进行预测。
多模态预测 (multi-modal prediction)
同时处理多种类型的数据（如视觉、声音、动作），实现多角度理解。
未来扩展方向之一。
Open Questions Unanswered questions from this research
- 1 尽管UniJEPA在多个任务中表现优异，但在极端复杂环境（如高动态、强光变化、长时间依赖）中的表现仍有限。如何增强模型的鲁棒性，特别是在真实世界中面对噪声、遮挡和多模态信息融合时，仍需深入研究。此外，模型在超大规模和多模态场景下的训练成本和调优难度较高，如何实现更高效的训练和推理机制，是未来的重要方向。
Applications
Immediate Applications
自主机器人感知与控制
利用UniJEPA进行环境理解和路径预测，支持机器人在复杂环境中的自主导航和操作，无需大量标注数据，提升自主性和适应性。
自动驾驶环境建模
在自动驾驶中，UniJEPA可以实现对道路场景的动态预测和路径规划，提高系统的反应速度和鲁棒性。
虚拟现实中的场景理解
通过高效的场景建模与预测，增强虚拟环境中的交互体验，实现更自然的虚拟交互。
Long-term Vision
通用自主智能体
结合多模态信息和强化学习，打造具有自主感知、推理和规划能力的通用智能体，广泛应用于机器人、自动驾驶和智能制造。
智能环境的自我维护与优化
未来系统能自主学习环境变化规律，进行主动维护和优化，实现长时间自主运行和自我进化。
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.
References (20)
Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets
A. Blattmann, Tim Dockhorn, Sumith Kulal et al.
Revisiting Feature Prediction for Learning Visual Representations from Video
Adrien Bardes, Q. Garrido, Jean Ponce et al.
A Simple Framework for Contrastive Learning of Visual Representations
Ting Chen, Simon Kornblith, Mohammad Norouzi et al.
V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning
Mahmoud Assran, Adrien Bardes, David Fan et al.
Back to the Features: DINO as a Foundation for Video World Models
Federico Baldassarre, Marc Szafraniec, Basile Terver et al.
DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning
Gaoyue Zhou, Hengkai Pan, Yann LeCun et al.
DINOv2: Learning Robust Visual Features without Supervision
M. Oquab, Timothée Darcet, Théo Moutakanni et al.
WorldSimBench: Towards Video Generation Models as World Simulators
Yiran Qin, Zhelun Shi, Jiwen Yu et al.
LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
Lucas Maes, Quentin Le Lidec, Damien Scieur et al.
Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
Mahmoud Assran, Quentin Duval, Ishan Misra et al.
Temporal Difference Learning for Model Predictive Control
Nicklas Hansen, Xiaolong Wang, H. Su
Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
Ze Liu, Yutong Lin, Yue Cao et al.
QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation
Dmitry Kalashnikov, A. Irpan, P. Pastor et al.
Mastering Diverse Domains through World Models
Danijar Hafner, J. Pašukonis, Jimmy Ba et al.
InstrucRobo: Object-centric multi-instruction decoupling model for explainable robotic manipulation
Panqi Yang, Haodong Jing, Nanning Zheng et al.
The “Something Something” Video Database for Learning and Evaluating Visual Common Sense
Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski et al.
Scaling Egocentric Vision: The EPIC-KITCHENS Dataset
D. Damen, Hazel Doughty, G. Farinella et al.
RT-1: Robotics Transformer for Real-World Control at Scale
Anthony Brohan, Noah Brown, Justice Carbajal et al.
SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation
Vijay Badrinarayanan, Alex Kendall, R. Cipolla
Learning Latent Dynamics for Planning from Pixels
Danijar Hafner, T. Lillicrap, Ian S. Fischer et al.