Joint Embedding Predictive Architecture (JEPA)

Kiến trúc dự đoán nhúng chung (JEPA) là một thiết kế học tự giám sát, trong đó một mạng học bằng cách dự đoán biểu diễn trừu tượng của một phần đầu vào bị che dựa trên phần nhìn thấy được. Với một ảnh có một số vùng bị che, JEPA dự đoán embedding của các vùng bị che mà không tái tạo pixel, nhờ đó tín hiệu huấn luyện tập trung vào nội dung ngữ nghĩa có thể dự đoán như hình dạng và bố cục đối tượng, đồng thời bỏ qua những chi tiết khó dự đoán như kết cấu lá hoặc nhiễu cảm biến. Yann LeCun đề xuất kiến trúc này trong bài báo quan điểm năm 2022 A Path Towards Autonomous Machine Intelligence như một thành phần xây dựng world model; các model I-JEPA và V-JEPA của Meta hiện thực hóa kiến trúc này cho ảnh và video.

Cách JEPA hoạt động#

JEPA có ba mạng. Context encoder, thường là một vision transformer (ViT), tạo embedding cho các patch nhìn thấy được. Target encoder tạo embedding cho toàn bộ đầu vào, trong đó các embedding tại vị trí bị che trở thành mục tiêu. Predictor, một transformer hẹp hơn, nhận embedding ngữ cảnh cùng các token mặt nạ đánh dấu vị trí của từng mục tiêu rồi đưa ra ước tính embedding mục tiêu. Hàm mất mát là khoảng cách L2 trung bình giữa embedding mục tiêu được dự đoán và thực tế; chỉ context encoder và predictor nhận gradient. Trong I-JEPA, mỗi ảnh huấn luyện tạo ra bốn khối mục tiêu, mỗi khối chiếm 15–20% ảnh, và một khối ngữ cảnh chiếm 85–100%; mọi phần chồng lấn với mục tiêu đều bị loại bỏ, buộc model suy luận các vùng ngữ nghĩa lớn từ ngữ cảnh thưa.

Ngăn hiện tượng sụp đổ biểu diễn#

Việc dự đoán embedding mở ra một lối tắt: nếu cả hai encoder đều xuất ra cùng một vector hằng số cho mọi đầu vào, dự đoán sẽ hoàn hảo và loss bằng 0, một lỗi được gọi là sụp đổ biểu diễn. I-JEPA và V-JEPA ngăn hiện tượng này bằng tính bất đối xứng. Không có gradient nào truyền vào target encoder; thay vào đó, trọng số của target encoder bám theo trung bình động hàm mũ (EMA) của trọng số context encoder, với hệ số momentum bắt đầu ở 0.996 và tăng tuyến tính lên 1.0 trong quá trình huấn luyện I-JEPA. Cùng cách cập nhật EMA này cũng giúp việc YOLO26 được huấn luyện trong Ultralytics diễn ra mượt hơn, trong đó trainer duy trì trung bình động của trọng số model, với hệ số decay tăng dần tới 0.9999, rồi lưu bản sao của model với các trọng số đã được lấy trung bình làm checkpoint. LeJEPA, được Randall Balestriero và LeCun công bố năm 2025, thay thế teacher EMA và stop-gradient bằng một hạng tử chính quy hóa có tên SIGReg, giúp đẩy các embedding về phía một phân phối Gaussian đẳng hướng.

Vòng lặp PyTorch dưới đây huấn luyện một JEPA nhỏ trên các vector tổng hợp, trong đó góc nhìn ngữ cảnh và góc nhìn mục tiêu cùng chia sẻ một nguyên nhân ẩn 8 chiều.

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}")

Kết quả in ra là step 1: latent loss 0.5372, target embedding std 0.547 và step 500: latent loss 0.0060, target embedding std 0.595. Sai số dự đoán trong không gian embedding giảm khoảng 90 lần trong khi độ phân tán của embedding mục tiêu vẫn ổn định, cho thấy loss đã giảm mà embedding không sụp đổ thành một hằng số.

JEPA so với các phương pháp liên quan#

• Autoencoder có mặt nạ (MAE): Che patch theo cách tương tự rồi tái tạo các pixel bị thiếu, khiến model tiêu tốn năng lực vào chi tiết cấp thấp. JEPA dự đoán nội dung bị thiếu trong không gian biểu diễn.
• Học tương phản: Kéo các embedding của hai góc nhìn đã tăng cường từ cùng một ảnh lại gần nhau và đẩy embedding của các ảnh khác ra xa, do đó cần ví dụ âm và các phép tăng cường được thiết kế thủ công như cắt ảnh và thay đổi màu. I-JEPA không cần cả hai.
• Các phương pháp tự chưng cất như DINO: Cũng huấn luyện student dựa trên teacher EMA, khớp embedding giữa các crop đã tăng cường của cùng một ảnh. Thay vào đó, JEPA dự đoán embedding của một vùng từ một vùng khác, qua đó học cấu trúc không gian và thời gian.
• World model tạo sinh: Dự đoán các khung hình tương lai trong không gian pixel, như các model tạo video. World model JEPA dự đoán embedding tương lai và có thể lập kế hoạch trực tiếp trong không gian tiềm ẩn.

Các model JEPA#

• I-JEPA (2023): Model ảnh huấn luyện ViT-Huge/14 trên ImageNet bằng 16 GPU A100 trong chưa đầy 72 giờ; các đặc trưng của model chuyển giao được sang phân loại tuyến tính, đếm đối tượng và dự đoán độ sâu mà không cần các phép tăng cường dữ liệu được thiết kế thủ công.
• V-JEPA (2024): Được huấn luyện trên 2 triệu video, chỉ với mục tiêu dự đoán đặc trưng; model lớn nhất, một ViT-H/16 được giữ cố định, đạt 81.9% trên Kinetics-400, 72.2% trên Something-Something v2 và 77.9% trên ImageNet-1k.
• V-JEPA 2 (2025): Được tiền huấn luyện trên hơn 1 triệu giờ video internet, model đạt độ chính xác top-1 77.3% trên Something-Something v2 và recall-at-5 39.7 trên tác vụ dự đoán hành động Epic-Kitchens-100. Biến thể có điều kiện hành động, V-JEPA 2-AC, được huấn luyện bổ sung trên chưa đầy 62 giờ video robot chưa gán nhãn từ dataset DROID, sau đó thực hiện gắp và đặt vật thể zero-shot trên cánh tay robot Franka tại hai phòng thí nghiệm bằng cách lập kế hoạch hướng đến mục tiêu hình ảnh.