To perform in-context learning, language models must extract signals from individual few-shot examples, aggregate these into a learned prediction rule, and then apply this rule to new examples. How is this implemented in the forward pass of modern transformer models? To study this, we consider a structured family of few-shot learning tasks for which the true prediction rule is to add an integer $k$ to the input. We find that Llama-3-8B attains high accuracy on this task for a range of $k$, and localize its few-shot ability to just three attention heads via a novel optimization approach. We further show the extracted signals lie in a six-dimensional subspace, where four of the dimensions track the unit digit and the other two dimensions track overall magnitude. We finally examine how these heads extract information from individual few-shot examples, identifying a self-correction mechanism in which mistakes from earlier examples are suppressed by later examples. Our results demonstrate how tracking low-dimensional subspaces across a forward pass can provide insight into fine-grained computational structures.

本研究解决了语言模型在上下文学习中提取少量示例信号的机制问题。提出了一种新颖的优化方法，揭示了Llama-3-8B模型在加法任务中的学习能力仅依赖于三个注意力头，并且信号在六维子空间中提取。重要发现是，这些注意力头通过一种自我校正机制来提升学习准确性，展示了在前向传播中低维子空间的追踪如何提供细致的计算结构洞察。

通过激活子空间理解加法的上下文学习