Attention Sinks

注意力汇聚 token 是一种token，它在许多注意力头和层中会获得很大比例的注意力权重，尽管其内容与查询无关；它通常是序列中的第一个 token。MIT、Meta AI 和卡内基梅隆大学的 Guangxuan Xiao 及其同事在发表于 ICLR 2024 的论文《带注意力汇聚的高效流式语言模型》中为这一现象命名。他们展示了当滑动 KV 缓存逐出首批 token 时，Llama 2、MPT、Falcon 和 Pythia 等大型语言模型（LLM）会出现性能崩溃；保留其中四个 token 则能恢复稳定生成。

汇聚 token 为何会形成#

自注意力会通过softmax将每个查询的分数转换为权重，而这些权重的总和始终为 1。当上下文中没有任何与某个注意力头相关的 token 时，权重仍必须分配到某处。在因果语言模型中，每个较晚的 token 都能看到最初的 token，因此训练会让这些初始 token 成为注意力头无须使用权重时的默认去向。在对 Llama 2 进行的 StreamingLLM 分析中，除最底部两层外，每一层的所有注意力头都会将注意力集中在初始 token 上。向 softmax 添加汇聚 logit 后，可以看出这一效果：

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

没有汇聚项时，该注意力头会将全部权重分配给六个无关值，输出范数为 0.83。汇聚 logit 为 3 时，汇聚项获得 77% 的权重，输出缩小到 0.19，因此该注意力头几乎不对当前查询产生作用。OpenAI 的开放权重 gpt-oss 模型会为每个注意力头加入可学习的汇聚 logit，让每个注意力头都有一个存放多余权重的位置，且这些权重不会影响输出。

使用 StreamingLLM 进行流式生成#

无限生成的模型最终会超出任何上下文窗口，也无法将每个 token 都保留在 KV 缓存中；而普通窗口注意力只保留最近的 token，一旦最初的 token 被逐出，就会失效。StreamingLLM 会永久保留最初四个 token 的键和值，并将它们与近期 token 的滚动窗口并列存放；它按 token 在缓存中的位置而非原始文本中的位置分配位置索引，并在每一步重新应用旋转位置嵌入（RoPE）：

这种方法使语言模型在 400 万个及更多 token 的序列上保持稳定，速度最高可达滑动窗口方案的 22.2 倍；后者需要为每个新 token 重新计算缓存。不过，窗口仍然限制着模型能够回忆的内容：从中间位置逐出的 token 将永久消失，因此流式聊天机器人或 AI 智能体只能记住最近几轮对话。作者还在预训练模型时，于每个样本开头添加了一个专用的可学习占位 token；该 token 单独承担了汇聚角色。

注意力汇聚与相关术语的区别#

• 滑动窗口注意力： 将每个 token 的注意力限制在近期邻居上。如果在推理时将其应用于经过全注意力训练的模型，那么汇聚 token 离开窗口后，模型就会失效。
• 上下文退化： 提示不断增长但仍未超出上下文窗口时，准确率随之下降的现象。

视觉模型中的汇聚 token#

视觉 Transformer（ViT）也表现出类似现象。在《视觉 Transformer 需要寄存器》一文中，Meta 和 Inria 的研究人员在监督式和自监督 ViT 中发现，低信息量背景图块上存在范数异常大的 token；这些 token 被重新用于内部计算。向输入添加少量额外的寄存器 token 后，这些 token 就承担了该作用，消除了特征图和注意力图中的伪影，并提升了密集预测性能。

Meta 的 SAM 2 可在 Ultralytics 包中运行；它采用类似的“保留锚点加窗口”方案来分配视频内存：对于每一帧新图像，其内存注意力会读取用户提示帧对应的内存库条目（这些条目会在整个视频期间保留），再加上最近 6 帧的条目。因此，无论视频运行多久，每帧的注意力开销都保持不变。