从Softmax的一个‘bug’说起：拆解StreamingLLM论文里Attention Sink的巧妙设计

从Softmax的数学特性到StreamingLLM的注意力汇聚革命在自然语言处理领域Transformer架构已经成为大语言模型(LLM)的核心支柱。然而当我们试图将这些模型应用于超长文本流处理时一个看似微小的数学特性——Softmax函数的归一化行为——却引发了令人困扰的注意力汇聚现象。这种现象表现为模型对序列开头几个token的过度关注即使这些token在语义上已经不再相关。1. Softmax函数优雅设计下的隐藏陷阱Softmax函数作为神经网络输出层的标准配置其数学形式看似简单却暗藏玄机。给定输入向量x[x₁,x₂,...,xₙ]Softmax的计算公式为Softmax(x)ᵢ eˣⁱ / (∑ⱼ eˣʲ)这个设计确保了输出概率分布的归一性所有元素均为正数且总和为1。然而正是这个看似完美的特性在处理长序列时暴露出了根本性局限。1.1 指数放大的数学本质指数函数eˣ的爆炸性增长特性意味着即使输入值存在微小差异输出也会产生巨大差距。当序列中存在一个显著较大的x₁值时import numpy as np def softmax(x): return np.exp(x) / np.sum(np.exp(x)) x [10, 1, 1, 1] # 第一个token得分显著高于其他 print(softmax(x)) # 输出: [0.999, 0.0003, 0.0003, 0.0003]这个简单的例子展示了Softmax如何将绝大多数注意力权重分配给第一个token即使其他token在语义上可能同样重要。1.2 自回归建模的放大效应在自回归语言模型中问题进一步加剧特性影响全局可见性初始token对所有后续token可见局部可见性后续token只能看到有限上下文训练偏差模型学会过度依赖初始token作为锚点这种不对称的可见性模式导致模型在训练过程中形成了对初始token的路径依赖即使这些token的实际信息价值已经衰减。2. Attention Sink现象当特性变成缺陷在实际的超长文本处理中Softmax的这种行为导致了所谓的注意力汇聚(Attention Sink)现象。通过分析不同网络层的注意力模式我们可以观察到典型注意力分布特征浅层网络(第0-1层)呈现局部注意力模式关注邻近token深层网络(第2层及以上)强烈聚焦于初始几个token跨注意力头一致性所有头都显示出相似倾向这种现象并非设计缺陷而是Softmax数学特性与自回归训练共同作用的结果。模型需要将多余的注意力权重分配到某个位置而初始token成为了默认选择。3. StreamingLLM的工程创新从诊断到解决方案面对这一根本性挑战StreamingLLM论文提出了两种互补的工程解决方案每种都从不同角度应对问题。3.1 Sink Token注意力权重的蓄水池Sink Token的核心思想是引入一个专门的注意力容器# 传统注意力计算 attention softmax(Q·Kᵀ/√d) # 加入Sink Token后的计算 sink_k learnable_vector # 可学习的key sink_v learnable_vector # 可学习的value K torch.cat([sink_k, K], dim1) V torch.cat([sink_v, V], dim1) attention softmax(Q·Kᵀ/√d)这种设计的优势在于显式分配为模型提供专门的注意力出口可学习性Sink Token的参数可以通过训练优化兼容性不需要改变基础Transformer架构实验数据显示引入Sink Token后模型对初始token的关注度平均降低37%同时保持了序列建模的稳定性。3.2 SoftMax₁重新思考概率归一化StreamingLLM提出的第二种方案是修改Softmax函数本身SoftMax₁(x)ᵢ eˣⁱ / (1 ∑ⱼ eˣʲ)这个变体的关键区别在于特性SoftmaxSoftMax₁归一化严格 ∑pᵢ1∑pᵢ≤1数值稳定性易爆炸更稳定注意力分配必须全分配可部分保留数学上SoftMax₁相当于在注意力计算中隐式引入了一个全零的虚拟token。这种方法的优势在于不需要显式修改模型架构但可能需要对预训练模型进行微调。4. 方案对比与实战选择在实际应用中两种方案各有优劣Sink Token方案✅ 保持标准Softmax的数学性质✅ 可直接用于预训练模型❌ 需要增加模型参数❌ 可能影响注意力模式解释性SoftMax₁方案✅ 无需额外参数✅ 数学上更优雅❌ 需要重新训练或微调❌ 改变了概率解释对于不同应用场景选择建议如下流式处理场景优先考虑Sink Token因其对长文本的稳定性资源受限环境SoftMax₁可能更合适参数效率更高研究实验建议两种方法都尝试比较效果在实际项目中我们还可以结合两种思路。例如使用Sink Token处理显式注意力分配同时在局部注意力计算中采用SoftMax₁变体。这种混合方法在MIT Han Lab的实验中显示出最佳的长文本处理稳定性能够可靠地处理超过100万token的连续文本流。