跳到正文
原文
arXiv· arXiv:2502.14010· Kayo Yin, Jacob Steinhardt·本站收录· 原文发表

哪些注意力头对上下文学习更关键?

Which Attention Heads Matter for In-Context Learning?

论文速读

可解释性

据论文 PDF 整理(AI 生成),以原文为准

在12个模型上,作者称少样本ICL更依赖FV heads而非induction heads。

问题
少样本ICL同时被归因于复制下一token的induction heads,和编码任务向量的FV heads。作者要分辨哪一类头在因果上支撑少样本准确率,以及两类分数为何相关。
方法
在12个模型上用重复随机序列的注意力算induction score,用打乱标签后的因果中介算FV score。对最高分头做均值消融,并排除另一类top 2%头以分开相关。另在7个Pythia的8个检查点上看头如何转变。
实验与结果
作者称FV消融比induction消融更降少样本准确率,大于1B且保留top 2% FV头时后者与随机消融相当。token-loss差值相反。许多FV头早期induction分高后下降,反向未见。
局限与可以继续做的
作者称猜想C1解释不了全程induction分低的FV头,并建议训练时去掉induction头再看FV。训练动态只含7个Pythia、各8个检查点;论文未报告消融重复次数。
实验设置Pythia 70M、Pythia 160M 等 · Appendix A.7的45个自然语言ICL任务、Todd et al. (2024)的ICL任务 等 · few-shot ICL accuracy、token-loss difference 等
被测模型
Pythia 70MPythia 160MPythia 410MPythia 1BPythia 1.4BPythia 2.8BPythia 6.9BGPT-2GPT-2 mediumGPT-2 largeGPT-2 XLLlama 2
基准
Appendix A.7的45个自然语言ICL任务Todd et al. (2024)的ICL任务Feng & Steinhardt (2024)的binding任务Pile合成的重复随机token序列
指标
few-shot ICL accuracytoken-loss differenceinduction scoreFV score
AI 导读Kayo Yin 与 Jacob Steinhardt 的 arXiv 论文比较 12 个语言模型中的 induction head 与 function vector(FV)head,通过消融发现少样本上下文学习(ICL)性能主要依赖 FV head,模型规模越大越明显。研究还发现两类注意力头彼此关联:训练中不少 FV head 起初表现为 induction head,之后才转向 FV 机制。作者据此推测,induction 机制有助于学习最终驱动 ICL 的更复杂的 FV 机制。

Kayo Yin 与 Jacob Steinhardt 的 arXiv 论文比较 12 个语言模型中的 induction head 与 function vector(FV)head,通过消融发现少样本上下文学习(ICL)性能主要依赖 FV head,模型规模越大越明显。研究还发现两类注意力头彼此关联:训练中不少 FV head 起初表现为 induction head,之后才转向 FV 机制。作者据此推测,induction 机制有助于学习最终驱动 ICL 的更复杂的 FV 机制。

深度解读

