Riepilogo Tecnico: Composizione di Esperti di Memoria per Modelli del Mondo Diffusivi

Enunciato del Problema

I modelli del mondo (WM) sono essenziali per l'apprendimento per rinforzo (RL) e la pianificazione, in quanto prevedono stati futuri plausibili coerenti con le osservazioni passate. Tuttavia, le architetture esistenti affrontano un compromesso fondamentale tra fedeltà della memoria ed efficienza computazionale. I modelli basati su Transformer preservano i dettagli locali ma soffrono di costi di attenzione quadratici, limitando la lunghezza del contesto. Al contrario, le reti ricorrenti e i modelli a spazio di stato (SSM) scalano in modo efficiente ma comprimono la storia in stati nascosti, scartando inevitabilmente dettagli fini nel tempo. Aumentare semplicemente la lunghezza del contesto in una singola architettura porta a un addestramento instabile e a costi di inferenza proibitivi. Inoltre, una composizione ingenua di più modelli spesso comporta un "collasso delle modalità", in cui la distribuzione combinata affina eccessivamente le regioni comuni e perde diversità.

Metodologia

Gli autori propongono la Composizione di Esperti di Memoria (CoME), un framework basato sulla diffusione che disaccoppia la coerenza futuro-passato da una singola architettura monolitica. Invece, CoME sfrutta un insieme di esperti specializzati, ciascuno gestisce un ruolo specifico della memoria, e li fonde probabilisticamente al momento dell'inferenza.

1. Prodotto di Esperti Contrastivi (PoCE)

Per affrontare le limitazioni del standard Prodotto di Esperti (PoE), che può sopprimere la diversità e amplificare modalità spurie, gli autori introducono una formulazione di Prodotto di Esperti Contrastivi.

• Meccanismo: Per un insieme di esperti eterogenei $p_i(x)$, la composizione è definita come $p(x) \propto \prod \tilde{p}_i(x)$, dove $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$.
• Baseline Contrastiva: $\bar{p}_i(x)$ rappresenta la distribuzione di baseline incondizionata dell'esperto. Il coefficiente $\alpha_i > 1$ agisce come fattore di contrasto.
• Giustificazione Teorica: A differenza dell'ingenuo temperamento (che riduce la varianza del kernel e distorce la geometria locale), questo mix contrastivo ridisegna i componenti del mix mantenendo la forma dei kernel. Sopprime efficacemente le modalità spurie (dove un esperto assegna un'alta probabilità senza prove forti) mantenendo modalità secondarie valide, a condizione che l'esperto condizionato assegni una probabilità più alta alle modalità valide rispetto alla baseline.

2. Esperti di Memoria Specializzati

Il framework istanzia tre esperti complementari:

• Memoria a Breve Termine (STM): Un modello diffusivo condizionato su una finestra di contesto recente ($10\text{--}100$ frame). Cattura dinamiche locali fini e dipendenze temporali immediate utilizzando un'architettura con lunghezza di contesto estesa (ad esempio, attenzione a finestra scorrevole).
• Memoria a Lungo Termine (LTM): Un modello diffusivo separato che memorizza la storia episodica ($100\text{--}1000$ frame) direttamente nei suoi pesi. Questo è ottenuto tramite fine-tuning al momento del test di adattatori LoRA leggeri sul contesto a lungo termine. Questo approccio ammortizza il costo del richiamo, consentendo il riutilizzo in tempo costante dell'esperienza passata senza immettere l'intera storia attraverso il modello ad ogni passo. Evita l'oblio catastrofico congelando il backbone e aggiornando solo gli adattatori.
• Memoria Spaziale a Lungo Termine (SLTM): Un'estensione della LTM che incorpora prior spaziali (ad esempio, pose della camera, mappe di punti) estratte dalla storia. Questo esperto impone coerenza geometrica e spaziale, disambiguando le posizioni in ambienti ripetitivi o visivamente simili.

