跳转至

In-Context Algebra

会议: ICLR2026
arXiv: 2512.16902
代码: algebra.baulab.info
领域: llm_nlp
关键词: in-context learning, mechanistic interpretability, symbolic reasoning, finite groups, transformer mechanisms

一句话总结

本文设计了一个 in-context 代数任务——令 token 成为纯变量、每条序列重新随机分配含义——发现 Transformer 在此设定下不再学习经典的傅里叶/几何表示,而是涌现出三种 符号推理机制(交换复制、单位元识别、闭包消去),并揭示了训练过程中这些能力按阶段性相变依次出现的规律。

研究背景与动机

  1. 固定嵌入的局限:先前大量模型可解释性研究(grokking, 模算术)表明 Transformer 在 token 嵌入中预编码了任务信息(如"108"编码了"能被2整除"),学到的是周期性/傅里叶基的几何策略。
  2. 真正的抽象推理:抽象推理的标志是能处理 含义事先未知的符号。若 token 不携带任何固定语义,模型会学什么策略?
  3. 纯变量设定:作者提出让每条序列中的 token 只是占位变量,通过一个随机映射 \(\varphi_s\) 将有限群元素映射到词表符号(且每条序列映射不同),迫使模型仅从上下文关系推理。
  4. 与 ICL 的关联:这是对 in-context learning 内在机制的深入探索——模型需在上下文中观察"事实"并推断代数结构,而非依赖参数记忆。
  5. 可解释性方法论:作者设计了 5 种目标数据分布 + 因果干预实验,为机制验证提供了严格的方法范式。
  6. 阶段性学习:训练过程中不同能力以相变的形式依次涌现,揭示了 Transformer 学习抽象运算的内在课程。

方法详解

任务框架

给定有限群集合 \(\mathcal{G} = \{G_1, G_2, \ldots, G_m\}\),每条训练序列 \(s\) 的生成过程:

  1. 采样群:从 \(\mathcal{G}\) 中采样子集 \(\mathcal{G}_s\),令 \(H_s = \bigcup \mathcal{G}_s\)\(|H_s| \leq N\)(词表大小)
  2. 随机映射:构造一一映射 \(\varphi_s: H_s \to V\),将群元素随机分配到变量 token
  3. 组装序列:从群中采样乘法事实 \(x \cdot y = z\),通过 \(\varphi_s\) 转为变量语句并拼接

序列格式:

\[s = v_{x_1} v_{y_1} = v_{z_1},\; v_{x_2} v_{y_2} = v_{z_2},\; \cdots,\; v_{x_k} v_{y_k} = v_{z_k}\]

每个事实含 4 个位置:左槽 \(v_{x_i}\)右槽 \(v_{y_i}\)预测 token "="、答案槽 \(v_{z_i}\)

模型配置

  • 架构:4 层自回归 Transformer,每层 8 个注意力头,隐藏维度 1024
  • 训练目标:标准 next-token prediction
  • 序列长度\(k = 200\) 个代数事实(约 1000 个 token)
  • 训练群\(\mathcal{G} = \{C_3, \ldots, C_{10}, D_3, D_4, D_5\}\)(循环群 + 二面体群,阶数 ≤ 10)
  • 词表\(N = 16\) 个变量 token + 特殊 token "=" 和 ","

五种假设机制与目标分布

为消歧模型可能使用的算法,作者设计了 5 种目标数据分布:

分布 测试机制 构造方式
\(\mathcal{D}_{\text{copy}}\) 逐字复制 序列中包含最终事实的副本
\(\mathcal{D}_{\text{commute}}\) 交换复制 包含交换事实 \(yx=z\),无逐字副本
\(\mathcal{D}_{\text{identity}}\) 单位元识别 最终事实含单位元 \(ey=y\),上下文有暴露单位元的事实
\(\mathcal{D}_{\text{associate}}\) 结合律组合 包含可通过结合律推导的最小事实集
\(\mathcal{D}_{\text{cancel}}\) 闭包消去 包含共享左/右槽的所有事实,利用消去律排除

因果验证方法

使用 间接效应 (Indirect Effect, IE) 量化组件重要性:

\[\text{IE}(l,h) = P(v_{\text{target}} \mid a_{s_{\text{clean}}}^{(l,h)} \to s_{\text{corrupt}}) - P(v_{\text{target}} \mid s_{\text{corrupt}})\]

通过在干净/损坏序列对之间 patch 注意力头激活,定位实现各机制的关键组件。

三大核心机制

1. 交换复制 (Commutative Copying):单一注意力头 (Layer 3, Head 6) 负责。当存在逐字副本时,该头注意答案槽并直接提升对应 token 的 logit;当仅存在交换事实 \(yx=z\) 时,该头转而注意交换事实的答案槽。

2. 单位元识别 (Identity Element Recognition):由 查询提升 + 单位元抑制 两个子机制协同完成。Head 3.1 提升查询中两个变量的 logit(查询提升),Head 3.6 抑制已识别的单位元 token 的 logit(单位元抑制),最终正确剩下非单位元变量作为答案。最终层注意力输出的 PCA 第一主成分可清晰分离单位元/非单位元事实。

