Joint Embedding Predictive Architecture (JEPA)

Предиктивная архитектура с совместными эмбеддингами (JEPA) — это подход к обучению без учителя, при котором сеть учится предсказывать абстрактное представление скрытой части входных данных по видимой части. Получив изображение с замаскированными областями, JEPA предсказывает эмбеддинги этих областей, не восстанавливая пиксели. Поэтому обучающий сигнал сосредоточен на предсказуемом семантическом содержимом, таком как форма и расположение объектов, и не учитывает непредсказуемые детали, например текстуру листьев или шум датчика. Янн Лекун предложил эту архитектуру в своей программной статье 2022 года A Path Towards Autonomous Machine Intelligence как строительный блок для моделей мира, а модели Meta I-JEPA и V-JEPA воплотили её в жизнь для изображений и видео.

Как работает JEPA#

В JEPA используются три сети. Энкодер контекста, обычно визуальный трансформер (ViT), кодирует видимые патчи. Целевой энкодер кодирует весь вход, а эмбеддинги скрытых областей становятся целевыми значениями. Предиктор — более узкий трансформер — получает эмбеддинги контекста и токены маски, указывающие положение каждой целевой области, а затем выдаёт оценку целевых эмбеддингов. Функция потерь — это среднее L2-расстояние между предсказанными и фактическими целевыми эмбеддингами; только энкодер контекста и предиктор получают градиенты. В I-JEPA при обучении для каждого изображения создаются четыре целевых блока, каждый из которых охватывает 15–20% изображения, и один блок контекста, охватывающий 85–100%; области пересечения с целевыми блоками исключаются, поэтому модели приходится выводить крупные семантические области по разреженному контексту.

Как избежать коллапса представлений#

Предсказание эмбеддингов открывает возможность воспользоваться коротким путём: если оба энкодера выдают один и тот же постоянный вектор для любого входа, предсказание идеально и потери равны нулю. Этот сбой называется коллапсом представлений. I-JEPA и V-JEPA предотвращают его за счёт асимметрии. Градиенты не передаются целевому энкодеру; вместо этого его веса отслеживают экспоненциальное скользящее среднее (EMA) весов энкодера контекста, причём импульс при обучении I-JEPA линейно растёт с 0.996 до 1.0. То же обновление EMA сглаживает обучение YOLO26 в Ultralytics: тренер поддерживает скользящее среднее весов модели с коэффициентом затухания, постепенно увеличивающимся до 0.9999, и сохраняет усреднённую копию в качестве контрольной точки. LeJEPA, опубликованная в 2025 году Рэндаллом Балестриеро и Лекуном, заменяет учителя EMA и остановку градиента на регуляризатор SIGReg, который приближает распределение эмбеддингов к изотропному гауссовскому распределению.

Приведённый ниже цикл на PyTorch обучает небольшую JEPA на синтетических векторах, в которых представления контекста и целевого объекта имеют общую скрытую причину размерностью 8.

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): модель изображений обучила ViT-Huge/14 на ImageNet менее чем за 72 часа с помощью 16 GPU A100; её признаки переносятся на линейную классификацию, подсчёт объектов и предсказание глубины без аугментаций данных, созданных вручную.
• V-JEPA (2024): обученная на 2 миллионах видеозаписей только с предсказанием признаков в качестве целевой задачи, её крупнейшая модель — замороженная ViT-H/16 — достигает 81.9% на Kinetics-400, 72.2% на Something-Something v2 и 77.9% на ImageNet-1k.
• V-JEPA 2 (2025): Модель, предварительно обученная на видео из интернета общей продолжительностью более 1 млн часов, достигает точности top-1 77.3% на наборе Something-Something v2 и значения recall-at-5 39.7 в задаче прогнозирования действий на Epic-Kitchens-100. Вариант V-JEPA 2-AC, обусловленный действиями, прошёл дообучение с использованием менее 62 часов неразмеченного видео с роботами из набора данных DROID, а затем в двух лабораториях без дополнительных примеров захватывал и размещал объекты с помощью роботизированных манипуляторов Franka, планируя действия с опорой на целевые изображения.