WorldKV: Efficient World Memory with World Retrieval and Compression

Riepilogo Tecnico: WorldKV

Enunciazione del Problema
I modelli di diffusione video autoregressivi hanno abilitato la generazione di mondi in tempo reale condizionati dall'azione, tuttavia mantenere un mondo persistente in cui il ritorno a un punto di vista precedentemente visto produca contenuti coerenti rimane una sfida aperta. Sebbene l'attenzione con cache KV completa preservi questa coerenza, comporta una crescita lineare dell'impronta di memoria e del costo di attenzione con la lunghezza della simulazione, rompendo i vincoli in tempo reale e portando a fallimenti per esaurimento della memoria (OOM). Al contrario, l'inferenza a finestra scorrevole ripristina il throughput scartando il contesto vecchio ma sacrifica la coerenza a lungo termine, causando ai modelli di allucinare nuovi contenuti o di subire derive quando si rivedono le scene. Sebbene osservazioni recenti suggeriscano che le cache KV in questi modelli funzionino come una forma emergente di memoria del mondo, l'accesso a questa memoria in modo efficiente senza il costo computazionale completo di attendere all'intera storia rimane irrisolto.

Metodologia: WorldKV
Gli autori propongono WorldKV, un framework senza addestramento progettato per abilitare una memoria a lungo termine efficiente nei modelli di mondo video autoregressivi. Opera al di sopra delle architetture di inferenza a finestra scorrevole esistenti e introduce due componenti complementari: Recupero del Mondo e Compressione del Mondo.

1. Recupero del Mondo:

   • Meccanismo: Invece di scartare i blocchi di cache KV espulsi, WorldKV li memorizza nella memoria GPU o CPU. Quando il modello rivisita una scena, recupera selettivamente i blocchi pertinenti alla scena basandosi sulla corrispondenza della telecamera o dell'azione e li reinserisce nella finestra di attenzione nativa.
   • Integrazione: I blocchi recuperati vengono inseriti nativamente senza re-codifica o modifiche architetturali. Il framework è agnostico rispetto all'algoritmo di recupero, supportando strategie basate sulla similarità telecamera/azione o sull'attenzione.
   • Struttura della Finestra di Attenzione: La finestra attiva è partizionata in quattro regioni: cache KV sink (frame iniziali), cache KV recuperate (dalla storia memorizzata), cache KV recenti (frame immediatamente precedenti) e il blocco di denoising corrente.

2. Compressione del Mondo:

   • Motivazione: Memorizzare tutte le cache espulse è intensivo in termini di memoria. Gli autori osservano che i frame temporalmente adiacenti condividono un sostanziale contenuto visivo, producendo cache KV quasi duplicate.
   • Meccanismo: All'interno di ogni blocco memorizzato (ad esempio, 3 frame), il primo frame è designato come ancora. Il sistema calcola la similarità coseno Key-Key tra le chiavi dei frame non-ancora e il frame ancora.
   • Potatura: I token con alta similarità all'ancora vengono potati come ridondanti, mentre i token a bassa similarità (che rappresentano regioni appena rivelate o contenuti dinamici) vengono mantenuti.
   • Efficienza: Questo processo comprime ogni blocco a circa la metà della sua dimensione originale (ad esempio, mantenendo l'ancora più il 25% dei token dagli altri due frame), permettendo di far rientrare circa il doppio della storia entro un budget di memoria fisso.

Contributi Chiave

• Recupero del Mondo: Un framework che memorizza e recupera selettivamente i blocchi di cache KV espulsi utilizzando segnali basati su telecamera/azione o attenzione, permettendo al modello di accedere al contesto a lungo termine senza modifiche architetturali.
• Compressione del Mondo: Un meccanismo di potatura basato sulla similarità delle chiavi che riduce lo storage per blocco di circa il 50%, permettendo una copertura storica doppia sotto gli stessi vincoli di memoria, preservando o migliorando la fedeltà del ritorno.
• Efficienza Senza Addestramento: L'approccio non richiede fine-tuning, distillazione o moduli di memoria dedicati, operando interamente al momento dell'inferenza su backbones congelati.

Risultati Sperimentali
Il framework è stato valutato su due modelli di mondo video autoregressivi di scale diverse: Matrix-Game-2.0 (1,3 miliardi di parametri, addestrato su sequenze brevi) e LingBot-World-Fast (14 miliardi di parametri, addestrato su video lunghi).

• Throughput vs Fedeltà: WorldKV eguaglia o supera la fedeltà di memoria dell'attenzione con KV completa mantenendo un throughput circa 2 volte superiore. Su LingBot-World-Fast, raggiunge prestazioni comparabili alla KV completa (che è la modalità nativa) ma con FPS significativamente più alti. Su Matrix-Game-2.0, supera sia le baseline a finestra scorrevole che quelle con KV completa; notevolmente, la KV completa performa peggio qui a causa di errori cumulativi derivanti dall'attenzione su cache accumulate fuori distribuzione, mentre il recupero selettivo di WorldKV mitiga questo problema.
• Confronto con SOTA: WorldKV supera le baseline addestrate sulla memoria (WorldPlay, Yume-1.5) su metriche come LPIPS, PSNR e FID senza alcun addestramento specifico per la memoria.
• Studi di Ablazione:
  • Compressione intra-blocco: Una compressione moderata (mantenendo l'ancora + 50% degli altri token) produce prestazioni ottimali; mantenere solo l'ancora degrada la fedeltà, mentre mantenere tutti i token offre rendimenti decrescenti.
  • Copertura inter-blocco: Comprimere più blocchi per adattarsi a un budget fisso (ad esempio, da 6 blocchi a un budget di 3 blocchi) migliora la fedeltà rispetto al recupero di meno blocchi a piena risoluzione, evidenziando il valore di una copertura storica più ampia.

Significato e Affermazioni
Il paper afferma che WorldKV dimostra che la cache KV stessa è una forma fondamentale ed emergente di memoria del mondo, anche in modelli non esplicitamente addestrati su sequenze lunghe. Gestendo efficientemente questa risorsa esistente attraverso recupero e compressione, WorldKV raggiunge un modello di mondo persistente che mantiene visite coerenti alle scene rispettando i vincoli di inferenza in tempo reale. Gli autori sottolineano che ciò è ottenuto senza la latenza della ricostruzione 3D o il sovraccarico di addestramento di moduli di memoria dedicati.

Limitazioni
Gli autori notano che WorldKV è un metodo di gestione al momento dell'inferenza; pertanto, la fedeltà visiva è limitata dalla qualità di generazione del modello pre-addestrato sottostante. In simulazioni estremamente lunghe, l'accumulo di errori può ancora verificarsi. Inoltre, sebbene lo scaricamento sulla CPU possa ridurre ulteriormente l'uso della VRAM, le attuali latenze di trasferimento host-dispositivo impediscono la generazione in tempo reale in quella specifica configurazione, che viene lasciata come lavoro futuro.