⚡ EB-JEPA

Energy-Based Joint-Embedding Predictive Architectures

Meta AI Research, FAIR

Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido, Tushar Nagarajan,
Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar

An open source library and tutorial for learning representations for
prediction and planning using joint embedding predictive architectures.

Each example is (almost) self-contained and training takes up to a few hours on a single GPU card.

📚 Examples

Image JEPA

Self-supervised representations from unlabeled images on CIFAR-10, evaluated on classification.

Video JEPA

Predict next image representation in a sequence.

AC Video JEPA

JEPA for world modeling + planning in Two Rooms environment.

Planning Episode	Task Definition
Successful planning episode	From init to goal state

🚀 Installation

We use uv for package management.

# Install dependencies
uv sync
# Option 1: Activate virtual environment
source .venv/bin/activate
python main.py
# Option 2: Run directly with uv
uv run python main.py

If you need conda-specific packages, you can use Conda + uv

# Create conda environment with Python 3.12
conda create -n eb_jepa python=3.12 -y
conda activate eb_jepa
# Install package in editable mode with dev dependencies (pytest, black, isort)
uv pip install -e . --group dev

Add these to your ~/.bashrc for persistent configuration.

# Required for SLURM jobs to find datasets
export EBJEPA_DSETS=/path/to/eb_jepa/datasets
# Optional: Directory for checkpoints and logs
export EBJEPA_CKPTS=/path/to/checkpoints

🏋️ Training

Quick Start

# Local training
python -m examples.{image_jepa,video_jepa,ac_video_jepa}.main

Our default configs are tuned for H100 GPUs. With older GPUs (e.g., A100, V100), you may need to reduce batch size to fit in memory.

📂 Folder Structure

All experiments use a unified folder structure:

checkpoints/
└── {example_name}/
    ├── dev_2026-01-16_00-10/                 # Single/local runs (dev_ prefix)
    │   └── {exp_name}_seed1/
    │
    ├── sweep_2026-01-16_00-10/         # Auto-named 3-seed sweep
    │   ├── {exp_name}_seed1/
    │   ├── {exp_name}_seed1000/
    │   └── {exp_name}_seed10000/
    │
    └── sweep_my_experiment/            # Custom-named sweep
        └── ...

{exp_name} encodes key hyperparameters to avoid folder collisions, e.g.:

• image_jepa: resnet_vicreg_proj_bs256_ep300_ph2048_po2048_std1.0_cov80.0
• video_jepa: resnet_bs64_lr0.001_std10.0_cov100.0
• ac_video_jepa: impala_cov8_std16_simt12_idm1

# 3 seeds with wandb averaging (recommended)
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml

# Custom sweep name
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --sweep my_experiment

# Single job
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --single

# Full hyperparameter sweep
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep

# With wandb sweep UI for hyperparameter analysis
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep

# Example: examples/image_jepa/cfgs/default.yaml
sweep:
  param_grid:
    loss.cov_coeff: [0.1, 1.0, 10.0, 100.0]
    loss.std_coeff: [1.0, 10.0]
    meta.seed: [1, 1000, 10000]

🧪 Running test cases

Libraries added to eb_jepa must have their own test cases. To run the tests:

# With uv sync installation
uv run pytest tests/
# With conda + uv installation (no .venv created)
pytest tests/

👩‍💻 Development

Before contributing, please format your code with the following tools:

# Remove unused imports
autoflake --remove-all-unused-imports -r --in-place .
# Sort imports
python -m isort eb_jepa examples tests
# Format code
python -m black eb_jepa examples tests

📚 Citing EB-JEPA

If you find this repository useful, please consider giving a ⭐ and citing:

@misc{terver2026lightweightlibraryenergybasedjointembedding,
      title={A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures},
      author={Basile Terver and Randall Balestriero and Megi Dervishi and David Fan and Quentin Garrido and Tushar Nagarajan and Koustuv Sinha and Wancong Zhang and Mike Rabbat and Yann LeCun and Amir Bar},
      year={2026},
      eprint={2602.03604},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2602.03604},
}

📄 License

EB-JEPA is Apache licensed. See LICENSE.