Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2506.05229 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866909639561445376 |
|---|---|
| author | Sivtsov, Danil Rodkin, Ivan Kuzmin, Gleb Kuratov, Yuri Oseledets, Ivan |
| author_facet | Sivtsov, Danil Rodkin, Ivan Kuzmin, Gleb Kuratov, Yuri Oseledets, Ivan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_05229 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts Sivtsov, Danil Rodkin, Ivan Kuzmin, Gleb Kuratov, Yuri Oseledets, Ivan Machine Learning Computation and Language 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. |
| title | Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2506.05229 |