Riepilogo Tecnico: Quando LeJEPA Impara un Modello del Mondo?

Enunciato del Problema

Il documento affronta un divario fondamentale nell'apprendimento auto-supervisionato (SSL): mentre le Architetture Predittive a Incorporamento Congiunto (JEPAs) hanno dimostrato successo empirico nell'apprendimento di rappresentazioni per visione, video e pianificazione, non esiste alcuna garanzia teorica che queste rappresentazioni costituiscano un "Modello del Mondo". Nello specifico, rimane incerto in quali condizioni una rappresentazione appresa recuperi fedelmente la struttura latente sottostante del mondo (ad esempio, posizione, velocità, colore) piuttosto che intrecciare questi fattori.

Una rappresentazione che mescola i veri gradi di libertà (ad esempio, mescolando la posizione di un oggetto con il suo colore) può performare bene su compiti ristretti, ma fallisce quando il mondo cambia o quando viene utilizzata per la generalizzazione composizionale. La domanda centrale è: Quando una rappresentazione appresa recupera linearmente le variabili latenti del mondo? Gli autori definiscono questa proprietà come identificabilità lineare, una condizione necessaria per un'analisi lineare affidabile e una pianificazione ottimale nello spazio latente.

Metodologia e Quadro Teorico

Il Modello del Mondo

Gli autori formalizzano il "Mondo" come un sistema con variabili latenti $z \in \mathbb{R}^n$ non osservate. I dati osservati $x$ sono generati da una funzione di miscelazione sconosciuta, potenzialmente altamente non lineare $g$, tale che $x = g(z)$. L'obiettivo di apprendimento coinvolge coppie positive $(z, z')$ generate da un processo stazionario con transizioni indipendenti e rumorose additivamente.

Il documento si concentra su una specifica classe di mondi in cui:

1. Indipendenza: Le variabili latenti e le loro transizioni sono indipendenti tra le dimensioni.
2. Stazionarietà: La distribuzione marginale delle variabili latenti rimane costante tra le visualizzazioni.
3. Rumore Additivo: Le transizioni seguono $z' = m(z) + \eta$, dove $\eta$ è rumore indipendente.

Crucialmente, gli autori specializzano questa classe nel Mondo Gaussiano, dove $z \sim \mathcal{N}(0, I_n)$. Sotto queste assunzioni, le dinamiche di transizione corrispondono a un processo di Ornstein–Uhlenbeck (OU): $z' = \rho z + \sqrt{1-\rho^2}\eta$.

L'Apprendista: LeJEPA

L'architettura di apprendimento è LeJEPA (JEPA di LeCun), che addestra un encoder $f$ per produrre incorporamenti $h(z) = f(g(z))$ minimizzando una funzione di perdita composta da due termini:

