多头注意力中 Token 选择的几何分析

摘要

我们提出了一个用于分析大语言模型（LLM）多头注意力的几何框架。在不改变机制本身的情况下，我们以 top-N 选择的视角审视标准注意力，并直接在 value 态空间中研究其行为。我们定义几何指标——精确率（Precision）、召回率（Recall）与 F 分数——来量化被选与未选 token 之间的可分性，并在有经验依据的假设（稳定的 value 范数与被压缩的 sink token、相似度的指数衰减、分段式注意力权重轮廓）下，推导出对维度与间隔有显式依赖的非渐近界。该理论预测存在一个小 N 的运行区间，其中非平凡可分性最强，并阐明了序列长度与 sink 相似度如何塑造这些指标。在实证上，跨 LLaMA-2-7B、Gemma-7B 与 Mistral-7B 的测量紧密贴合理论包络：top-N 选择使可分性更锐利，sink 相似度与 Recall 相关。我们还发现 LLaMA-2-7B 中的注意力头专门化为三种模式——Retriever、Mixer、Reset——具有不同的几何特征。总体而言，注意力表现为一个结构化的几何分类器，其 token 选择具有可测量的判据，为 LLM 提供注意力头层面的可解释性，并为几何感知的稀疏化与 LLM 注意力设计提供参考。

多头注意力中 Token 选择的几何分析 — 图1：价值-状态空间中几何可分性的二维示例。左：token 嵌入位于一个圆上。中：在按注意力权重 \alpha_{i} 缩放之后，被选中（洋红色星形）与未被选中（黑色圆点）的点都向原点移动。右：可以看到位于 B_{r_{\min}}(s) 与 B_{r_{\max}}(s) 中的点，以及 r_{\min} 和 r_{\max}。