3. 闭包消去 (Closure-Based Cancellation):计算 \(S_{\text{closure}} - S_{\text{cancel}}\)。闭包子机制追踪与查询变量同群的所有元素;消去子机制利用消去律排除已出现在相同左/右槽事实中的答案。作者训练了 16 维子空间 \(W\),对闭包集进行因果干预,达到 99.8% 的干预准确率。

实验关键数据

主要结果:算法覆盖率与模型性能

机制 训练数据覆盖率 (AUC) Hold-out 覆盖率 (AUC)
逐字复制 67.9%
交换复制 +12.1%
单位元识别 +4.2% 28.7%
闭包消去 +2.7% +39.1%
结合律组合 +3.6% +16.9%
总覆盖 90.4% 84.7%
模型实际准确率 92.4% 87.3%

各目标分布上的模型准确率

数据分布 \(k=50\) \(k=100\)
逐字复制 \(\mathcal{D}_{\text{copy}}\) ~100% 100.0%
交换复制 \(\mathcal{D}_{\text{commute}}\) ~97% 99.0%
单位元识别 \(\mathcal{D}_{\text{identity}}\) ~98% 100.0%
闭包消去 \(\mathcal{D}_{\text{cancel}}\) ~95% 97.0%
结合律组合 \(\mathcal{D}_{\text{associate}}\) ~55% 60.2%

泛化能力

  • 未见群泛化:对训练中未见的所有 8 阶群达到近乎完美的准确率
  • 半群泛化:在非群结构如半群上仍有非平凡准确率
  • 拟群/岩浆:拟群表现下降,岩浆(magma)几乎失败,表明模型依赖群的结构性质

消融分析:阶段性相变

训练过程中的技能习得呈现清晰的 五阶段相变

阶段 习得技能 训练步
结构 token 预测 ("=", ",") 最早出现
➁➂ 群闭包 + 查询提升(单位元50%准确率) 第二阶段
➃➄ 逐字复制 + 交换复制 急剧下降
➅➆ 闭包消去 + 完整单位元识别 联合渐进提升
结合律组合 最后出现

关键发现: - 复制是基础:消去和单位元识别建立在复制能力之上 - 联合涌现:单位元抑制和消去子空间执行类似的"抑制"功能,因此同步学习 - 结合律最难:是最后习得的能力,且准确率仅 60%

因果干预结果

实验 关键头 AIE
逐字复制 Head 3.6 0.91
交换复制 Head 3.6 0.48
其他头最高 < 0.08
闭包子空间干预准确率 16-dim \(W\) 99.8%

亮点

  • 极简而深刻的实验设计:纯变量代数设定巧妙隔离了"嵌入先验"与"上下文推理",是 ICL 机制研究的优质范式
  • 完整的机制解剖:从假设→目标分布设计→覆盖率分析→因果干预→子空间探测,形成闭环验证
  • 相变与课程学习的自然对应:揭示了 Transformer 自发的分阶段技能习得,具有理论启发性
  • 符号 vs 几何策略的依赖性:论证了推理策略取决于任务结构——固定 token 导致几何策略,纯变量导致符号策略
  • 代码与数据开源,实验可复现

局限性 / 可改进方向

  1. 模型规模受限:仅在 4 层小 Transformer 上验证,未知是否适用于大规模预训练 LLM
  2. 结合律学习不足:模型仅学到 60% 的结合律准确率,说明多步推理仍是挑战
  3. 任务偏理想化:有限群代数距离自然语言推理仍有距离,结论是否迁移到更复杂场景需验证
  4. 词表规模限制:仅 16 个变量 token,扩展到更大词表时机制是否保持一致尚未研究
  5. 未探索 CoT:结合链式思维提示可能提升结合律等复杂推理的表现

与相关工作的对比

对比维度 本文 先前工作 (Nanda et al., Zhong et al.)
Token 含义 每序列随机变化的纯变量 固定语义(如数字)
习得策略 符号推理(复制、消去) 傅里叶基 / 几何表示
是否涌现 grokking 阶段性相变但非传统 grokking 典型 grokking
泛化能力 泛化到未见群 泛化到同分布数据
分析方法 因果干预 + 子空间探测 权重/嵌入分析

与 Akyürek et al. (2024) 的 induction head / n-gram head 分析互补——本文的 copying head (3.6) 具有类似的 n-gram 匹配行为,但进一步展示了交换复制和单位元抑制等更高级功能。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ — 纯变量代数设定首次提出,揭示了与传统固定 token 完全不同的推理机制
  • 实验充分度: ⭐⭐⭐⭐⭐ — 5种目标分布 + 因果干预 + 子空间探测 + 阶段分析,验证极为充分
  • 写作质量: ⭐⭐⭐⭐⭐ — 图表精美,逻辑清晰,Figure 4/5/6 的可视化尤其出色
  • 价值: ⭐⭐⭐⭐ — 对 ICL 机制研究有重要启示,但与实际 LLM 应用的桥接尚待验证