Backbone a confronto per i World Model: RNN, Transformer e S4

L'articolo "Facing Off World Model Backbones: RNNs, Transformers, and S4" (Deng, Park e Ahn, 2023) si interroga su quale architettura di modellazione sequenziale sia più efficace per i world model latenti nel model-based reinforcement learning (MBRL). Introduce S4WM, che applica i layer Structured State Space Sequence (S4) all'immaginazione visiva latente ad alta dimensionalità.

Lavori correlati

• World model latenti / MBRL. Le fondamenta più vicine sono agenti a dinamica latente come PlaNet, Dreamer, DreamerV2 e DreamerV3, che apprendono rappresentazioni stocastiche compatte dello stato e predicono traiettorie future latenti per ottimizzare la politica. L'articolo considera i modelli ricorrenti di tipo RSSM — in particolare Dreamer — come la famiglia di baseline dominante.

• Reti neurali ricorrenti (RNN). I modelli a dinamica basati su RNN, LSTM e GRU forniscono aggiornamenti ricorrenti online dello stato e sono stati fondamentali per la modellazione visiva del mondo. Il loro principale limite è l'addestramento sequenziale nel tempo, che riduce il parallelismo temporale e rende più difficile modellare contesti molto lunghi.

• World model basati su Transformer. I modelli a dinamica basati su Transformer, incluso TransDreamer, sostituiscono la ricorrenza con l'attenzione, consentendo un addestramento parallelo e un accesso diretto a storie lunghe. Tuttavia, il loro costo di attenzione quadratico nella lunghezza della sequenza rende costose sia le rollout immaginate lunghe sia la modellazione visiva su contesti estesi.

• Modelli a spazio di stato strutturato (S4). S4 è un modello sequenziale a spazio di stato strutturato progettato per dipendenze a lungo raggio. Combina un'interpretazione ricorrente a spazio di stato con una forma convoluzionale che permette un addestramento parallelo della sequenza. L'articolo sostiene che S4 rappresenti un compromesso promettente: inferenza ricorrente a tempo lineare e training parallelizzabile.

• La novità di S4WM. A differenza dell'uso di S4 per benchmark sequenziali convenzionali, S4WM lo integra in un world model visivo latente, con codifica/decodifica delle immagini e dinamica latente stocastica, per poi confrontarlo con backbone RNN e Transformer all'interno di un framework comune di modellazione del mondo. Il contributo dichiarato è il primo world model basato su S4 in grado di generare sequenze di immagini ad alta dimensionalità tramite immaginazione latente.

Una sintesi ideale per una revisione della letteratura:

I world model precedenti hanno utilizzato principalmente dinamiche latenti ricorrenti o Transformer. I modelli basati su RNN sono efficienti nell'inferenza ma si addestrano in modo sequenziale, mentre i world model basati su Transformer parallelizzano l'addestramento ma hanno un costo temporale quadratico. S4WM esplora i layer a spazio di stato strutturato come backbone alternativo, pensato per mantenere il training parallelo supportando al contempo una modellazione efficiente della dinamica su orizzonti lunghi.