Riepilogo Tecnico: Sub-JEPA

Enunciato del Problema

Le Architetture Predittive a Incorporamento Congiunto (JEPAs) offrono un quadro per l'apprendimento di modelli del mondo prevedendo future rappresentazioni latenti senza ricostruire i pixel grezzi. Tuttavia, l'addestramento di questi modelli end-to-end affronta un compromesso bias-varianza critico. Senza vincoli sufficienti, l'encoder collassa in soluzioni banali (mappando tutti gli input su embedding identici); al contrario, vincoli eccessivi sopprimono la ricchezza delle rappresentazioni apprese.

Un lavoro recente, LeWorldModel (LeWM), ha affrontato questo problema applicando un prior gaussiano isotropo agli embedding latenti, consentendo un addestramento end-to-end stabile. Tuttavia, gli autori identificano una limitazione fondamentale in questo approccio: i compiti di controllo naturale risiedono tipicamente su varietà a bassa dimensionalità all'interno di spazi ambientali ad alta dimensionalità. Imporre un prior gaussiano isotropo globale direttamente nello spazio ambientale ad alta dimensionalità introduce un bias eccessivamente forte che non corrisponde alla geometria intrinseca della dinamica sottostante, potenzialmente ostacolando le prestazioni su compiti con bassa dimensionalità intrinseca.

Metodologia: Sub-JEPA

Gli autori propongono Sub-JEPA (Regolarizzazione Gaussiana in Sottospazio), un metodo che cerca un punto operativo favorevole sulla frontiera bias-varianza spostando i vincoli gaussiani dallo spazio ambientale in molteplici sottospazi a bassa dimensionalità.

Meccanismo Principale

1. Proiezione su Sottospazio Ortogonale: Invece di regolarizzare l'intero vettore latente $z$ a $D$ dimensioni, il metodo proietta $z$ su $K$ matrici di proiezione random, fissate e a righe ortonormali $\{P_k\}$. Ogni proiezione mappa il vettore latente in un sottospazio a dimensionalità inferiore $d_s$ (dove $d_s \approx D/K$).
   $$z^{(k)} = P_k z$$
2. Regolarizzazione Gaussiana Multi-Sottospazio: Il vincolo gaussiano viene applicato indipendentemente all'interno di ciascuno dei $K$ sottospazi.
   • All'interno di ciascun sottospazio $k$, gli embedding proiettati sono ulteriormente proiettati su $M$ vettori unitari casuali.
   • Viene calcolata la statistica di normalità di Epps–Pulley per queste proiezioni 1D per imporre la gaussianità.
   • La perdita totale di regolarizzazione $L_{reg}$ è la media di queste statistiche su tutti i sottospazi e le direzioni casuali.
3. Obiettivo di Addestramento: Il modello viene addestrato end-to-end utilizzando una perdita combinata:
   $$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$$
   dove $L_{pred}$ è l'errore di previsione dello stato latente (MSE) e $\lambda$ è il peso della regolarizzazione.

Scelte Progettuali

• Proiezioni Fisse: Le matrici di proiezione sono inizializzate in modo ortogonale e mantenute fisse. Ciò impedisce al regolarizzatore di adattarsi alla distribuzione latente in evoluzione, garantendo che il vincolo rimanga stabile e indipendente.
• Ortogonalità: L'uso di proiezioni ortogonali assicura che i sottospazi forniscano visioni non ridondanti e bilanciate della rappresentazione latente, preservando la struttura geometrica.

Contributi Chiave

1. Identificazione del Bias: Il documento identifica che il prior gaussiano isotropo globale in LeWM introduce un bias eccessivo per compiti con bassa dimensionalità intrinseca, non corrispondendo alla geometria sottostante del compito.
2. Quadro Sub-JEPA: Gli autori propongono un metodo che rilassa il prior globale applicando la regolarizzazione gaussiana su molteplici viste proiettate casuali, a bassa dimensionalità e a righe ortonormali. Ciò preserva l'effetto anti-collasso consentendo allo spazio latente di conformarsi meglio alla dimensionalità intrinseca della dinamica.
3. Validazione Empirica: Esperimenti estesi su quattro ambienti di controllo continuo dimostrano che Sub-JEPA supera costantemente LeWM. I guadagni di prestazioni si sono rivelati correlati con riduzioni del rank effettivo degli embedding appresi, confermando che il metodo rispetta meglio la dimensionalità intrinseca.

Risultati Sperimentali

Il metodo è stato valutato su quattro benchmark di controllo continuo: Two-Room (bassa dimensionalità intrinseca), Reacher, PushT e OGB-Cube (3D visivamente ricco).

• Prestazioni di Pianificazione: Sub-JEPA ha superato costantemente LeWM in tutti gli ambienti. Il miglioramento è stato più pronunciato su Two-Room, dove il prior gaussiano a spazio completo di LeWM era più disallineato rispetto alla bassa dimensionalità del compito. Sub-JEPA ha inoltre superato PLDM (una baseline multi-obiettivo) e DINO-WM (una baseline con encoder fisso) su diversi compiti senza fare affidamento su backbones visivi pre-addestrati.
• Analisi del Rank Effettivo: È stata osservata una correlazione diretta tra la riduzione del rank effettivo (da LeWM a Sub-JEPA) e l'aumento del tasso di successo nella pianificazione. I compiti con maggiori compressioni di rank (Two-Room, OGB-Cube) hanno mostrato i maggiori guadagni di prestazioni.
• Decodificabilità Fisica: Esperimenti di probing sull'ambiente PushT hanno mostrato che le rappresentazioni di Sub-JEPA hanno preservato le informazioni sullo stato fisico (posizione dell'agente, posizione del blocco) tanto quanto o meglio di LeWM, con una decodificabilità non lineare migliorata.
• Qualità della Traiettoria Latente:
  • Coerenza Temporale: Le visualizzazioni hanno mostrato che Sub-JEPA produce traiettorie latenti temporalmente più coerenti rispetto all'evoluzione meno organizzata in LeWM.
  • Raddrizzamento del Percorso: Sub-JEPA ha raggiunto una maggiore "linearità" nei vettori di velocità latente, indicando dinamiche più lisce e che evolvono in modo più lineare, cruciale per la pianificazione.
  • Stabilità a Lungo Orizzonte: Nelle simulazioni in anello aperto, Sub-JEPA ha mantenuto la geometria della scena e ha mostrato meno deriva spaziale su orizzonti lunghi rispetto a LeWM.

Significato e Affermazioni

Il documento afferma che Sub-JEPA offre una soluzione semplice ma efficace al compromesso bias-varianza nei modelli del mondo basati su JEPA. Spostando la regolarizzazione gaussiana in sottospazi, il metodo:

• Rilassa il prior isotropo globale senza perdere i benefici anti-collasso.
• Corrisponde meglio alla dimensionalità intrinseca della dinamica sottostante, in particolare per compiti in cui lo spazio degli stati è a bassa dimensionalità.
• Funge da baseline solida per la futura ricerca sui modelli del mondo basati su JEPA, dimostrando che i vincoli strutturali possono essere semplificati e resi più flessibili senza sacrificare la stabilità.

Gli autori sottolineano che il loro approccio mantiene la semplicità della ricetta di addestramento di LeWM (utilizzando solo la perdita di previsione e un regolarizzatore) introducendo al contempo un prior strutturale più flessibile che produce prestazioni di pianificazione e qualità della rappresentazione superiori.