[控场AI]
· 2 分钟阅读· 1,374 字

大模型上下文学习机制被破解:门控键值记忆电路

大模型上下文学习机制被破解:门控键值记忆电路

研究揭示ICL在任意标签下由门控键值记忆电路实现,并具因果可验证性。

该研究在三个3-4B参数Transformer中发现了少样本句子分类在任意标签下的完整ICL电路。机制核心是门控键值记忆:中间层注意力头将演示类别写为键,上层头匹配查询类别并复制标签;早期层信号编码标签是否重复并门控读出,另一组头抑制最近标签。研究通过激活修补、头消融和引导因果验证了因果性,并显示简化softmax读出模型可跨提示分布预测行为。该电路在三个模型中保守,可迁移到语义概念,在0.6B到8B模型家族中出现并在指令调优后保留。

引言:ICL 的机制谜题

上下文学习(In-Context Learning, ICL)让预训练语言模型仅凭提示中的示例就能完成新任务,无需更新参数。当标签任意时,模型必须获取一个它不可能记住的映射,学习发生在单次前向传播中。此前研究找到了候选部件,但未阐明它们如何组合。最新研究给出了少样本句子分类中任意标签下的完整电路,揭示了 ICL 的工作机制。

研究概览

残差流(Residual Stream) 是 Transformer 中贯穿各层的信息主干通道。每一层的注意力头和 MLP 子层都会读取残差流、计算后将结果加回残差流,因此它相当于一个共享的“工作内存”。研究者常把某个概念(如“标签身份”)在残差流中的表示称为该概念的编码。理解残差流有助于把握文中“标签 token 的残差流编码标签身份”这一说法:它意味着标签信息以线性可读的形式存在于这一通道中,可被后续层直接利用。

电路的核心:门控键值记忆

研究者在三个 3-4B 参数量的 Transformer 中发现了这一电路。模型在标签 token 位置实现键值记忆:每个标签 token 的残差流编码标签身份,中间层注意力头将前一个演示的类别写为键。查询时,上层头将查询类别与存储的键匹配并复制出对应标签。

一个早期层信号编码是否有标签重复,门控该读出过程;若无此信号,读取头会关注提示开头,无标签被偏好。另一组头抑制最近出现的标签。每一步都通过激活修补、头消融和引导因果建立。

激活修补(Activation Patching) 是一种因果干预方法:先让模型在原始输入上运行并记录某处激活,再在改变输入(如替换演示标签)后将该激活替换回去,观察输出是否恢复。若恢复,说明该激活对行为有因果作用。头消融(Head Ablation) 则直接移除或置零某个注意力头,检验其必要性。引导因果(Guided Causal) 通常指结合梯度或方向引导来定位并验证因果路径。三者共同构成从相关性走向因果性的证据链。

因果验证与泛化

研究通过多种因果干预验证了机制。一个带有两个常数的 softmax 读出模型,在一个提示分布上拟合,能预测另一个分布。该计算在三个模型中保守,命名头可复现行为,从头训练的 Transformer 也长出相同部件。电路可迁移到语义概念和有意义标签,因果门控在 0.6B 到 8B 模型家族中出现,并在指令调优后保留。

Softmax 读出模型 在这里指一种简化计算模型:读取头对键做 softmax 归一化后加权复制值,其中两个常数(如温度或偏置)刻画匹配的锐度。研究者用该模型拟合一个提示分布上的行为,再预测另一个分布,以此检验机制假设的泛化力。若简化模型能跨分布预测,说明真实电路与这一读出形式在功能上等价,从而支持“门控联想查找”的解释。

意义与展望

ICL 在此设定下是一个门控联想查找,由少数可识别组件构建,从演示因果追踪到预测。这为理解大模型能力提供了基础,可能启发更高效架构设计。

机制示意

#上下文学习#ICL#Transformer#键值记忆#因果验证#门控机制#大模型机制
分享:

相关推荐