arXiv:2505.05145cs.LGcs.AI2025-05被引 11

通过分析注意力头的激活子空间,揭示了大模型如何用少量样例实现加法推理。

Understanding In-context Learning of Addition via Activation Subspaces

  • 定位到仅少数注意力头负责少样本学习,用子空间分解揭示其计算机制。
  • 在Llama-3-8B上发现三个头的六维子空间可完整捕捉加法规则,其中四维追踪个位数周期变化。
  • 提出信息流动数学关系,发现后续示例能自我修正前期错误,提升推理鲁棒性。

为实现少样本学习,语言模型从少量输入-标签对中提取信号,聚合形成预测规则并应用于新输入。这一过程在现代Transformer模型的前向传播中如何实现?我们研究了一类结构化的少样本任务,其真实规则是将整数 $k$ 加到输入上。提出一种新优化方法,将模型的少样本能力定位到少数注意力头。通过降维与分解深入分析单个注意力头,发现对于 Llama-3-8B-instruct,其任务机制可压缩为三个注意力头构成的六维子空间:其中四个维度以周期为 $2$、$5$ 和 $10$ 的三角函数追踪个位数字,两个维度以低频成分追踪数值量级。为进一步理解机制,我们推导出注意力头中“聚合”与“提取”子空间之间的数学恒等式,从而追踪个体示例到最终聚合概念的信息流。由此识别出一种自校正机制:早期演示所学错误被后期演示抑制。结果表明,通过追踪前向传播中局部化注意力头的低维子空间,可深入揭示语言模型内部精细的计算结构。

原文摘要 · Abstract (English)

To perform few-shot learning, language models extract signals from a few input-label pairs, aggregate these into a learned prediction rule, and apply this rule to new inputs. How is this implemented in the forward pass of modern transformer models? To explore this question, we study a structured family of few-shot learning tasks for which the true prediction rule is to add an integer $k$ to the input. We introduce a novel optimization method that localizes the model's few-shot ability to only a few attention heads. We then perform an in-depth analysis of individual heads, via dimensionality reduction and decomposition. As an example, on Llama-3-8B-instruct, we reduce its mechanism on our tasks to just three attention heads with six-dimensional subspaces, where four dimensions track the unit digit with trigonometric functions at periods $2$, $5$, and $10$, and two dimensions track magnitude with low-frequency components. To deepen our understanding of the mechanism, we also derive a mathematical identity relating ``aggregation'' and ``extraction'' subspaces for attention heads, allowing us to track the flow of information from individual examples to a final aggregated concept. Using this, we identify a self-correction mechanism where mistakes learned from earlier demonstrations are suppressed by later demonstrations. Our results demonstrate how tracking low-dimensional subspaces of localized heads across a forward pass can provide insight into fine-grained computational structures in language models.

少样本学习注意力机制子空间分析推理机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。