El Agrupamiento Diagonal Desbloquea el Paralelismo en Transformadores de Memoria Recurrente para Contextos Largos
Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts
June 5, 2025
Autores: Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, Ivan Oseledets
cs.AI
Resumen
Los modelos Transformer enfrentan dificultades en la inferencia de contextos largos debido a su complejidad cuadrática en tiempo y lineal en memoria. Los Transformers con Memoria Recurrente (RMTs) ofrecen una solución al reducir el costo asintótico a tiempo lineal y uso de memoria constante. Sin embargo, su mecanismo de actualización de memoria conduce a una ejecución secuencial, lo que genera un cuello de botella en el rendimiento.
Introducimos el *Diagonal Batching*, un esquema de planificación que desbloquea el paralelismo entre segmentos en los RMTs mientras preserva la recurrencia exacta. Este enfoque elimina la restricción secuencial, permitiendo una inferencia eficiente en GPU incluso para entradas de contexto largo individuales, sin necesidad de técnicas complejas de batching y pipelining. Dado que esta técnica es únicamente una reordenación computacional en tiempo de ejecución, los modelos RMT existentes pueden adoptarla sin necesidad de reentrenamiento.
Aplicado a un modelo LLaMA-1B ARMT, el *Diagonal Batching* logra una aceleración de 3.3x en comparación con la implementación estándar de atención completa de LLaMA-1B y una aceleración de 1.8x sobre la implementación secuencial de RMT en secuencias de 131,072 tokens. Al eliminar el cuello de botella secuencial, el *Diagonal Batching* reduce el costo y la latencia de inferencia, fortaleciendo así a los RMTs como una solución práctica para aplicaciones del mundo real con contextos largos.
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.