ReplaySSM riduce il traffico dello stato ricorrente nel decoding dei modelli linguistici

Quando i livelli state-space di un modello scrivono un grande stato ricorrente dopo ogni token, è possibile ridurre queste scritture senza cambiare gli output? La domanda conta perché questi aggiornamenti possono rappresentare una parte importante della latenza di decoding, anche nei modelli ibridi in cui il costo dell’attenzione cresce con il contesto. Il paper mette alla prova un modo diverso di rappresentare e aggiornare lo stato.

ReplaySSM conserva gli input recenti accanto a un checkpoint e calcola gli output tra aggiornamenti periodici dello stato, riducendo il traffico di memoria e facilitando il rollback.

Il paper chiede se i livelli state-space debbano scrivere tutto il loro stato aggiornato dopo ogni token, anche quando ogni aggiornamento richiede pochi calcoli.

ReplaySSM conserva gli input recenti accanto a un checkpoint e lo aggiorna periodicamente, con l’obiettivo di ridurre queste scritture mantenendo gli stessi output.

Su quattro modelli ibridi a batch size 256, il passo di decoding dell’intero modello è stato da 1,20 a 1,47 volte più veloce rispetto al decoding convenzionale ottimizzato di vLLM.

Il risultato supporta l’utilità di cambiare la gestione dello stato in questi test; esperimenti separati mostrano anche perché è importante evitare che le letture a bassa precisione alimentino gli aggiornamenti ricorrenti.

I risultati provengono da quattro modelli ibridi, ciascuno testato su una sola GPU e senza parallelismo tensoriale; gli esperimenti a bassa precisione hanno usato una quantizzazione emulata. Le misurazioni non stabiliscono che altrove si ottengano le stesse prestazioni o la stessa accuratezza.