提出可高效推理的循环变压器,提升有效深度并降低内存开销。
The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
- 每层自回归记忆,实现层间递归但保持标准解码复杂度。
- 序列长度N时显存流量从Θ(N²)降至Θ(N log N),算力利用率提升至Θ(N/log N)。
- 减少层数即可超越同参数量Transformer,适合低延迟部署场景。
Transformers虽并行处理标记但时间深度受限:位置t处各层仅依赖前一层计算的键值对,深度受层数限制。循环模型具备无限时间深度,但存在优化不稳定且难以利用现代加速器的问题。本文提出循环变压器,通过让每层基于自身激活计算键值对,实现层间递归记忆,同时保留标准自回归解码开销。该结构在弱假设下可模拟传统Transformer与逐标记递归更新,且避免优化不稳定性。朴素实现预填充/训练时因键值顺序暴露而带宽受限,有效算术强度接近1;我们提出一种精确分块算法,在保持数学等价性的同时将显存流量从Θ(N²)降至Θ(N log N),使有效算术强度升至Θ(N/log N)(序列长度为N)。在150M和300M参数的C4预训练中,循环变压器在相同参数量下优于基准Transformer,且用更少层数达成性能提升,表明递归可实现深度与宽度的权衡,从而降低KV缓存内存占用与推理延迟。
原文摘要 · Abstract (English)
Transformers process tokens in parallel but are temporally shallow: at position $t$, each layer attends to key-value pairs computed based on the previous layer, yielding a depth capped by the number of layers. Recurrent models offer unbounded temporal depth but suffer from optimization instability and historically underutilize modern accelerators. We introduce the Recurrent Transformer, a simple architectural change where each layer attends to key-value pairs computed off its own activations, yielding layerwise recurrent memory while preserving standard autoregressive decoding cost. We show that the architecture can emulate both (i) a conventional Transformer and (ii) token-to-token recurrent updates under mild assumptions, while avoiding optimization instability. Naively, prefill/training appears bandwidth-bound with effective arithmetic intensity near $1$ because keys and values are revealed sequentially; we give an exact tiling-based algorithm that preserves the mathematical computation while reducing HBM traffic from $Θ(N^2)$ to $Θ(N\log N)$, increasing effective arithmetic intensity to $Θ(N/\log N)$ for sequence length $N$. On 150M and 300M parameter C4 pretraining, Recurrent Transformers improve cross-entropy over a parameter-matched Transformer baseline and achieve the improvement with fewer layers (fixed parameters), suggesting that recurrence can trade depth for width, thus reducing KV cache memory footprint and inference latency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。