arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2609.29812cs.LG

FlashLoop:通过惰性更新实现快速且内存高效的循环Transformer

FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates

Wanqi Yang, Shiwei Liu

首次发表
浏览论文内容

中文总结 AI 辅助

本文提出FlashLoop,一种无需训练的推理框架,通过token稀疏更新、稀疏注意力和KV残差量化减少循环Transformer中的跨循环冗余,实现无损精度下高达1.64倍加速和6倍KV缓存内存缩减。

中文摘要 AI 辅助

循环Transformer作为一种参数高效的方法,通过重复应用共享Transformer块来增加计算深度,已引起广泛关注。然而,其相对于传统Transformer的实际优势仍存在争议:每增加一次循环就会多一次Transformer前向传播,并需要缓存另一组KV状态,导致推理FLOPs和KV缓存内存随循环深度持续增长。在循环次数多和上下文长的情况下,这种开销变得尤为严重,阻碍了循环Transformer的参数效率转化为实际的推理效率。在本文中,我们发现循环引入的额外计算和存储大部分是冗余的。随着循环的进行,状态变化越来越集中于一小部分token上;注意力输出差异主要由稀疏且稳定的键列子集主导;相邻循环之间的KV残差越来越适合低位量化。基于这些观察,我们提出了FlashLoop,一个无需训练的推理框架,通过token稀疏更新、稀疏注意力和KV残差量化来减少跨循环冗余。在多个循环Transformer模型上,FlashLoop在实现无损精度的同时,实现了高达1.64倍的端到端加速和高达6倍的KV缓存内存减少,显著提高了将循环Transformer扩展到更大计算深度和更长上下文的实用性。

英文摘要

Looped Transformers have attracted substantial attention as a parameter-efficient approach to increasing computational depth through repeated application of shared Transformer blocks. However, their practical advantages over conventional Transformers remain under debate: each additional loop incurs another Transformer pass and requires caching another set of KV states, causing inference FLOPs and KV-cache memory to grow continuously with loop depth. This overhead becomes particularly severe at large loop counts and long context, preventing the parameter efficiency of Looped Transformers from translating into practical inference efficiency. In this paper, we find that much of the additional computation and storage introduced by looping is redundant. As recurrence proceeds, state changes become increasingly concentrated on a small subset of tokens; attention-output differences are dominated by a sparse and stable subset of key columns; and KV residuals between adjacent loops become progressively more amenable to low-bit quantization. Building on these observations, we introduce FlashLoop, a training-free inference framework that reduces cross-loop redundancy through token-sparse updates, sparse attention, and KV-residual quantization. Across several Looped Transformers models, FlashLoop delivers lossless accuracy while achieving up to 1.64$\times$ end-to-end speedup and up to 6$\times$ KV-cache memory reduction, substantially improving the practicality of scaling Looped Transformers to greater computational depths and longer context.

发表机构

  • ELLIS Institute Tübingen(ELLIS 图宾根研究所)
  • Max Planck Institute for Intelligent Systems(马克斯·普朗克智能系统研究所)
  • Tübingen AI Center(图宾根人工智能中心)

机构由 AI 辅助整理,请以论文原文为准。

补充信息

↑