3.3 KiB
3.3 KiB
title, created, updated, type, tags, sources, venue, authors
| title | created | updated | type | tags | sources | venue | authors | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| StreamingLLM: 基于注意力汇的高效流式语言模型 | 2026-05-14 | 2026-05-14 | paper |
|
|
ICLR 2024 |
|
StreamingLLM: 基于注意力汇的高效流式语言模型
Efficient Streaming Language Models with Attention Sinks
guangxuan-xiao, yuandong-tian, beidi-chen, song-han, mike-lewis — ICLR 2024
核心问题
如何让 LLM 在流式场景(多轮对话、实时生成)中处理无限长输入序列,同时保持效率和性能?
现有方案的问题:
- Dense Attention:O(T²) 复杂度,且超出预训练窗口后性能崩溃
- Window Attention:只缓存最近 Token 的 KV,但当初始 Token 被逐出时模型崩溃
- Sliding Window + Re-computation:性能好但太慢,每个新 Token 都要重建 KV
核心发现:Attention Sink(注意力汇)
作者发现了一个关键现象:LLM 的初始 Token 在所有层和注意力头上都吸引了不成比例的高注意力分数,无论其语义是否相关。
原因:SoftMax 函数要求所有 token 的注意力分数之和为 1。即使当前 query 没有强语义匹配,模型也必须把多余的注意力值 "倾倒" 到某处。由于自回归建模的特性,初始 Token 对所有后续 Token 都可见,因此被自然训练为注意力的 "汇" (sink)。
→ 这就是为什么 Window Attention 一旦驱逐初始 Token 就会崩溃:它破坏了 SoftMax 分母的分布。
详见 attention-sinks。
方法:StreamingLLM
基于上述洞察,StreamingLLM 的 KV 缓存分为两部分:
- Attention Sinks(4 个初始 Token):稳定注意力计算,锚定 SoftMax 分布
- Rolling KV Cache(最近 Token):负责语言建模的语义内容
关键设计:位置编码在 cache 内部分配(而非原始文本),这对 RoPE 和 ALiBi 的性能至关重要。
- 无需微调,直接适用于任何使用相对位置编码的自回归 LLM
- 支持 RoPE(Llama-2, Falcon, Pythia)和 ALiBi(MPT)
- 与上下文窗口扩展方法正交,可互补使用
关键结果
- 在 Llama-2 (7/13/70B)、MPT (7/30B)、Falcon (7/40B)、Pythia (2.9/6.9/12B) 上验证
- 稳定语言建模至 400 万 Token 以上
- 相比 Sliding Window + Re-computation 加速 最高 22.2×
- 4 个初始 Token 即可恢复性能,更多 Token 边际收益递减
Sink Token 预训练
- 用 1.6 亿参数模型从头预训练验证
- 在所有训练样本前添加一个可学习的专用 Sink Token → 只需要这一个 Token 即可稳定流式推理
- 相比之下,Vanilla 模型需要 4 个初始 Token 作为 attention sinks
- SoftMax-off-by-One(Zero Sink)有改善但不够充分
→ 建议未来的 LLM 预训练统一加入 Sink Token
概念网络
- attention-sinks — 核心现象
- window-attention — 失败的基线
- length-extrapolation — 问题域
- rolling-kv-cache — StreamingLLM 机制
- sink-token — 预训练改进
- kv-cache-bottleneck — 相关优化
- rotary-position-embedding — 位置编码