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

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.

• The paper presents EB-JEPA, a modular library that learns high-level semantic representations via energy-based joint-embedding predictive models.
• It demonstrates efficient training on limited resources across static images, video prediction, and action-conditioned world modeling with robust regularization.
• Empirical results on CIFAR-10 and Moving MNIST reveal strong performance, hyperparameter stability, and improved planning success with inverse dynamics.

Modular Energy-Based Joint-Embedding Predictive Architectures for Visual Representations and World Modeling

Overview of EB-JEPA Library and Objectives

The paper presents EB-JEPA, an open-source, modular library designed for accessible experimentation and education with Joint-Embedding Predictive Architectures (JEPAs). The central premise behind JEPA is to learn predictive models in a high-level representation space rather than in the raw observation space, thereby focusing on extracting task-relevant, semantic abstractions critical for downstream applications in perception and control. This approach contrasts with generative models that reconstruct every observation pixel, which are often computationally expensive and susceptible to modeling irrelevant variations.

EB-JEPA operationalizes these principles through a concise codebase implementing three escalating modalities: self-supervised learning on static images, temporal prediction in videos, and action-conditioned world modeling for goal-directed planning. Each implementation is engineered for training on limited resources (single GPU, several hours), dramatically lowering the research and educational entry barriers for modern energy-based and joint-embedding learning.

Figure 1: EB-JEPA provides a modular foundation and self-contained examples for learning representations via joint-embedding predictive objectives across (a) images, (b) temporal videos, (c) action-conditioned world modeling, and (d) goal-driven planning.

Unified JEPA Framework and Energy-Based Modeling

The JEPA framework unifies representation learning and prediction across diverse settings (static, sequential, action-conditioned) via an energy-based formulation. For a state $x$ (optionally conditioned on actions $a$), an encoder $f_\theta$ maps inputs to latent vectors, and a predictor $g_\phi$ is trained to predict future latent abstractions based on current and past representations, optionally conditioned by action encodings $q_\omega(a)$. The scalar energy function quantifies compatibility via prediction error in latent space, with the general objective:

$$\mathcal{L} = \mathcal{L}_{\text{pred}}(g_\phi(z, u), z') + \lambda \, \mathcal{R}(z)$$

Collapse prevention—ensuring models do not output trivial, non-informative representations—is enforced not by sampling explicit negatives but via regularization objectives. The library implements both VICReg’s dual variance-covariance regularization and SIGReg’s isotropic Gaussianity encouragement, the latter motivated by recent theory on optimal latent spaces for downstream prediction risk.

Empirical Findings: Representation Learning and Regularization

On CIFAR-10, linear probing of JEPA representations yields accuracies up to 91%, with comprehensive ablations showing that:

• Learned projectors give a ~3% gain, especially with bottleneck architectures for SIGReg.
• SIGReg shows greater hyperparameter stability (i.e., robust performance with minimal tuning) relative to VICReg, which achieves similar peak performance but is sensitive to parameterization.

  Figure 2: Hyperparameter sensitivity on CIFAR-10: SIGReg is markedly more stable, while VICReg is more sensitive to tuning.

• In both image and video settings, regularization in a projected embedding space (as opposed to the encoder output) is more effective.
• Multistep prediction (with $k$-step rollouts) during training, rather than one-step, aligns learning dynamics with autoregressive evaluation and mitigates compounding error (“exposure bias”).

Figure 3: Multistep rollout training on video-JEPA: variance-covariance loss, prediction loss, and mean Average Precision over epochs, showing benefits of recursive prediction.

Video Prediction and Temporal Dynamics

Applying JEPA to video modeling on the Moving MNIST dataset demonstrates prolonged temporal coherence in prediction. Models exhibit robust multistep latent-space prediction, with the ability to track object trajectories and dynamics over extended horizons.

Figure 4: Moving MNIST video-JEPA: input sequence, 1-step prediction, and autoregressive multi-step rollout—the model correctly infers digit motion and preserves trajectory consistency.

Ablation experiments confirm that increasing the rollout horizon during training yields improved downstream detection scores, attributable to better learning of temporal dependencies and reduced mismatch between training and inference.

Action-Conditioned World Models and Planning

The action-conditioned extension (AC-video-JEPA) instantiates latent world models by integrating action encodings, learning latent dynamics suitable for planning—critical for model-based RL and robotics. The method supports goal-conditioned planning by searching for action sequences that minimize the latent-space distance to the goal embedding, using MPPI for trajectory optimization.

Evaluation on the Two Rooms navigation task shows a 97% success rate on environment variants with randomized wall positions, confirming the method’s practical effectiveness even under substantial stochasticity and non-monotonicity in optimal trajectories.

Figure 5: Three successful AC-video-JEPA planning rollouts (Two Rooms, randomized walls): initial state, entire predicted trajectory, and goal state; all solutions efficiently reach their goals within task constraints.

Ablation of loss components demonstrates that inverse dynamics modeling (IDM) is essential for collapse prevention (success drops to 1% without it), while variance, covariance, and temporal smoothness each supply substantial complementary gains. Additionally, using a cumulative cost over all predicted states, rather than just a final-state-matching cost, is empirically superior for path efficiency and robustness to prediction noise.

Implications, Theoretical Insights, and Future Directions

EB-JEPA’s design facilitates rapid mechanistic research on self-supervision, energy-based modeling, and world models:

• Regularization and collapse: The detailed ablations highlight the nuanced interplay between variance, covariance, temporal, and inverse dynamics regularizers. The testbed supports further theory-driven deconstruction and principled hyperparameter tuning or adaptation strategies.
• Hierarchical world models: The modular split between encoding, prediction, and regularization naturally extends to multi-scale/hierarchical temporal modeling, an open direction for enabling agents to reason across long horizons.
• Learned cost/value functions: While current planning uses generic latent distances, the modularity of the cost function interface positions the method for future integration with value learning, reward propagation, or imitation learning.
• Bridging algorithmic prototyping and large-scale evaluation: The EB-JEPA library's small-scale, readable codebase complements large JEPA repositories for empirically grounding new algorithmic ideas before scaling.

Conclusion

The EB-JEPA library provides a unified, energy-based framework for joint-embedding self-supervised learning across images, videos, and action-conditioned environments, balancing accessibility and pedagogical clarity with extensive empirical rigor. The results underscore the importance of carefully designed regularization in energy-based and joint-embedding approaches, and demonstrate that sophisticated planning and world modeling are attainable outside large-scale setups. This positions EB-JEPA as an enabling resource for foundational algorithmic and theoretical advances in self-supervised and model-based reinforcement learning.