A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

TL;DR

EB-JEPA presents a lightweight, open-source library for learning representations and world models using Joint-Embedding Predictive Architectures, which train by predicting in a learned representation space rather than reconstructing pixels. The framework unifies an energy-based objective with regularization to prevent collapse, covering image representation learning (Image-JEPA), video prediction (Video-JEPA), and action-conditioned world modeling (AC-video-JEPA). The authors implement modular components (encoders, predictors, regularizers, planners) and demonstrate single-GPU training on CIFAR-10, Moving MNIST, and Two Rooms, achieving strong signals such as about 91% linear probing accuracy on CIFAR-10 and 97% planning success. Through exhaustive ablations, they show the critical role of all regularizers (variance, covariance, temporal similarity, inverse dynamics) and provide a practical template for rapid JEPA experimentation at small scale.

Abstract

We present EB-JEPA, an open-source library for learning representations and world models using Joint-Embedding Predictive Architectures (JEPAs). JEPAs learn to predict in representation space rather than pixel space, avoiding the pitfalls of generative modeling while capturing semantically meaningful features suitable for downstream tasks. Our library provides modular, self-contained implementations that illustrate how representation learning techniques developed for image-level self-supervised learning can transfer to video, where temporal dynamics add complexity, and ultimately to action-conditioned world models, where the model must additionally learn to predict the effects of control inputs. Each example is designed for single-GPU training within a few hours, making energy-based self-supervised learning accessible for research and education. We provide ablations of JEA components on CIFAR-10. Probing these representations yields 91% accuracy, indicating that the model learns useful features. Extending to video, we include a multi-step prediction example on Moving MNIST that demonstrates how the same principles scale to temporal modeling. Finally, we show how these representations can drive action-conditioned world models, achieving a 97% planning success rate on the Two Rooms navigation task. Comprehensive ablations reveal the critical importance of each regularization component for preventing representation collapse. Code is available at https://github.com/facebookresearch/eb_jepa.

Figures (5)

• Figure 1: EB-JEPA is a modular code base and tutorial, providing self-contained implementations of Joint-Embedding Predictive Architecture for (a) self-supervised image representation learning (b) video prediction in latent space, and (c) action-conditioned world models that enable goal-directed planning (d).
• Figure 2: Hyperparameter sensitivity comparison between SIGReg and VICReg on CIFAR-10. SIGReg demonstrates greater stability across different hyperparameter configurations, while VICReg achieves similar peak performance but requires more careful tuning.
• Figure 3: Video-JEPA training dynamics and multistep rollout ablation. (a) Training dynamics over 50 epochs: variance-covariance regularization loss $\mathcal{R}$ (left), prediction loss $\mathcal{L}_{\text{pred}}$ (center), and mean Average Precision (right). (b) Training with $k$-step recursive predictions achieves significantly better Average Precision compared to single-step predictions, demonstrating improved temporal understanding, with a Pareto optimum around $k=4$ rollout steps.
• Figure 4: Video JEPA visualization on Moving MNIST. From left to right: input frames, 1-step prediction visualization, and full autoregressive rollout. The model maintains coherent predictions of digit motion over extended horizons, correctly capturing trajectory and dynamics.
• Figure 5: Visualization of three successful planning evaluation episodes of our AC-video-JEPA on the Two Rooms environment with random wall. From left to right: initial frame (red), full episode outputted by the planning optimization procedure, goal frame used to define planning cost (red). Each episodes allows a maximum of 200 steps in the environment.