Joint Embedding Predictive Architecture (JEPA)

공동 임베딩 예측 아키텍처(JEPA)는 자기 지도 학습 설계의 하나로, 네트워크가 입력의 보이는 부분을 바탕으로 숨겨진 부분의 추상 표현을 예측하며 학습합니다. 일부 영역이 마스킹된 이미지가 주어지면 JEPA는 해당 영역의 픽셀을 복원하지 않고 임베딩을 예측합니다. 따라서 학습 신호는 객체의 형태와 배치처럼 예측 가능한 의미 정보에 집중되고, 나뭇잎의 질감이나 센서 노이즈처럼 예측할 수 없는 세부 정보는 건너뜁니다. Yann LeCun은 A Path Towards Autonomous Machine Intelligence라는 2022년 입장 논문에서 월드 모델의 구성 요소로 이 아키텍처를 제안했으며, Meta의 I-JEPA와 V-JEPA 모델은 이를 이미지와 비디오에 적용했습니다.

JEPA 작동 방식#

JEPA는 세 개의 네트워크로 구성됩니다. 일반적으로 비전 트랜스포머(ViT)인 컨텍스트 인코더는 보이는 패치를 임베딩합니다. 타깃 인코더는 전체 입력을 임베딩하며, 숨겨진 위치의 임베딩이 타깃이 됩니다. 더 좁은 트랜스포머인 예측기는 컨텍스트 임베딩과 각 타깃 위치를 표시하는 마스크 토큰을 입력받아 타깃 임베딩의 추정치를 출력합니다. 손실 함수는 예측된 타깃 임베딩과 실제 타깃 임베딩 사이의 평균 L2 거리를 사용하며, 컨텍스트 인코더와 예측기에만 그래디언트가 전달됩니다. I-JEPA에서는 각 학습 이미지에서 이미지의 15~20%를 차지하는 타깃 블록 4개와 85~100%를 차지하는 컨텍스트 블록 1개를 생성합니다. 타깃과 겹치는 부분은 제거하므로 모델은 희소한 컨텍스트만으로 크고 의미 있는 영역을 추론해야 합니다.

표현 붕괴 방지#

임베딩을 예측하면 지름길이 생길 수 있습니다. 두 인코더가 모든 입력에 대해 같은 상수 벡터를 출력하면 예측이 완벽하고 손실이 0이 되는데, 이를 표현 붕괴라고 합니다. I-JEPA와 V-JEPA는 비대칭 구조로 이를 방지합니다. 타깃 인코더에는 그래디언트가 전달되지 않으며, 대신 타깃 인코더의 가중치는 컨텍스트 인코더 가중치의 지수 이동 평균(EMA)을 따릅니다. I-JEPA 학습에서 모멘텀은 0.996으로 시작해 선형적으로 증가하여 1.0에 도달합니다. 동일한 EMA 업데이트는 Ultralytics의 YOLO26 학습에도 사용됩니다. 여기서 트레이너는 감쇠율이 0.9999를 향해 증가하는 모델 가중치의 이동 평균을 유지하고, 평균화된 사본을 체크포인트로 저장합니다. Randall Balestriero와 LeCun이 2025년에 발표한 LeJEPA는 EMA 티처와 그래디언트 차단을 SIGReg라는 정규화 항으로 대체해 임베딩이 등방성 가우스 분포를 따르도록 합니다.

아래 PyTorch 루프는 컨텍스트 뷰와 타깃 뷰가 공유하는 숨겨진 8차원 원인을 사용해 합성 벡터에서 소규모 JEPA를 학습합니다.

import copy

import torch
from torch import nn

torch.manual_seed(0)
latent = torch.randn(512, 8)  # hidden cause shared by both views
x, y = latent @ torch.randn(8, 64), latent @ torch.randn(8, 64)  # context view x, target view y

encoder = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32))
target_encoder = copy.deepcopy(encoder).requires_grad_(False)  # updated only by EMA
predictor = nn.Linear(32, 32)
optimizer = torch.optim.Adam([*encoder.parameters(), *predictor.parameters()], lr=1e-3)

for step in range(1, 501):
    # Predict the target embedding from the context embedding; raw inputs are never reconstructed
    loss = nn.functional.mse_loss(predictor(encoder(x)), target_encoder(y))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    with torch.no_grad():  # EMA update of the target encoder with momentum 0.996
        for target_param, param in zip(target_encoder.parameters(), encoder.parameters()):
            target_param.lerp_(param, 0.004)
    if step in {1, 500}:
        print(f"step {step}: latent loss {loss.item():.4f}, target embedding std {target_encoder(y).std(0).mean():.3f}")

step 1: latent loss 0.5372, target embedding std 0.547과 step 500: latent loss 0.0060, target embedding std 0.595을 출력합니다. 임베딩 공간의 예측 오류는 약 90배 감소하는 동안 타깃 임베딩의 분산은 유지됩니다. 즉, 임베딩이 상수로 붕괴하지 않고 손실이 감소했습니다.

JEPA와 관련 접근법 비교#

• 마스크드 오토인코더(MAE): 같은 방식으로 패치를 마스킹하고 누락된 픽셀을 복원하므로 모델 용량을 저수준 세부 정보에 사용합니다. JEPA는 표현 공간에서 누락된 내용을 예측합니다.
• 대조 학습: 이미지의 증강된 두 뷰에 대한 임베딩을 서로 끌어당기고 다른 이미지의 임베딩과는 밀어내므로, 음성 예시와 크롭 및 색상 지터링 같은 수작업 증강이 필요합니다. I-JEPA는 둘 다 사용하지 않습니다.
• DINO와 같은 자기 증류 방법: EMA 티처를 기준으로 학생 모델을 학습하며 동일 이미지의 증강된 크롭 사이에서 임베딩을 일치시킵니다. 반면 JEPA는 한 영역의 임베딩을 다른 영역으로부터 예측하여 공간 및 시간 구조를 학습합니다.
• 생성형 월드 모델: 비디오 생성 모델처럼 픽셀 공간에서 미래 프레임을 예측합니다. JEPA 월드 모델은 미래 임베딩을 예측하고 잠재 공간에서 직접 계획할 수 있습니다.

JEPA 모델#

• I-JEPA (2023): 이미지 모델은 A100 GPU 16개를 사용해 72시간 이내에 ImageNet의 ViT-Huge/14를 학습했으며, 수작업 데이터 증강 없이도 해당 특징을 선형 분류, 객체 수 세기, 깊이 예측에 전이할 수 있습니다.
• V-JEPA (2024): 특징 예측만을 목표로 비디오 200만 개에서 학습했습니다. 가장 큰 모델인 고정된 ViT-H/16은 Kinetics-400에서 81.9%, Something-Something v2에서 72.2%, ImageNet-1k에서 77.9%를 달성합니다.
• V-JEPA 2 (2025): 인터넷 비디오 100만 시간 이상으로 사전 학습했으며, Something-Something v2에서 top-1 정확도 77.3%, Epic-Kitchens-100 동작 예측에서 recall-at-5 39.7을 달성했습니다. 동작 조건부 변형인 V-JEPA 2-AC는 DROID 데이터셋의 레이블이 없는 로봇 비디오 62시간 미만으로 후속 학습한 뒤, 두 연구소의 Franka 로봇 팔을 사용해 이미지 목표를 향해 계획하며 제로샷 방식으로 물체를 집어 옮겼습니다.