循环 Transformer：更大的有效深度和高效解码

Transformer 并行处理词元，但时间上很浅：在位置 $t$ 处，每个层都会处理基于前一层计算的键值对，产生由层数限制的深度。循环模型提供无限的时间深度，但存在优化不稳定的问题，并且历史上没有充分利用现代加速器。我们引入了循环 Transformer，这是一个简单的架构更改，其中每个层都处理根据其自身激活计算的键值对，产生分层循环内存，同时保留标准自回归解码成本。我们证明该架构可以在温和的假设下模拟（i）传统的 Transformer 和（ii）词元到词元的循环更新，同时避免优化不稳定。天真地说，预填充/训练似乎受到带宽限制，有效算术强度接近1，因为键和值是按顺序显示的；我们给出了一种精确的基于平铺的算法，该算法保留了数学计算，同时将 HBM​​ 流量从 $θ(N^2)$ 减少到 $θ(N\log N)$，将序列长度 $N$ 的有效算术强度增加到 $θ(N/\log N)$。在 150M 和 300M 参数 C4 预训练上，循环 Transformer 比参数匹配的 Transformer 基线提高了交叉熵，并以更少的层数（固定参数）实现了改进，这表明循环可以用深度换取宽度，从而减少 KV 缓存内存占用和推理延迟。