에너지 기반 조인트 임베딩 예측 아키텍처를 위한 경량 라이브러리

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

본 논문에서는 조인트 임베딩 예측 아키텍처(JEPA)를 사용하여 표현 학습 및 세계 모델을 학습할 수 있는 오픈 소스 라이브러리인 EB-JEPA를 소개합니다. JEPA는 픽셀 공간 대신 표현 공간에서 예측을 학습하며, 생성 모델링의 단점을 피하면서 다운스트림 작업에 적합한 의미론적으로 중요한 특징을 포착합니다. 당사의 라이브러리는 모듈화되고 독립적인 구현을 제공하며, 이미지 수준의 자기 지도 학습을 위해 개발된 표현 학습 기술이 시간적 동적 특성이 더 복잡한 비디오로, 그리고 제어 입력의 효과를 예측해야 하는 행동 기반 세계 모델로 어떻게 적용될 수 있는지를 보여줍니다. 각 예제는 몇 시간 이내에 단일 GPU로 훈련할 수 있도록 설계되어, 에너지 기반 자기 지도 학습을 연구 및 교육에 쉽게 활용할 수 있도록 합니다. CIFAR-10 데이터셋에 대한 JEA 구성 요소의 성능 분석을 제공합니다. 이러한 표현을 분석한 결과, 91%의 정확도를 보여주며, 모델이 유용한 특징을 학습한다는 것을 나타냅니다. 비디오 데이터로 확장하여, 동일한 원리가 시간 모델링으로 어떻게 확장될 수 있는지 보여주는 Moving MNIST 데이터셋에 대한 다단계 예측 예제를 포함합니다. 마지막으로, 이러한 표현이 행동 기반 세계 모델을 어떻게 구동할 수 있는지 보여주며, Two Rooms 내비게이션 작업에서 97%의 계획 성공률을 달성합니다. 포괄적인 분석을 통해, 표현 붕괴를 방지하기 위한 각 정규화 구성 요소의 중요성을 밝힙니다. 코드 및 관련 자료는 https://github.com/facebookresearch/eb_jepa 에서 확인할 수 있습니다.

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.