1. Perdita di Allineamento: Minimizza la distanza tra gli incorporamenti di coppie positive, $E[\|h(z') - h(z)\|^2]$. Ciò incoraggia la rappresentazione a catturare caratteristiche lente e prevedibili.
2. SIGReg (Regolarizzazione Isotropa Gaussiana Schizzata): Regolarizza esplicitamente la distribuzione degli incorporamenti affinché sia una Gaussiana isotropa, $h(z) \sim \mathcal{N}(0, I_n)$, prevenendo il collasso della rappresentazione.

Il problema di ottimizzazione è:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{soggetto a} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Approccio Teorico

Gli autori impiegano un'analisi spettrale dell'operatore di transizione $T$, definito come $(Th)(z) = E[h(z')|z]$.

• Sviluppo di Hermite: Per latenti gaussiani, le autofunzioni dell'operatore di transizione sono i polinomi di Hermite. Gli autovalori decadono geometricamente come $\rho^d$, dove $d$ è il grado del polinomio.
• Penalità Spettrale: La perdita di allineamento penalizza i componenti di grado superiore (non lineari) più pesantemente rispetto ai componenti lineari, poiché si decorrelano più rapidamente sotto la transizione.
• Teoria di Sturm-Liouville: Per distribuzioni generali (non gaussiane), le autofunzioni sono caratterizzate dalla teoria di Sturm-Liouville. La prima autofunzione non costante è sempre monotona, ma la linearità è garantita solo se la distribuzione è Gaussiana.

Contributi Chiave

1. Direzione Diretta: Identificabilità Lineare per Mondi Gaussiani

Il Teorema 1 dimostra che se le variabili latenti sono gaussiane e la transizione è un processo OU, qualsiasi minimizzatore dell'obiettivo LeJEPA (soggetto a vincoli di gaussianità) deve essere una trasformazione ortogonale lineare delle vere latenti: $h(z) = Qz$ per qualche matrice ortogonale $Q$.

• Meccanismo: La prova si basa su una decomposizione spettrale in cui qualsiasi non linearità riduce strettamente la correlazione tra le coppie positive. Poiché la perdita di allineamento mira a massimizzare questa correlazione, l'ottimo globale forza la rappresentazione a essere puramente lineare (polinomi di Hermite di grado 1).
• Risultato: La rappresentazione appresa recupera le variabili latenti a meno di una rotazione globale, e le dinamiche di transizione sono recuperate esattamente nello spazio latente.

2. Direzione Inversa: Unicità della Gaussiana

Il Teorema 2 stabilisce che la distribuzione Gaussiana è la unica distribuzione latente per cui vale l'identificabilità lineare in questa classe di mondi.

• Meccanismo: Utilizzando la teoria di Sturm-Liouville, gli autori mostrano che richiedere che la prima autofunzione dell'operatore di transizione sia affine (lineare) forza la funzione di punteggio della distribuzione latente ad essere lineare. Questa condizione caratterizza univocamente la distribuzione Gaussiana.
• Implicazione: Se le variabili latenti non sono gaussiane, LeJEPA non può garantire l'identificabilità lineare; la rappresentazione appresa sarà necessariamente una trasformazione non lineare delle latenti.

3. Identificabilità Approssimata

Il Teorema 3 fornisce un limite quantitativo sull'errore di recupero quando gli obiettivi sono soddisfatti solo approssimativamente (cioè, allineamento imperfetto e sbiancamento imperfetto).

• L'errore scala in modo elegante con il divario di allineamento $\delta$ e l'errore di sbiancamento $\epsilon$.
• Il limite evidenzia che il divario di allineamento è il fattore dominante nell'errore di recupero, mentre lo sbiancamento approssimato è relativamente "gratuito" da ottenere.

4. Pianificazione Ottimale nello Spazio Latente

Il Teorema 4 collega l'identificabilità lineare alla teoria del controllo. Dimostra che se l'encoder è linearmente identificabile ($h(z) = Qz$), allora qualsiasi problema di controllo ottimale con costi invarianti rispetto a trasformazioni ortogonali (ad esempio, distanza euclidea verso un obiettivo) può essere risolto in modo ottimale nello spazio latente appreso. Le azioni ottimali e le funzioni valore nello spazio latente mappano direttamente sul mondo reale, richiedendo nessuna modifica alla funzione di costo.

Risultati Sperimentali

Gli autori validano la loro teoria attraverso diversi scenari sperimentali:

1. Simulazioni 2D: Utilizzando varie funzioni di miscelazione non lineari (spirale, taglio, RealNVP), LeJEPA recupera con successo le variabili latenti a meno di rotazione quando le latenti sono gaussiane, confermando il Teorema 1.
2. Scalabilità ad Alte Dimensioni: Esperimenti fino a $N=1024$ dimensioni mostrano che i regolarizzatori basati su statistiche di batch (SIGReg, VICReg) mantengono un'identificabilità lineare quasi perfetta ($R^2 > 0.999$), mentre i metodi contrastivi (InfoNCE) degradano su larga scala a causa di disallineamenti nella larghezza del kernel.
3. Ablazioni Distribuzionali: Un'analisi della forma della distribuzione latente (da code pesanti a uniforme) dimostra che il recupero lineare raggiunge un picco netto alla Gaussiana ($\alpha=2$) e degrada per distribuzioni non gaussiane, validando il Teorema 2.
4. Controllo Robotico Basato su Pixel: Utilizzando il DeepMind Control Suite Reacher, gli autori confrontano l'addestramento su transizioni sintetiche gaussiane (OU) rispetto a traiettorie RL reali.
   • Condizione OU: Raggiunge un'alta identificabilità lineare ($R^2 \approx 0.95$).
   • Condizione Traiettoria: Le politiche del mondo reale inducono marginali non gaussiani e anisotropi, risultando in un'identificabilità significativamente inferiore ($R^2 < 0.5$).
5. Validazione della Pianificazione: Nel compito Reacher, l'interpolazione a linea retta nello spazio latente appreso dell'encoder addestrato su gaussiane decodifica in percorsi quasi rettilinei nello spazio delle giunture (pianificazione ottimale). Al contrario, l'encoder addestrato su traiettorie produce percorsi curvi e subottimali, confermando il Teorema 4.

Significato e Affermazioni

Il documento afferma di fornire il primo risultato di identificabilità per le JEPAs, trasformando una ricetta empiricamente di successo in una garanzia matematica.

• Garanzia Fondamentale: Il lavoro stabilisce che LeJEPA recupera provatamente la struttura del mondo (variabili latenti e dinamiche) sotto condizioni specifiche e ben definite (latenti gaussiani, rumore additivo stazionario).
• Ruolo della Gaussianità: Il documento inverte la narrazione classica dell'ICA lineare. Nell'ICA lineare, la gaussianità è la distribuzione in cui la separazione delle sorgenti fallisce. In questo contesto non lineare e auto-supervisionato, la gaussianità è la unica distribuzione che abilita l'identificabilità lineare.
• Implicazioni Pratiche: La teoria spiega perché le strategie di esplorazione che approssimano camminate casuali isotrope (mantenendo i dati nel regime gaussiano) sono cruciali per l'apprendimento di modelli del mondo utili, mentre le politiche orientate agli obiettivi (che collassano i dati in regioni a bassa entropia) possono rompere l'identificabilità.
• Fondamento per la Pianificazione: Dimostrando che l'identificabilità lineare abilita una pianificazione ottimale con costi invarianti per rotazione, il documento fornisce una giustificazione teorica per l'uso di spazi latenti appresi per il controllo senza necessità di mappare indietro sulle coordinate fisiche reali.

Gli autori notano le limitazioni: la teoria assume che la dimensione dell'encoder corrisponda alla dimensione latente e si basa sull'ottimalità a livello di popolazione. Tuttavia, dimostrano che le garanzie degradano in modo elegante con campioni finiti ed errori di ottimizzazione, e verificano formalmente tutti i teoremi principali utilizzando il dimostratore di teoremi Lean 4.