Riepilogo Tecnico: UWM-JEPA

Enunciato del Problema

Le Architetture Predittive di Incorporamento Congiunto (JEPAs) prevedono stati futuri nello spazio latente anziché ricostruire osservazioni grezze, separando efficacemente l'apprendimento della rappresentazione dalla decodifica generativa. Tuttavia, in ambienti parzialmente osservabili, una rappresentazione latente vettoriale standard manca della struttura interna necessaria per trasportare la "credenza" su continuazioni nascoste. Quando la storia delle osservazioni correnti non determina univocamente lo stato nascosto, il futuro non è un singolo punto nello spazio delle rappresentazioni, ma una distribuzione sugli stati nascosti compatibili.

I predittori vettoriali standard addestrati contro target a Media Mobile Esponenziale (EMA) possono minimizzare la perdita, ma la funzione di perdita da sola non specifica quali incertezze debbano essere preservate, quali debbano essere mediate o come le azioni controfattuali debbano spostare lo stato di credenza. Di conseguenza, le JEPAs a latente vettoriale spesso falliscono nel mantenere dinamiche sensibili alle azioni o nel preservare informazioni rilevanti per il compito durante un "rollout cieco" (immaginazione di stati futuri senza nuove osservazioni).

Metodologia

Gli autori introducono UWM-JEPA (Unitary World Model JEPA), una modifica del framework JEPA che tratta il latente come uno stato di credenza esplicito anziché come un punto deterministico.

1. Latente a Stato di Credenza

Invece di un vettore, l'encoder online $E_\theta$ mappa il contesto osservato su una matrice densità $\rho_t$ in uno spazio latente bipartito $\mathcal{H}_T = \mathcal{H}_S \otimes \mathcal{H}_E$ (Sistema $\otimes$ Ambiente).

• Dimensioni: Il fattore sistema ($d_{sys}=8$) è letto dalla funzione di perdita e dalle sonde; il fattore ambiente ($d_{env}=2$) agisce come registro per trasportare modi nascosti e correlazioni sistema-ambiente.
• Vincoli: La matrice densità è vincolata ad essere hermitiana, semi-definita positiva e a traccia unitaria. Questi vincoli sono imposti per costruzione: lo stato iniziale è massimamente misto, gli aggiornamenti delle osservazioni applicano operatori di misura positivi e i passaggi di previsione usano coniugazione unitaria.

2. Predittore Unitario

Il predittore avanza lo stato latente tramite coniugazione unitaria:
$$\hat{\rho}_{t+k|t} = U_k \rho_t (U_k)^\dagger$$
dove $U = \exp(-iH\Delta t)$ e $H$ è un generatore hermitiano.

• Condizionamento all'Azione: Il generatore è lineare nell'azione $a$: $H(a) = H_0 + a H_1$. Ciò permette alla traiettoria unitaria di essere guidata dalle azioni.
• Obiettivo di Addestramento: Il modello utilizza l'impalcatura JEPA standard: un encoder target EMA $E_\xi$ con stop-gradient e una perdita di corrispondenza nello spazio latente valutata dopo la traccia parziale sul fattore sistema e un testata di proiezione.
• Target Controfattuali: Per garantire che il predittore impari a utilizzare il canale delle azioni, il target di addestramento è generato da un rollout di simulatore sotto una sequenza di azioni controfattuali $\tilde{a}$, anziché la traiettoria osservata con insegnamento forzato.

3. Garanzia Teorica

L'affermazione teorica centrale è la non-dimenticanza dello stato congiunto. Poiché la coniugazione unitaria è un'isometria, essa preserva esattamente lo spettro (autovalori), la purezza e l'entropia di von Neumann della matrice densità dello stato congiunto sistema-ambiente durante il rollout cieco.

• Limitazione: Questa garanzia si applica allo stato congiunto completo. Lo stato ridotto del sistema (ottenuto tramite traccia parziale) non è unitario, il che significa che le informazioni possono spostarsi nelle correlazioni sistema-ambiente e diventare inaccessibili alle sonde lineari, anche se lo stato congiunto rimane non dissipativo.

Risultati Chiave

1. La Sensibilità all'Azione Richiede Target Controfattuali

