Il Batching Diagonale Sblocca il Parallelismo nei Trasformatori con Memoria Ricorrente per Contesti Lunghi

Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

June 5, 2025

Autori: Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, Ivan Oseledets

cs.AI

Abstract

I modelli Transformer faticano nell'inferenza su contesti lunghi a causa della loro complessità quadratica nel tempo e lineare nella memoria. I Recurrent Memory Transformers (RMT) offrono una soluzione riducendo il costo asintotico a tempo lineare e utilizzo di memoria costante. Tuttavia, il loro meccanismo di aggiornamento della memoria porta a un'esecuzione sequenziale, creando un collo di bottiglia nelle prestazioni. Introduciamo il Diagonal Batching, uno schema di pianificazione che sblocca il parallelismo tra segmenti negli RMT preservando la ricorrenza esatta. Questo approccio elimina il vincolo sequenziale, consentendo un'inferenza efficiente su GPU anche per singoli input a contesto lungo, senza ricorrere a tecniche complesse di batching e pipelining. Poiché la tecnica è puramente un riordinamento dei calcoli in fase di esecuzione, i modelli RMT esistenti possono adottarla senza necessità di riaddestramento. Applicato a un modello LLaMA-1B ARMT, il Diagonal Batching ottiene un miglioramento di 3,3x rispetto al LLaMA-1B con attenzione completa standard e un miglioramento di 1,8x rispetto all'implementazione sequenziale di RMT su sequenze di 131.072 token. Rimuovendo il collo di bottiglia sequenziale, il Diagonal Batching riduce i costi e la latenza dell'inferenza, rafforzando così gli RMT come soluzione pratica per applicazioni reali con contesti lunghi.

English

Transformer models struggle with long-context inference due to their quadratic time and linear memory complexity. Recurrent Memory Transformers (RMTs) offer a solution by reducing the asymptotic cost to linear time and constant memory usage. However, their memory update mechanism leads to sequential execution, causing a performance bottleneck. We introduce Diagonal Batching, a scheduling scheme that unlocks parallelism across segments in RMTs while preserving exact recurrence. This approach eliminates the sequential constraint, enabling efficient GPU inference even for single long-context inputs without complex batching and pipelining techniques. Because the technique is purely a run-time computation reordering, existing RMT models adopt it with no retraining. Applied to a LLaMA-1B ARMT model, Diagonal Batching yields a 3.3x speedup over standard full-attention LLaMA-1B and a 1.8x speedup over the sequential RMT implementation on 131,072-token sequences. By removing sequential bottleneck, Diagonal Batching reduces inference cost and latency, thereby strengthening RMTs as a practical solution for real-world, long-context applications.