通过激活子空间理解加法的注意力机制学习
机器学习
2025-10-10 v3 人工智能
计算与语言
摘要
为了进行少样本学习,语言模型从少量输入-标签对中提取信号,将其聚合为学习到的预测规则,并将该规则应用于新输入。如何在现代变换器模型的前向传播中实现这一点?为探索此问题,我们研究了一类具有结构化少样本学习任务的集合,其真实的预测规则是将整数 加到输入上。我们引入一种新型优化方法,将模型的少样本能力局部化到仅少数注意力头。随后,我们通过降维和分解对单个注意力头进行深入分析。以 Llama-3-8B-instruct 为例,我们将其在本任务上的机制简化为仅三个注意力头,每个注意力头的子空间维度为六维,其中四个维度通过周期为 、 和 的三角函数跟踪个位数字,两个维度则通过低频分量跟踪数值大小。为深化对该机制的理解,我们还推导了一个关于注意力头“聚合”子空间与“提取”子空间之间数学关系的等式,使我们能够跟踪从单个示例到最终聚合概念的信息流。通过该等式,我们识别出一种自我纠正机制:早期演示中学习到的错误被后续演示所抑制。我们的结果表明,跟踪局部分析头在前向传播中的低维子空间,可提供对语言模型中细粒度计算结构的洞见。
引用
@article{arxiv.2505.05145,
title = {Understanding In-context Learning of Addition via Activation Subspaces},
author = {Xinyan Hu and Kayo Yin and Michael I. Jordan and Jacob Steinhardt and Lijie Chen},
journal= {arXiv preprint arXiv:2505.05145},
year = {2025}
}