In-Context Algebra¶
会议: ICLR2026
arXiv: 2512.16902
代码: algebra.baulab.info
领域: 其他
关键词: in-context learning, mechanistic interpretability, symbolic reasoning, finite groups, transformer mechanisms
一句话总结¶
本文设计了一个 in-context 代数任务——令 token 成为纯变量、每条序列重新随机分配含义——发现 Transformer 在此设定下不再学习经典的傅里叶/几何表示,而是涌现出三种 符号推理机制(交换复制、单位元识别、闭包消去),并揭示了训练过程中这些能力按阶段性相变依次出现的规律。
研究背景与动机¶
固定嵌入的局限:先前大量模型可解释性研究(grokking, 模算术)表明 Transformer 在 token 嵌入中预编码了任务信息(如"108"编码了"能被2整除"),学到的是周期性/傅里叶基的几何策略。
真正的抽象推理:抽象推理的标志是能处理 含义事先未知的符号。若 token 不携带任何固定语义,模型会学什么策略?
纯变量设定:作者提出让每条序列中的 token 只是占位变量,通过一个随机映射 \(\varphi_s\) 将有限群元素映射到词表符号(且每条序列映射不同),迫使模型仅从上下文关系推理。
与 ICL 的关联:这是对 in-context learning 内在机制的深入探索——模型需在上下文中观察"事实"并推断代数结构,而非依赖参数记忆。
可解释性方法论:作者设计了 5 种目标数据分布 + 因果干预实验,为机制验证提供了严格的方法范式。
阶段性学习:训练过程中不同能力以相变的形式依次涌现,揭示了 Transformer 学习抽象运算的内在课程。
方法详解¶
整体框架¶
作者构造了一个纯变量的有限群乘法任务:每条序列把群元素经一次性随机映射打散成无固定语义的 token,再用一个 4 层自回归 Transformer 做 next-token prediction。模型既无法依赖嵌入里的语义先验,只能从上下文里的乘法事实推断代数结构,作者再用五种目标分布配合因果干预,把模型实际用到的算法逐一拆解出来。
关键设计¶
1. 纯变量代数任务:把语义先验从 token 里彻底剥离。 以往可解释性研究里 token 含义是固定的(比如数字"108"天然编码"能被 2 整除"),模型于是学到傅里叶基这类几何捷径,看不出它是否真在"推理"。本文给定有限群集合 \(\mathcal{G}=\{C_3,\ldots,C_{10},D_3,D_4,D_5\}\),每条序列 \(s\) 先采子集 \(\mathcal{G}_s\) 使并集 \(H_s=\bigcup\mathcal{G}_s\) 满足 \(|H_s|\le N=16\),再造一一映射 \(\varphi_s:H_s\to V\) 把群元素随机分配到 16 个变量 token——关键在于每条序列的 \(\varphi_s\) 都不同,同一个 token 在不同序列里含义完全不同。序列由 \(k=200\) 条乘法事实拼成,形如 $\(s = v_{x_1} v_{y_1} = v_{z_1},\; \cdots,\; v_{x_k} v_{y_k} = v_{z_k}\)$ 每条事实占 4 个位置(左槽 \(v_{x_i}\)、右槽 \(v_{y_i}\)、预测点 "="、答案槽 \(v_{z_i}\)),模型须在 "=" 处预测答案。这样模型唯一能依赖的就是上下文里事实之间的关系,强迫它做符号推理而非套用几何模板。
2. 五种目标分布:把"它到底用了哪种算法"变成可控实验。 模型答对一道题可能来自好几种不同算法,光看准确率无法区分。作者为每种假想机制构造一个目标数据分布,在该分布上只有对应算法才可解、其它捷径被堵死:\(\mathcal{D}_{\text{copy}}\) 上下文里直接放最终事实的逐字副本(测逐字复制);\(\mathcal{D}_{\text{commute}}\) 只放交换事实 \(yx=z\) 而无逐字副本(测交换复制);\(\mathcal{D}_{\text{identity}}\) 让最终事实含单位元 \(ey=y\) 且上下文有暴露单位元的事实(测单位元识别);\(\mathcal{D}_{\text{cancel}}\) 放齐共享左/右槽的事实、逼模型用消去律排除(测闭包消去);\(\mathcal{D}_{\text{associate}}\) 只给可经结合律推导的最小事实集(测结合律组合)。在哪个分布上准确率高,就说明模型掌握了对应机制。
3. 因果验证:用间接效应与激活 patching 定位实现机制的注意力头。 知道模型会某种算法,还要知道是哪些组件在执行它。作者在一对干净/损坏序列之间做激活 patching,用间接效应 (Indirect Effect) 量化某个头 \((l,h)\) 的贡献:$\(\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}})\)$ 即把干净序列里该头的激活移植到损坏序列后,正确 token 概率回升多少。对涉及集合状态的机制(如闭包),作者进一步训练一个 16 维子空间 \(W\) 做子空间级干预,用以验证模型是否真的在内部维护了某个代数集合。
4. 三大涌现机制:复制、单位元识别、闭包消去。 上述方法最终解剖出模型实际依赖的三种符号机制。交换复制由单个头 (Layer 3, Head 6) 承担:有逐字副本时它注意答案槽、直接抬高对应 token 的 logit;只剩交换事实 \(yx=z\) 时它转而注意交换事实的答案槽——同一个头自适应地处理两种情况。单位元识别靠"查询提升 + 单位元抑制"两个子机制配合:Head 3.1 抬高查询中两个变量的 logit,Head 3.6 压低已识别出的单位元 token,剩下的非单位元变量便是答案;最终层注意力输出的 PCA 第一主成分能干净地分开含/不含单位元的事实。闭包消去对应一个 \(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 导致几何策略,纯变量导致符号策略
- 代码与数据开源,实验可复现
局限与展望¶
- 模型规模受限:仅在 4 层小 Transformer 上验证,未知是否适用于大规模预训练 LLM
- 结合律学习不足:模型仅学到 60% 的结合律准确率,说明多步推理仍是挑战
- 任务偏理想化:有限群代数距离自然语言推理仍有距离,结论是否迁移到更复杂场景需验证
- 词表规模限制:仅 16 个变量 token,扩展到更大词表时机制是否保持一致尚未研究
- 未探索 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 应用的桥接尚待验证