Sintesi Tecnica: "Il tuo JEPA Probabilistico è Segretamente un Modello di Markov Nascosto"

Definizione del Problema

Le architetture Joint-Embedding Predictive (JEPA) apprendono prevedendo rappresentazioni target da contesti osservati piuttosto che ricostruendo le osservazioni grezze. La Variational JEPA (VJEPA) estende questo concetto prevedendo una distribuzione condizionale su stati latenti futuri, e l'estensione Predictive Information Bottleneck (PIB) rende ulteriormente stocastica la rappresentazione corrente per controllare la ritenzione delle informazioni.

Nonostante le loro somiglianze strutturali con i modelli a spazio di stato, una precisa mappatura probabilistica tra le JEPA temporali completamente stocastiche e i classici Modelli di Markov Nascosti (HMM) è rimasta poco chiara. Nello specifico, rimane ambiguo se l'encoder di contesto stocastico in una JEPA agisca come un modello di emissione o come una distribuzione di filtraggio, e in quali condizioni la dinamica latente appresa soddisfi i rigorosi requisiti di coerenza di un HMM (ad esempio, coerenza di Chapman–Kolmogorov, equivalenza di filtraggio).

Metodologia

Il paper propone un'interpretazione a spazio di stato di una PIB-VJEPA completa e indicizzata nel tempo, scomponendo la sua architettura in tre ruoli analoghi a un HMM:

1. Inferenza dello Stato (Filtraggio): L'encoder di contesto stocastico $q_\theta(Z_t | X_{\le t})$ inferisce una credenza sullo stato latente corrente dalla storia delle osservazioni.
2. Transizione di Stato: Un predittore probabilistico $p_\phi(Z_{t+1} | Z_t, \xi_t)$ propaga lo stato latente in avanti.
3. Emissione: Un decoder $p_\psi(X_{t+1} | Z_{t+1})$, un encoder target invertibile o un'implicita condizione indotta, fornisce la mappatura stato-osservazione.

Per rendere questa corrispondenza concreta e analiticamente trattabile, l'autore introduce la Markov-Chain JEPA (MCJEPA). In questa formulazione:

• Lo stato latente $Z_t$ è categorico (discreto).
• Il predittore è sostituito da una matrice di transizione appresa e row-stocastica $A$.
• Le predizioni multi-step sono calcolate tramite potenze di matrici ($\hat{q}_{t+h} = q_t A^h$), garantendo l'esatta coerenza di Chapman–Kolmogorov (ovvero, le predizioni dirette e composte su più orizzonti sono identiche).

Il paper generalizza questo framework a:

• Transizioni condizionate: Dove $A$ è una funzione di informazioni laterali $\xi_t$.
• Stati continui: Sostituendo la matrice con un kernel di Markov (ad esempio, transizioni Gaussiane).
• Tempo continuo: Utilizzando matrici generatrici o equazioni differenziali stocastiche.
• Limiti deterministici: Mostrando che la standard JEPA deterministica è un caso degenero in cui il kernel di transizione diventa una Delta di Dirac.

L'autore definisce quattro livelli progressivamente più forti di corrispondenza tra JEPA e HMM:

1. Corrispondenza Computazionale: Condivisione della pipeline encode-transizione-emissione.
2. Rappresentazione Completa di Emissione: Possesso di una valida condizionale stato-osservazione.
3. Equivalenza HMM a Livello di Sequenza: Soddisfacimento delle condizioni di Markov, coerenza marginale e coerenza di filtraggio (formalizzate nel Teorema 1).
4. Equivalenza di Modello e Obiettivo: Addestramento del modello utilizzando la verosimiglianza di sequenza in stile HMM oltre agli obiettivi JEPA.

Inoltre, il paper interpreta l'apprendimento del Predictive Information Bottleneck come un processo di Markovizzazione: la compressione promuove la minimalità (rimozione della storia ridondante), mentre la prevedibilità residua testa la sufficienza (assicurarsi che non venga scartata alcuna informazione rilevante per la transizione).

Contributi Chiave

1. Costruzione di MCJEPA: L'introduzione di MCJEPA, che impone l'esatta coerenza di percorso multi-orizzonte attraverso una matrice di transizione condivisa, fornendo un ponte trasparente tra predittori neurali e catene di Markov classiche.
2. Condizioni Formali di Equivalenza HMM: La derivazione delle condizioni sufficienti (Teorema 1) sotto le quali una JEPA temporale probabilistica ammette una rappresentazione esatta di HMM a livello di sequenza, chiarendo la distinzione tra il ruolo dell'encoder (filtraggio) e la direzione di emissione.
3. Framework Diagnostico: Lo sviluppo di diagnostiche di "prevedibilità residua" per testare se uno stato appreso sia Markov-sufficiente. Se la storia precedente migliora la predizione oltre lo stato attuale, la rappresentazione è insufficiente.
4. Continuum degli Obiettivi: La dimostrazione che la stessa architettura di stato latente può essere addestrata tramite un continuum di obiettivi: pura predizione latente JEPA, pura verosimiglianza di sequenza HMM, o un ibrido di entrambi.

