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

AI summary

Overview

Research area: Self-supervised representation learning and world models for computer vision and control, specifically Joint-Embedding Predictive Architectures (JEPAs) framed through Energy-Based Models (EBMs).

Technical level: Intermediate. The paper assumes familiarity with self-supervised learning concepts (contrastive methods, stop-gradient, VICReg, EMA), but the library itself is designed to lower the barrier to entry for newcomers.

Scope: This is a software-and-tutorial paper: it introduces an open-source library (eb_jepa) providing three modular, single-GPU JEPA implementations—image representation learning, video prediction, and action-conditioned planning—along with ablations that quantify the role of each regularization component.

What This Paper Is About

JEPAs learn by predicting in a learned latent representation space rather than reconstructing pixels, which avoids spending compute on task-irrelevant visual detail. Despite growing interest, production-scale JEPA codebases are hard to navigate and typically require large distributed training setups or frozen pre-trained encoders, making the paradigm difficult to study or teach. This paper introduces EB-JEPA: a lightweight, modular library with self-contained examples that reproduce JEPA principles on modest hardware, plus systematic ablations showing which design choices (predictor structure, projectors, regularization) actually matter.

Key Contributions

1. Three progressively complex, accessible implementations: Image-JEPA (view-invariance on CIFAR-10), Video-JEPA (temporal prediction on Moving MNIST), and Action-Conditioned Video-JEPA (goal-conditioned planning in the Two Rooms navigation environment). Each trains on a single GPU within a few hours.

2. A modular architecture: Encoders (ResNet, ViT, IMPALA), predictors (UNet-style spatial, GRU temporal), regularizers (VICReg, SIGReg, temporal similarity, inverse dynamics), and planners (MPPI, CEM) are exposed as interchangeable components.

3. Systematic ablations of JEPA components: Quantifies the impact of projector design, regularizer choice, loss coefficients, and predicting multiple rollout steps, including the finding that each regularization term is essential to prevent representation collapse.

4. An educational resource: Clear documentation and concise code intended to help researchers and students understand energy-based self-supervised learning and world modeling without needing large-scale infrastructure.

Main Findings

• Image representation quality: Linear probing on CIFAR-10 reaches roughly 90–91% accuracy with both SIGReg and VICReg, competitive with prior self-supervised methods on this benchmark. SIGReg peaks at 91.02%, VICReg at 90.12%.

• Projectors matter: Regularizing in a learned projected space rather than directly on encoder outputs yields roughly a 3-point accuracy improvement for both regularizers. SIGReg prefers a bottleneck projector (e.g., 2048 → 128), while VICReg prefers a wider output (e.g., 2048 → 1024).

• SIGReg is more forgiving to tune: With only one hyperparameter (λ), SIGReg maintains high accuracy across a wider range of settings; VICReg reaches similar peak performance but is more sensitive to its standard-deviation and covariance coefficients. A mismatched λ (0.1) collapses SIGReg accuracy to 27%.

• Multi-step rollout training improves video prediction: Training with k-step recursive prediction (Pareto optimum near k = 4) significantly outperforms single-step training on downstream detection Average Precision, reducing the mismatch between teacher-forced training and autoregressive inference.

• Strong planning performance: The action-conditioned model reaches 97% ± 2% success on the Two Rooms task with randomized wall positions, using MPPI with cumulative cost. CEM is close behind at 96% ± 2%; final-state-only cost drops to 89%.

• Every regularizer is load-bearing: Removing the inverse dynamics model loss collapses planning success to 1% ± 1%. Removing variance or covariance terms each drops success to roughly 46–47%. Removing temporal similarity drops it to 61%. This is the paper's most striking ablation result.

• Planning cost design matters: A cumulative distance over all imagined states (rather than only the final state) both encourages efficient paths and is more robust to compounding prediction errors, worth an 8-point improvement.

Methodology in Plain English

The authors frame JEPA training as energy-based learning: a model assigns low "energy" to compatible input–output pairs and high energy to incompatible ones. In all three examples, energy equals prediction error in latent space—the distance between what the encoder says a target should look like and what the predictor guesses it will look like.

