FlashLoop：通过延迟更新实现快速且内存高效的循环Transformer

循环 Transformer 作为一种参数有效的方法，通过重复应用共享 Transformer 块来增加计算深度，引起了广泛关注。然而，它们相对于传统 Transformer 的实际优势仍然存在争议：每个额外的循环都会引发另一次 Transformer 传递，并且需要缓存另一组 KV 状态，导致推理 FLOP 和 KV 缓存内存随着循环深度而不断增长。这种开销在大循环计数和长上下文时变得尤其严重，从而阻止了 Looped Transformer 的参数效率转化为实际的推理效率。在本文中，我们发现循环引入的许多额外计算和存储都是多余的。随着循环的进行，状态变化变得越来越集中在一小部分词元上；注意输出差异主要由关键列的稀疏且稳定的子集主导；相邻循环之间的 KV 残差逐渐变得更适合低位量化。基于这些观察，我们引入了 FlashLoop，这是一种免训练推理框架，可通过词元稀疏更新、稀疏注意力和 KV 残差量化来减少跨循环冗余。在多个 Looped Transformer 模型中，\textsc{FlashLoop} 提供无损精度，同时实现高达 1.64$\times$ 的端到端加速和高达 6$\times$ KV 缓存内存减少，从而大大提高了将 Looped Transformer 扩展到更大计算深度和更长上下文的实用性。