Risultati Sperimentali

L'autore ha condotto quattro esperimenti controllati su processi sintetici con strutture latenti note (ground-truth):

• Esperimento 1 (Recupero di HMM Finiti): MCJEPA ha recuperato con successo una matrice di transizione HMM a 4 stati con alta accuratezza (ARI $\approx$ 0.98) e quasi zero disaccordo di percorso ($D_{path} = 0$) grazie al vincolo della matrice condivisa. Al contrario, i predittori specifici per orizzonte (che apprendono matrici indipendenti per ogni step) hanno ottenuto una maggiore flessibilità in regimi ambigui ma non sono riusciti a mantenere una legge di transizione coerente e composizionale.
• Esperimento 2 (Filtraggio vs. Evidenza Locale): Utilizzando posteriore oracle, lo studio ha mostrato che il filtraggio dipendente dalla storia ($p(S_t | X_{\le t})$) supera significativamente l'evidenza locale ($p(S_t | X_t)$) all'aumentare dell'ambiguità di emissione, validando l'interpretazione dell'encoder di contesto JEPA come un filtro ammortizzato.
• Esperimento 3 (Compressione Predittiva): In un processo binario del secondo ordine, il modello ha compresso con successo una storia di 3 bit in una rappresentazione minima sufficiente a 4 stati ($Z_t = (X_{t-1}, X_t)$), rimuovendo la ridondanza di $X_{t-2}$ senza perdere potere predittivo. Ulteriore compressione a un singolo bit ha comportato una perdita di informazione predittiva, rilevata da un aumento netto del guadagno di storia residua ($\Delta_{hist}$).
• Esperimento 4 (Confronto degli Obiettivi): Confrontando i regimi di addestramento JEPA-only, HMM-style e Hybrid sulla stessa architettura:
  • JEPA-only ha appreso uno stato latente significativo ma con una NLL di sequenza e un errore di transizione più elevati.
  • HMM-style (verosimiglianza di sequenza + distillazione del filtro) ha ottenuto il più forte recupero dello stato (ARI $\approx$ 0.9945) e la NLL di sequenza più bassa.
  • Hybrid (addestramento con obiettivi JEza + HMM) ha quasi colmato il divario rispetto alle prestazioni del regime HMM-style (98.6% del divario di NLL di sequenza), mantenendo l'originale obiettivo di predizione latente JEPA.
  • Fondamentalmente, il regime HMM-style ha ottenuto risultati solidi senza la perdita di predizione latente JEPA ($L_{MC}$), provando che l'architettura è flessibile rispetto all'obiettivo di addestramento.

Significato e Rivendicazioni

Il paper sostiene che la JEPA temporale probabilistica non è semplicemente un'alternativa agli HMM, ma può essere interpretata rigorosamente come tale sotto specifiche condizioni strutturali e di addestramento.

• Insight Strutturale: L'encoder di contesto stocastico in PIB-VJEPA svolge il ruolo di una distribuzione di filtraggio (inferenza degli stati nascosti dalla storia), non di un modello di emissione. La direzione di emissione è fornita separatamente da un decoder o da un encoder inverso.
• Principio di Markovizzazione: L'apprendimento del Predictive Information Bottleneck è inquadrato come la costruzione di uno stato di Markov compatto e sufficiente. La compressione guida la minimalità, mentre le diagnostiche residue assicurano la sufficienza.
• Dipendenza dall'Obiettivo: L' "HMM-ness" di una JEPA non è determinata unicamente dalla sua architettura, ma dal suo obiettivo di addestramento. Sebbene l'architettura supporti la semantica HMM, l'addestramento standard JEPA ottimizza l'allineamento latente, non la verosimiglianza di sequenza. La piena equivalenza richiede una supervisione esplicita a livello di sequenza.
• Generalità: La corrispondenza si estende oltre MCJEPA ai predittori neurali generali (continui, ricorrenti o di ordine superiore) a patto che lo stato latente (possibilmente aumentato con memoria) soddisfi la proprietà di Markov del primo ordine.

Il paper conclude che questi risultati forniscono un'interpretazione a spazio di stato principled per la JEPA temporale, distinguendo tra ruoli computazionali, completezza rappresentativa e la semantica probabilistica imposta dall'obiettivo di addestramento.