Joint Embedding Predictive Architecture (JEPA)

Un'architettura predittiva a embedding congiunti (JEPA) è un approccio di apprendimento auto-supervisionato in cui una rete apprende predicendo la rappresentazione astratta di una parte nascosta dell'input a partire da una parte visibile. Data un'immagine con alcune regioni mascherate, una JEPA predice gli embedding delle regioni mascherate senza ricostruirne i pixel; il segnale di addestramento si concentra così sui contenuti semantici prevedibili, come la forma e la disposizione degli oggetti, e trascura i dettagli imprevedibili, come la trama delle foglie o il rumore del sensore. Yann LeCun ha proposto l'architettura nel suo documento programmatico del 2022 A Path Towards Autonomous Machine Intelligence come elemento costitutivo dei modelli del mondo; i modelli I-JEPA e V-JEPA di Meta l'hanno poi applicata alle immagini e ai video.

Come funziona JEPA#

Una JEPA è composta da tre reti. Un encoder di contesto, in genere un vision transformer (ViT), codifica le patch visibili. Un encoder target codifica l'input completo e i suoi embedding nelle posizioni nascoste diventano i target. Un predittore, un Transformer più stretto, prende gli embedding di contesto e i token maschera che indicano la posizione di ogni target e ne stima gli embedding. La funzione di perdita è la distanza L2 media tra gli embedding predetti e quelli effettivi; solo l'encoder di contesto e il predittore ricevono gradienti. In I-JEPA, ogni immagine di addestramento produce quattro blocchi target, ciascuno corrispondente al 15–20% dell'immagine, e un blocco di contesto corrispondente all'85–100%, eliminando qualsiasi sovrapposizione con i target; il modello deve quindi dedurre grandi regioni semantiche da un contesto sparso.

Evitare il collasso delle rappresentazioni#

Predire gli embedding apre una scorciatoia: se entrambi gli encoder restituiscono lo stesso vettore costante per ogni input, la predizione è perfetta e la perdita è zero; questo problema è noto come collasso delle rappresentazioni. I-JEPA e V-JEPA lo prevengono introducendo un'asimmetria. Nessun gradiente viene propagato all'encoder target; i suoi pesi seguono invece una media mobile esponenziale (EMA) dei pesi dell'encoder di contesto, con un momentum che parte da 0.996 e sale linearmente fino a 1.0 durante l'addestramento di I-JEPA. Lo stesso aggiornamento EMA stabilizza il training di YOLO26 in Ultralytics: il trainer mantiene una media mobile dei pesi del modello, con un decadimento che aumenta gradualmente fino a 0.9999, e salva la copia mediata come checkpoint. LeJEPA, pubblicata nel 2025 da Randall Balestriero e LeCun, sostituisce l'insegnante EMA e l'arresto del gradiente con un regolarizzatore chiamato SIGReg, che spinge gli embedding verso una distribuzione gaussiana isotropa.

Il ciclo PyTorch qui sotto addestra una piccola JEPA su vettori sintetici, nei quali una vista di contesto e una vista target condividono una causa nascosta a 8 dimensioni.

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

Stampa step 1: latent loss 0.5372, target embedding std 0.547 e step 500: latent loss 0.0060, target embedding std 0.595. L'errore di predizione nello spazio degli embedding diminuisce di circa 90 volte, mentre la dispersione degli embedding target resta stabile; la perdita si è quindi ridotta senza che gli embedding collassassero in una costante.

JEPA a confronto con gli approcci correlati#

• Autoencoder mascherati (MAE): mascherano le patch allo stesso modo e ricostruiscono i pixel mancanti, consumando capacità del modello sui dettagli di basso livello. JEPA predice i contenuti mancanti nello spazio delle rappresentazioni.
• Apprendimento contrastivo: avvicina gli embedding di due viste aumentate di un'immagine e allontana quelli di altre immagini, quindi richiede esempi negativi e trasformazioni manuali come ritagli e variazioni dei colori. I-JEPA non usa né gli uni né le altre.
• Metodi di auto-distillazione come DINO: addestrano anch'essi uno studente confrontandolo con un insegnante EMA e abbinano gli embedding di ritagli aumentati della stessa immagine. JEPA predice invece l'embedding di una regione a partire da un'altra, apprendendo così la struttura spaziale e temporale.
• Modelli generativi del mondo: predicono i frame futuri nello spazio dei pixel, come fanno i generatori video. Un modello del mondo JEPA predice gli embedding futuri e può pianificare direttamente in quello spazio latente.

Modelli JEPA#

• I-JEPA (2023): il modello per immagini ha addestrato un ViT-Huge/14 su ImageNet con 16 GPU A100 in meno di 72 ore; le sue feature possono essere trasferite alla classificazione lineare, al conteggio degli oggetti e alla predizione della profondità senza aumenti dei dati definiti manualmente.
• V-JEPA (2024): addestrato su 2 milioni di video usando la predizione delle feature come unico obiettivo, il suo modello più grande, un ViT-H/16 mantenuto congelato, raggiunge l'81.9% su Kinetics-400, il 72.2% su Something-Something v2 e il 77.9% su ImageNet-1k.
• V-JEPA 2 (2025): preaddestrata su oltre 1 milione di ore di video provenienti da Internet, raggiunge un'accuratezza top-1 del 77.3% su Something-Something v2 e un recall-at-5 di 39.7 sull'anticipazione delle azioni di Epic-Kitchens-100. La sua variante condizionata dalle azioni, V-JEPA 2-AC, è stata sottoposta a post-addestramento su meno di 62 ore di video robotici senza etichette del dataset DROID; poi ha raccolto e posizionato oggetti zero-shot con bracci robotici Franka in due laboratori, pianificando il movimento verso obiettivi nell'immagine.