Joint Embedding Predictive Architecture (JEPA)

Joint Embedding Predictive Architecture（JEPA）は、自己教師あり学習の設計手法です。ネットワークは、入力の可視部分から、隠された部分の抽象表現を予測して学習します。一部の領域をマスクした画像を受け取ると、JEPAはピクセルを再構成せずに、マスクされた領域の埋め込みを予測します。そのため、学習信号は物体の形状や配置など予測可能な意味的内容に集中し、葉のテクスチャやセンサーノイズなど予測しにくい細部は対象外になります。Yann LeCunは、世界モデルの構成要素として、2022年の立場表明論文A Path Towards Autonomous Machine Intelligenceでこのアーキテクチャを提案しました。MetaのI-JEPAとV-JEPAは、画像と動画を対象に実用化したモデルです。

JEPAの仕組み#

JEPAは3つのネットワークで構成されます。コンテキストエンコーダーは、通常ビジョントランスフォーマー（ViT）であり、可視パッチを埋め込みに変換します。ターゲットエンコーダーは入力全体を埋め込みに変換し、隠れた位置の埋め込みをターゲットとして使います。より小型のTransformerである予測器は、コンテキスト埋め込みと各ターゲットの位置を示すマスクトークンを受け取り、ターゲット埋め込みの推定値を出力します。損失関数には予測されたターゲット埋め込みと実際のターゲット埋め込みの平均L2距離を使い、コンテキストエンコーダーと予測器だけが勾配を受け取ります。I-JEPAでは、各学習画像から、画像の15～20%をそれぞれ覆うターゲットブロックを4つと、85～100%を覆うコンテキストブロックを1つ生成し、ターゲットと重なる部分は除外します。そのため、モデルはまばらなコンテキストから広い意味領域を推測する必要があります。

表現崩壊の回避#

埋め込みを予測する方法には近道があり、両方のエンコーダーがあらゆる入力に対して同じ定数ベクトルを出力すれば予測は完璧で、損失はゼロになります。この失敗は表現崩壊と呼ばれます。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分の1に減少する一方、ターゲット埋め込みの広がりは維持されます。つまり、埋め込みが定数に崩壊することなく損失が低下しました。

JEPAと関連手法の比較#

• Masked autoencoder（MAE）：同じ方法でパッチをマスクし、欠損ピクセルを再構成するため、モデルの容量を低レベルの細部に使います。JEPAは表現空間で欠損内容を予測します。
• 対照学習：画像の2つの拡張ビューの埋め込みを近づけ、他の画像の埋め込みを遠ざけます。そのため、負例と、クロップや色の揺らぎなどの手動設計によるデータ拡張が必要です。I-JEPAではどちらも使いません。
• DINOなどの自己蒸留手法：EMA教師モデルに対して生徒モデルを学習させ、同じ画像の拡張クロップ間で埋め込みを一致させます。一方、JEPAはある領域の埋め込みから別の領域の埋め込みを予測し、空間構造と時間構造を学習します。
• 生成型世界モデル：動画生成モデルと同様に、ピクセル空間で将来のフレームを予測します。JEPAの世界モデルは将来の埋め込みを予測し、潜在空間で直接計画を立てられます。

JEPAモデル#

• I-JEPA（2023）：この画像モデルは、16基のA100 GPUを使い、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時間未満使って追加学習した後、画像目標に向けて計画することで、2つの研究室でFranka ロボットアーム上の物体を、ゼロショットでつかんで配置しました。