• Fallimento con Insegnamento Forzato: Quando addestrato con target a insegnamento forzato (corrispondenti al futuro osservato), il canale delle azioni diventa inerte. Il rapporto tra il termine dipendente dall'azione e il termine delle dinamiche di base collassa a $\|H_1\|/\|H_0\| \approx 0.03$. Il predittore impara a corrispondere al target indipendentemente dall'azione perché il target contiene già l'effetto dell'azione.
• Successo Controfattuale: Sostituire il target con un rollout di simulatore sotto azioni controfattuali campionate ripristina la sensibilità all'azione, con $\|H_1\|/\|H_0\| = 1.00 \pm 0.15$.
• Controlli di Perturbazione: UWM-JEPA addestrato con target controfattuali (UWM-JEPA-CF) mostra un degrado monotono dell'accuratezza di previsione man mano che le azioni vengono perturbate (invertite, mescolate, negate), mentre un LSTM-JEPA-CF con parametri corrispondenti collassa all'accuratezza della classe maggioritaria (0.53) in tutte le condizioni.

2. Rollout Cieco e Mantenimento a Breve Orizzonte

• Indicatore di Velocità Nascosta: In un compito che richiede una simulazione in avanti di 5 passaggi per prevedere il segno della velocità nascosta (con l'osservazione target mascherata), UWM-JEPA-CF raggiunge un'accuratezza del 0.77. Il LSTM-JEPA-CF corrispondente fallisce, prevedendo la classe maggioritaria per tutti gli esempi.
• Mantenimento delle Informazioni ($\Delta R^2$): Misurando la perdita di informazioni rispetto a una sonda teacher durante il rollout cieco:
  • A orizzonti brevi ($k=1, 3$), UWM-JEPA perde meno di 10 punti di $R^2$ (rispettivamente 0.023 e 0.069).
  • La baseline LSTM-JEPA perde significativamente di più (0.413 e 0.678).
  • All'orizzonte addestrato ($k=5$), LSTM si riallinea, ma UWM mantiene prestazioni superiori nei passaggi intermedi.
  • A orizzonti lunghi ($k>10$), entrambi i modelli degradano, indicando che il mantenimento a lungo orizzonte rimane una sfida aperta.

3. La Rappresentazione del Contesto Non è il Motore

• Equivalenza dell'Encoder: Una sonda lineare tenuta da parte sulle rappresentazioni dell'encoder congelato mostra che UWM-JEPA e LSTM-JEPA performano identicamente ($R^2 = 0.901 \pm 0.032$ contro $0.894 \pm 0.021$, $p=0.70$).
• Conclusione: Il divario di prestazioni non è dovuto al fatto che UWM-JEPA abbia un encoder di contesto superiore. La separazione nasce interamente dalla capacità del predittore di mantenere la struttura dello stato di credenza durante il rollout.

Significato e Affermazioni

Il documento presenta UWM-JEPA come il primo modello del mondo della famiglia JEPA in cui il latente è trattato esplicitamente come una credenza su continuazioni nascoste (una matrice densità) e il predittore è vincolato a evolvere questa credenza senza dissipare lo spettro dello stato congiunto.

• Portata Modesta: Gli autori dichiarano esplicitamente che questo non è un pianificatore a scopo generale, un encoder universalmente più forte o una soluzione competitiva per la pianificazione a lungo orizzonte. La proprietà di "non-dimenticanza" è esatta solo per lo stato congiunto, non necessariamente per lo stato ridotto utilizzato per compiti a valle.
• Contributo Principale: Il lavoro dimostra che affinché i modelli del mondo JEPA possano immaginare in condizioni di osservabilità parziale, la geometria latente e le dinamiche del predittore contano più della sola capacità di codifica del contesto congelata.
• Implicazione Pratica: Lo studio evidenzia che l'addestramento contro target controfattuali è essenziale per imparare predittori sensibili alle azioni, una scoperta che si estende oltre la specifica parametrizzazione unitaria.

I risultati sono validati su ambienti strutturati parzialmente osservabili di tipo oscillatore e CartPole. Gli autori riconoscono che l'approccio attuale richiede l'accesso al simulatore per l'addestramento controfattuale e non si estende ancora alla previsione visiva non strutturata o alle catene di Markov nascoste.