6个问题,点开问题读完整回答

  1. 这篇论文试图解决什么问题?

    少样本ICL主要由induction heads还是FV heads驱动。

    少样本 ICL 主要由 induction heads 还是 function vector(FV)heads 驱动?

    • 使用场景:作者称大语言模型只靠提示中的少量示例就能做新任务,这一能力已被用来适配预训练模型。两条机制是:induction heads 找到当前 token 上次出现后的下一个 token 并复制;FV heads 的激活形成任务的紧凑编码,加回计算后可在没有示例时恢复 ICL。
    • 现有分歧:作者把 induction 电路视为少样本 ICL 关键机制的观点,归到 Olsson et al.(2022)、Crosbie 与 Shutova(2024)、Dong et al.(2022)。Todd et al.(2024)与 Hendel et al.(2023)则提出 function vectors。作者认为两边结果分叉,是因为 ICL 的操作定义、指标和模型规模并不相同。
    • 核心观察:两类头若是不同机制,位置和成员应分开;若分数相关,不控制相关会把 FV 样行为算进 induction 消融。
    • 本文工作:在 12 个 decoder-only 模型(70M 到 7B)和附录 A.7 的 45 个自然语言 ICL 任务上,比较两类头的位置、重叠、消融后的少样本准确率,以及 Pythia 训练中的转变。
  2. 有哪些相关研究?

    分歧来自指标、两类头的相关,以及先前研究的模型规模。

    引言和 §2 把相关工作分成 induction 电路、function vectors,以及 ICL 如何被定义。

    Induction heads

    • Elhage et al.(2021)——在 transformer circuits 中识别 induction heads:处理某 token 时注意它上次出现之后的 token。
    • Olsson et al.(2022)——在 1–3 层、只有注意力的小模型上,用上下文第 500 个与第 50 个 token 的损失差衡量 ICL,并称去掉 induction heads 会损害该指标。作者指出该指标并不紧密跟踪少样本表现。
    • Crosbie 与 Shutova(2024)、Bansal et al.(2023)——报告去掉 induction heads 会降低少样本准确率。作者认为这类消融(含他们自己未做排除的初轮消融)未控制与 FV 的相关。Singh et al.(2024)与 Olsson et al. 被作者归为把机制分析放在小模型上的工作。

    Function vectors

    • Todd et al.(2024)——用因果中介从特定注意力头取出任务的紧凑表示,加回后可在没有示例时恢复 ICL。本文的 FV score 沿用该框架。
    • Hendel et al.(2023)——同期提出 in-context learning 会形成 task vectors。论文把它与 Todd et al. 并列为 FV 机制来源,未单独比较二者。

    ICL 的定义

    • Brown(2020)、Dong et al.(2022)、Wei et al.(2023)——被引为“不更新参数、从提示少样本学习”这一用法。作者采用该定义,并把文中的 in-context learning 固定为这一含义。
    • Akyürek et al.(2023)、Von Oswald et al.(2023)——只在引言中被列为研究 ICL 机制的工作。论文未与它们做方法对比。

    基线与数据

    • 消融对照是无消融(clean)和随机头消融(random),另有排除另一类 top 2% 头的设置。任务来自 Todd et al.(2024),以及 Feng 与 Steinhardt(2024)的 binding 任务;token-loss 用 Pile(Gao et al., 2020)。

    作者把与先前结论的分歧主要归于三点:少样本准确率与 token-loss 差值不是同一现象、未控制两类头相关、先前机制研究偏小模型。作者认为其中指标选择影响最大(§2.3)。

  3. 论文如何解决这个问题?

    用两种分数标头,再用带排除的均值消融比因果贡献。

    作者用 induction score 和 FV score 标出两类头,以均值消融比较对少样本准确率的因果贡献,并用 Pythia 检查点看头如何转变。

    识别 induction heads

    • 用 TransformerLens。序列是均匀采样的 50 个随机 token,再原样重复一次。
    • 分数是处理重复段时,对原序列中下一 token 的注意力权重之和;对 1,000 条序列取平均,再除以总注意力质量,得到 0 到 1。
    • 定义为 SI(a,r)=∑i=150 ar′i → ri+1,即头 a 在重复段位置 i 对原序列下一 token 的注意力。

    识别 FV heads

    • 沿用 Todd et al. 的因果中介。对任务 t,先算头在该任务提示上的平均激活;再随机打乱输出标签,输入保持不变。
    • 把该头激活换成任务条件平均激活,看正确输出 y 上的模型量相对打乱提示升高多少。每条提示为 10 对示例加 1 个测试例。
    • 定义为 SFV(a|p̃)=f(p̃|a:=āt)[y]−f(p̃)[y],即替换激活前后,模型赋予正确输出的量之差。
    • §2.2 写在 37 个任务、每任务 100 条上取平均。§4.1 为避免识别与评测泄漏,把这 37 个任务分成 26 个测 FV score、11 个用于评测。两处任务数不一致,论文未解释。

    代表集与消融

    • 各类取分数最高的 2% 作为代表集,沿用 Todd et al.(2024)。直接重叠定义为 100×|IH∩ FV|/|IH| 的文字形式,即两类 top 2% 交集占 induction 代表集的百分比。
    • 均值消融把头输出换成它在任务数据上的平均输出。作者称这是为避开零消融的分布外效应。附录 A.3 另做随机头输出替换和零向量替换。
    • 排除消融:消融 n 个 FV 头时,只取 FV score 最高、且不在 induction score top 2% 内的 n 个头;反向同样。消融比例为 1%–20%。

    两种评测量

    • 少样本准确率:10 对示例加查询,看模型是否给出正确输出。评测共 19 个任务(上述 11 个再加 8 个新任务),每任务 100 条。8 个新任务中,4 个是 Todd et al. 任务的变体,4 个是 Feng 与 Steinhardt(2024)的 binding 任务。
    • token-loss 差值:Pile 上第 50 个 token 的损失减去第 500 个 token 的损失,10,000 条平均。作者相对 Olsson et al. 的差值取反,使更高表示更好,并建议不要再把该差值叫 ICL score。

    训练过程

    • 在 7 个 Pythia 模型的 8 个中间检查点上,记录 top 2% 头的平均分数、单个头的两条分数,以及检查点的少样本准确率。论文写总训练步为 143,000。
  4. 论文做了哪些实验?

    作者称少样本准确率更依赖FV heads,token-loss差值则更受induction消融影响。

    实验设置

    • 模型:Table 2 的 12 个均为被分析对象。Pythia 70M、160M、410M、1B、1.4B、2.8B、6.9B(均为 deduped);GPT-2(117M)、GPT-2 medium(345M)、GPT-2 large(774M)、GPT-2 XL(1.6B);Llama 2(7B,huggingface 名为 Llama-2-7b)。层数、头数和头维度见表 2。
    • 任务:附录 A.7 列 45 个自然语言 ICL 任务。少样本评测为 19 个任务、每任务 100 条、10 个示例。FV score 的任务数,§2.2 写 37,§4.1 写 26。induction score 用 1,000 条合成序列;token-loss 用 Pile 的 10,000 条。
    • 对照:clean;随机头消融。排除设置保留另一类的 top 2%。
    • 指标:少样本准确率;token-loss 差值;induction score;FV score。论文未使用单独的 LLM 评审模型,也未写准确率是严格字符串匹配还是别的判定。
    • 消融范围:分数最高的 1%–20% 头。论文未报告随机种子或重复次数。
    • 训练:7 个 Pythia,各 8 个检查点。Figure 4 的代表模型是 Pythia 160M、Pythia 1B、Pythia 6.9B。

    主结果

    正文未列出各消融比例的准确率读数。Figure 1a、Figure 4 与附录 Figure 8–9 是曲线,转换文本里没有逐点数值,下表只收正文写明的比较。

    条件少样本准确率token-loss 差值
    直接消融,12 个模型引言称 FV 消融更伤准确率,且 12 个模型上一致、更大模型更明显。§4.2 称 induction 消融仍大于随机,比例升高后向 FV 消融靠拢未作为该组主结论
    排除消融,参数量超过 1B作者称保留 top 2% FV 时,induction 消融与随机相当;反向仍降低准确率,差距随规模变大(Figure 4 中行)未按超过 1B 单独报告
    70M;token-loss 低于 160M§2.3 与附录 A.2 称 70M 上两类因果效应相近作者称低于 160M 时,两类相对随机都看不出更大影响
    token-loss 超过 345M未按此阈值单独报告作者称 induction 消融影响大于 FV,差距随规模缩小(Figure 4 下行)
    • 作者认为直接消融里 induction 的影响,主要来自同时具有高 FV score 的头;只要留下 top 2% FV 头,消融 induction 头后仍能做 ICL。Figure 1a 图注针对 Pythia 6.9B 称,induction 消融相对随机只有很小额外影响。这与 §4.2 对直接消融的总述不是同一句话。
    • 作者认为少样本准确率与 token-loss 差值测的是不同现象,并用来调和与 Olsson et al. 的表面矛盾。
    • 作者称 FV 相对 induction 的重要性随规模上升,并称这挑战可解释性中强版本的普遍性假设。

    训练中的转变

    • 测的是 7 个 Pythia、8 个检查点上 top 2% 的平均 induction score、FV score 和少样本准确率(Figure 5,附录 A.11)。
    • 作者称 induction 头约在第 1,000 步出现,分数先陡升,再进入平台或略降;FV 头约在第 16,000 步出现,之后持续上升。少样本准确率约在 induction 头出现时开始上升,之后继续渐增。
    • 单个头(Figure 6):作者称许多最终 FV score 高的头,早期 induction score 也高,随后前者上升、后者下降。作者称没有看到 induction 头在训练中发展出高 FV score。Figure 1b 以 Pythia 6.9B 的 Head 8.17 为例,图中未标逐点分数。

    把 FV 加回隐藏状态

    • 附录 A.6:取 FV score 最高的 2% 头,把它们在任务提示上的平均输出相加得到任务的 FV,加到第 |L|/3 层。40 个任务、每任务 1,000 条先前未见示例。对照是打乱标签,以及随机头的输出。
    • 作者称在多数模型上,加入 FV 能恢复未打乱提示上的表现;例外是 Pythia 2.8B。Figure 14 图注写 recovers most of the model accuracy,并点名该例外。正文未给各模型准确率。
    • 作者认为一个可能原因是头维度:Table 2 中 Pythia 2.8B 为 80,参数量相近而头维度为 128 的模型不同。这是作者的可能解释。

    其他消融与分析

    • 重叠:12 个模型中 7 个的 top 2% 直接重叠为 0,其余为 5%–15%;作者称多数模型中,FV 头的 induction score 约在第 90–95 百分位,反向亦然(Figure 3)。左图灰色为与随机头的重叠,正文未给该系列数值。
    • 层位置:作者称 induction 头多在中前层,FV 头略深(Figure 2、附录 Figure 19),并称“因此更抽象”仍是推测。附录 Figure 19 图注附有 t-stat 与 p-value,如 GPT-2 为 -2.83、0.05,Llama 2 7B 为 0.22、0.83,Pythia 70M 为 nan。正文未说明检验对象。
    • 附录 A.5:induction score 跨规模大体相近;Pythia 1B、1.4B 的最大 FV score 更高,头维度分别为 256 和 128。Pythia 70M 与 Llama 2 的 FV score 很低。作者把高头维度相对参数量标为 FV 强度的一个预测(文中 H6),并称 Llama 2 的低分可能与架构有关,还要更多实验才能确认。
    • 附录 A.3:作者称随机替换和零消融下结论仍成立;正文未给这两类消融的准确率。
    • 附录 A.10:被消融的两类集合重叠随消融头数增加。作者用此说明排除消融有必要。
    • 附录 Figure 15–18 按任务画消融曲线,红线是训练中见过的输出的随机基线;正文未汇总各任务数字。
  5. 有什么可以进一步探索的点?

    作者称猜想C1解释不了低induction的FV头,并建议训练中做干预。

    作者指出的局限与后续方向

    • §6:猜想 C1(induction 头是 FV 头的早期形式)解释不了训练全程 induction score 都低的 FV 头。
    • §6:要再核验 C1,未来工作可以在训练中去掉 induction heads,看 FV heads 是否仍会发展。
    • §6:猜想 C2(FV 头是 induction 与另一机制的组合)面临作者写出的困难:去掉没有高 induction score 的 FV 头仍会损害 ICL,这与“关键功能依赖组合机制”不好调和。
    • §7:若 induction 头是 FV 头的前驱,是什么使这一步成为必要。
    • §7:训练完成后仍然存在的 induction heads 起什么作用。
    • §7:是否还有其他机制能更完整地解释 ICL。

    实验覆盖范围

    • 被分析模型为 Table 2 的 12 个 decoder-only 模型:7 个 Pythia(70M–6.9B)、4 个 GPT-2(117M–1.6B)、Llama 2 7B。
    • 少样本准确率是 19 个任务、每任务 100 条、10 个示例;消融比例 1%–20%,含均值消融,以及附录 A.3 的随机替换和零消融。
    • 训练曲线只在 7 个 Pythia 上,每个 8 个检查点(§5)。
    • induction score 用 1,000 条长度为 50+50 的合成序列;token-loss 用 Pile 的 10,000 条;加回 FV 的实验是 40 个任务、每任务 1,000 条(附录 A.6)。
    • 论文未报告消融重复次数,也未在正文列出 Figure 4 各点的准确率。
  6. 总结一下论文的主要内容

    作者称少样本ICL更随FV heads变化,部分FV头由induction头转来。

    这项工作比较 12 个 decoder-only 模型里,两类被用来解释 ICL 的注意力头。

    • 问题:少样本 ICL 既被联系到复制“上次出现后的下一 token”的 induction heads,也被联系到携带任务向量的 FV heads。作者要看哪一类在因果上支撑少样本准确率。
    • 做法:用重复随机序列上的注意力定义 induction score,用打乱标签后的因果中介定义 FV score。对最高分头做均值消融,并用排除另一类 top 2% 的设置分开相关;再看 Pythia 检查点上单个头的两条分数。
    • 两类头:top 2% 的直接重叠在 12 个模型中有 7 个为 0,其余为 5%–15%;作者称多数模型里,一类头的另一分数约在第 90–95 百分位(Figure 3)。induction 头层位略早,FV 头略深。
    • 消融:作者称直接消融时,FV 头对少样本准确率的影响大于 induction 头,且更大模型上更明显;排除后,超过 1B 的模型上,induction 消融与随机消融相当,FV 消融仍降低准确率(§4.2)。70M 上作者称两类效应相近。
    • 另一指标与训练:超过 345M 时,token-loss 差值更受 induction 消融影响。作者称 induction 头约于第 1,000 步出现,FV 头约于第 16,000 步出现;许多 FV 头由高 induction score 转向高 FV score,反向未观察到。
    • 作者的结论:作者认为少样本 ICL 主要由 FV heads 驱动,并推测更简单的 induction 机制为后来的 FV 机制提供了前驱。作者还认为,指标定义和机制重叠会改变可解释性结论,结果挑战强版本的跨规模普遍性。
阅读原文arxiv.org