에너지 기반 공동 임베딩 예측 라이브러리

초록

EB‑JEPA는 이미지, 비디오, 행동 조건 세계 모델을 단일 GPU에서 몇 시간 안에 학습할 수 있도록 설계된 오픈소스 라이브러리이다. 표준 이미지 셀프‑슈퍼바이즈드 학습에서 시작해 시간적 예측과 행동‑조건 플래닝까지 단계적으로 확장하며, 정규화 기반의 붕괴 방지 기법(VICReg, SIGReg)과 다단계 롤아웃 손실을 결합한다. CIFAR‑10에서 91 % 선형 프로빙 정확도, Moving MNIST에서 장기 예측 시각화, Two Rooms 환경에서 97 % 플래닝 성공률을 기록한다.

상세 분석

EB‑JEPA 논문은 최근 각광받는 Joint‑Embedding Predictive Architecture(JEPA)를 실제 연구와 교육 현장에서 손쉽게 활용할 수 있도록 구현한 라이브러리를 제시한다. 핵심 아이디어는 “표현 공간에서 예측한다”는 점으로, 픽셀‑레벨 재구성 방식이 초래하는 불필요한 계산 비용과 의미 없는 저차원 특징을 회피한다. 이를 위해 저자는 에너지 기반 모델(EBM) 관점에서 손실을 정의하고, 예측 오차를 에너지로, 정규화 항을 추가해 표현 붕괴를 방지한다.

세 가지 학습 설정이 순차적으로 소개된다. 첫 번째는 이미지‑JEPA로, 두 개의 랜덤 뷰를 인코더에 통과시켜 얻은 임베딩 사이의 L2 거리를 최소화하고, VICReg 혹은 SIGReg을 통해 분산·공분산 혹은 가우시안성 검정을 적용한다. 두 번째는 비디오‑JEPA로, 시간적 리셉티브 필드를 가진 인코더와 시퀀스 예측기를 사용해 t 시점의 표현을 입력으로 t+1 시점의 목표 표현을 예측한다. 여기서는 다단계(k‑step) 롤아웃 손실을 도입해 훈련 단계와 추론 단계 사이의 exposure bias를 감소시킨다. 세 번째는 행동‑조건 비디오‑JEPA(액션‑조건 세계 모델)로, 행동 인코더가 제공하는 제어 벡터를 예측기에 결합해 미래 상태를 예측한다. 이 단계에서는 추가 정규화 항인 temporal similarity loss와 inverse dynamics loss를 도입해 행동-상태 간 일관성을 강화한다.

정규화 기법에 대한 비교 실험도 눈에 띈다. VICReg은 분산·공분산 두 항을 사용해 각 차원의 스케일과 상관관계를 조절하지만, 하이퍼파라미터 민감도가 높다. 반면 SIGReg은 1차원 랜덤 프로젝션에 대한 가우시안성 검정을 기반으로 하여 단일 λ 파라미터만 필요하고, 다양한 설정에서 안정적인 학습을 보인다. 실험 결과, CIFAR‑10에서 SIGReg이 VICReg보다 하이퍼파라미터 탐색 범위가 넓음에도 일관된 성능을 유지한다는 점이 강조된다.

예측 성능 평가에서는 Moving MNIST을 사용해 1‑step, 4‑step, 8‑step 롤아웃을 시각화했으며, 다단계 손실을 적용한 모델이 장기 예측에서 흐릿해지는 현상이 크게 감소한다. 플래닝 실험에서는 Two Rooms 환경에서 목표 관측을 입력으로 MPPI 기반 샘플링 플래너가 행동 시퀀스를 최적화하도록 설계되었다. 에너지 함수는 목표 상태와 현재 롤아웃 상태 간 L2 거리를 누적해 정의되며, 최소 에너지 경로가 목표 도달을 의미한다. 97 %의 성공률은 비교적 간단한 환경에서도 학습된 잠재 역학 모델이 충분히 정확함을 증명한다.

코드 구조는 크게 인코더, 예측기, 정규화 모듈, 플래너 네 개의 재사용 가능한 컴포넌트로 나뉜다. 각 모듈은 PyTorch 기반으로 구현돼 있으며, 하이퍼파라미터 테이블과 상세 주석이 제공돼 초보 연구자도 빠르게 실험을 재현할 수 있다. 또한, 학습 파이프라인은 단일 GPU에서 4‑8 GB 메모리 사용량으로 수시간 내에 수렴하도록 최적화돼 있다.

전반적으로 EB‑JEPA는 이론적 배경(EBM, 정규화 기반 붕괴 방지)과 실용적 구현(모듈화, 멀티스텝 손실, 행동‑조건 플래닝)을 성공적으로 결합했다. 이는 최신 자기지도 학습 기법을 비디오와 강화학습 영역으로 확장하려는 연구자들에게 중요한 베이스라인이 될 것으로 기대된다.