3. Generazione Unificata

La distribuzione finale è un prodotto del prior pre-addestrato e dei tre esperti contrastivi (STM, LTM, SLTM). Durante il campionamento, le funzioni di punteggio di questi esperti sono combinate additivamente (nello spazio logaritmico) per guidare il processo di denoising, assicurando che il futuro generato rispetti le dinamiche recenti, la storia a lungo termine e la coerenza spaziale.

Contributi Chiave

1. Architettura di Memoria Disaccoppiata: Il documento propone di risolvere il compromesso della memoria distribuendo la memoria su esperti specializzati piuttosto che su un singolo backbone, formulato come un problema di Prodotto di Esperti.
2. Prodotto di Esperti Contrastivi (PoCE): Una strategia innovativa per fondere esperti eterogenei. Gli autori forniscono prove teoriche ed empiriche che questo metodo previene il collasso delle modalità e l'eccessiva sicurezza fattorizzando le modalità ridondanti, a differenza del PoE standard.
3. Memoria Diffusiva Esterna: L'introduzione di un modulo LTM che utilizza il fine-tuning al momento del test di adattatori LoRA per memorizzare conoscenze episodiche nei pesi. Questo è esteso alla SLTM per incorporare la struttura spaziale, migliorando l'accuratezza nei compiti di navigazione.

Risultati Sperimentali

Il metodo è stato valutato su benchmark sintetici (Memory Maze, DMLab, Minecraft) e del mondo reale (RealEstate10K, RECON).

• Qualità di Ricostruzione: Sul dataset Memory Maze, CoME ha ottenuto prestazioni all'avanguardia in LPIPS (0,097), SSIM (0,892) e PSNR (23,07), superando significativamente le baseline inclusi transformer con attenzione completa, SSM e modelli a finestra scorrevole.
• Scalabilità: A differenza dei modelli con attenzione completa che richiedono $60\times$ più potenza di calcolo rispetto alla baseline STM, CoME ottiene una coerenza superiore con una frazione del costo di addestramento e inferenza.
• Lunghezza del Contesto: Il modello non ha mostrato saturazione nella coerenza delle prestazioni fino a 480 frame di contesto, dimostrando l'efficacia del componente LTM.
• Pianificazione e Navigazione: Nei compiti di navigazione condizionati all'obiettivo (benchmark RECON), CoME ha ridotto l'Errore di Traiettoria Assoluto (ATE) a 0,96, superando i modelli del mondo di navigazione specializzati (NWM) e altre baseline.
• Capacità di Richiamo: In un esperimento di traiettoria inversa su RealEstate10K, CoME ha richiamato con successo i frame iniziali dopo roll-out in avanti, mentre le baseline senza memoria generavano scene plausibili ma incoerenti.
• Studi di Ablazione: Rimuovere la formulazione contrastiva ha causato il crollo delle prestazioni, confermando la necessità del meccanismo PoCE. Aumentare il rango LoRA e la lunghezza del contesto per LTM ha migliorato costantemente i risultati senza overfitting.

Significato e Affermazioni

Il documento afferma di stabilire un nuovo paradigma per la costruzione di modelli del mondo diffusivi potenziati dalla memoria. Distribuendo la memoria su esperti specializzati e utilizzando una strategia di composizione contrastiva, CoME bilancia l'efficienza con un'alta fedeltà nella coerenza temporale. Gli autori sostengono che questo approccio permette ai modelli del mondo di scalare a contesti lunghi senza incorrere in costi quadratici o subire la perdita di dettagli intrinseca nella compressione ricorrente. Il metodo è presentato come un framework generale applicabile a vari modelli del mondo basati sulla diffusione, migliorando la loro affidabilità per la pianificazione e il processo decisionale in contesti di apprendimento per rinforzo. Il lavoro sottolinea che la memoria dovrebbe essere una proprietà distribuita del sistema piuttosto che un vincolo di una singola architettura.