The central risk in this setup is collapse: the trivial solution where the encoder outputs the same vector for every input and the prediction error is always zero. The library prevents this with explicit regularizers. VICReg uses two terms—a variance term that forces each feature dimension to spread out across a batch, and a covariance term that decorrelates dimensions so the model uses its full capacity. SIGReg takes a different route: it tests whether the embedding distribution matches an isotropic Gaussian along random one-dimensional projections, which theory suggests is optimal for downstream tasks.

Starting from this shared foundation, complexity increases in three stages. Image-JEPA creates two augmented views of one image and trains the encoder to produce matching representations (the predictor is the identity). Video-JEPA adds a temporal predictor that consumes a window of past frame representations and predicts the next one. Action-Conditioned Video-JEPA adds an action encoder, so the predictor learns how control inputs change the scene—a latent world model.

For video and action-conditioned settings, the authors train with k-step rollouts: the model predicts several steps forward recursively, so training resembles inference and exposure bias shrinks. For planning, they encode a goal image and use MPPI (Model Predictive Path Integral), a sampling-based optimizer, to search for action sequences whose imagined latent trajectories land closest to the goal, weighting samples by exponentiated negative energy.

Why This Matters

Impact on research. JEPA is a theory-heavy paradigm whose central concepts—collapse prevention, predictive latent spaces, energy-based objectives—are difficult to grasp without runnable code. EB-JEPA makes those mechanisms directly inspectable and cheap to modify, which should accelerate algorithmic experimentation (e.g., testing new regularizers, predictors, or hierarchies) before committing to expensive distributed runs. Its ablation results also serve as a practical checklist: each regularization term is not decorative, and the inverse dynamics loss in particular is the difference between a working world model and a degenerate one.

Real-world applications (each enabled by the underlying JEPA world-modeling capability, not necessarily by this specific library):

• Robot navigation and manipulation: Action-conditioned latent world models enable planning that generalizes across randomized environments without pixel-level simulation.
• Autonomous driving: Predicting latent future states conditioned on steering/throttle lets planners reason about consequences cheaply compared to pixel-space video prediction.
• Video understanding and surveillance: Temporal latent prediction supports anomaly detection and activity recognition without reconstructing every frame.
• Compute-efficient self-supervised pretraining: View-invariance learning at 90%+ linear probe accuracy on a single GPU is attractive for labs and organizations without large clusters.

Industry relevance. The library is released by Meta FAIR with Yann LeCun as a co-author, signaling continued investment in the JEPA research direction as an alternative to generative pixel-space modeling. The single-GPU, few-hours-per-example design makes it practical for small teams, academic groups, and internal prototyping, and the paper explicitly positions EB-JEPA as an upstream step in a two-stage workflow: validate ideas cheaply here, then scale on production codebases like stable-pretraining or JEPA-WMs.

Future Directions

1. Better regularization theory. It remains unclear why particular combinations of variance, covariance, temporal similarity, and inverse dynamics losses work, or when each is strictly necessary. The library is framed as a testbed for controlled studies of these dynamics and for automatic hyperparameter selection.

2. Hierarchical world models. Current JEPAs predict at a single temporal resolution, but long-horizon planning likely benefits from reasoning at multiple timescales—fine-grained dynamics for local control, coarse abstractions for distant goals. The modular encoder/predictor/regularizer split is a natural starting point.

3. Learned cost and value functions. Planning currently relies on a simple distance-to-goal in latent space, which may be suboptimal for complex tasks. Learning task-specific costs or value functions from demonstrations or sparse rewards could bridge pure world modeling and reward-driven reinforcement learning.

4. Scaling validated ideas. Determining which small-scale findings hold at scale—and how they interact with frozen pre-trained backbones, distributed training, and more complex benchmarks—remains an open transfer question that this library deliberately does not answer.

Target Audience

• Graduate students and newcomers to self-supervised learning who want a concrete, runnable entry point into JEPA and energy-based representation learning.
• Researchers prototyping JEPA variants who need a fast, single-GPU testbed before scaling to production infrastructure.
• Instructors and course designers teaching representation learning or world models, for whom the three-example progression provides a ready-made curriculum.
• Practitioners in robotics and control interested in latent world models for planning, who can adapt the action-conditioned example to their own environments.

Readers looking for state-of-the-art benchmark numbers on large-scale vision tasks will not find them here—the paper is explicit that its contribution is accessibility, modularity, and insight into component design, not peak performance on complex benchmarks.

Authors’ 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.