Диагональное пакетирование раскрывает параллелизм в рекуррентных трансформерах с памятью для работы с длинными контекстами

Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

June 5, 2025

Авторы: Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, Ivan Oseledets

cs.AI

Аннотация

Трансформерные модели сталкиваются с трудностями при работе с длинными контекстами из-за их квадратичной временной и линейной памяти сложности. Рекуррентные Трансформеры с Памятью (RMT) предлагают решение, снижая асимптотическую стоимость до линейного времени и постоянного использования памяти. Однако их механизм обновления памяти приводит к последовательному выполнению, создавая узкое место в производительности. Мы представляем Диагональное Батчирование — схему планирования, которая раскрывает параллелизм между сегментами в RMT, сохраняя точную рекуррентность. Этот подход устраняет последовательное ограничение, позволяя эффективное выполнение на GPU даже для одиночных длинных контекстов без сложных техник батчинга и конвейеризации. Поскольку этот метод представляет собой исключительно переупорядочивание вычислений во время выполнения, существующие модели RMT могут его использовать без необходимости переобучения. Примененное к модели LLaMA-1B ARMT, Диагональное Батчирование обеспечивает ускорение в 3.3 раза по сравнению со стандартной LLaMA-1B с полным вниманием и в 1.8 раза по сравнению с последовательной реализацией RMT на последовательностях из 131 072 токенов. Устраняя последовательное узкое место, Диагональное Батчирование снижает стоимость и задержку вывода, укрепляя RMT как практическое решение для реальных приложений с длинными контекстами.

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.