Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models¶

会议: NeurIPS 2025
arXiv: 2504.15471
代码: https://github.com/tylerachang/bigram-subnetworks
领域: 可解释性
关键词: Bigram子网络, 机制可解释性, 连续稀疏化, 残差流, 最小电路

一句话总结¶

通过连续稀疏化在Transformer语言模型中找到仅包含~10M参数的bigram子网络，它们集中在第一个MLP层，足以复现bigram预测（\(r>0.95\)），且被消融后模型性能大幅下降，证明这些子网络是语言模型中既必要又充分的最小next-token预测电路。

领域现状：机制可解释性研究已发现了induction head、name mover head等特定电路，但这些电路通常只覆盖特定行为。缺乏一个定义在整个输入空间上的"最小基础电路"，作为研究更复杂电路的起点。

现有痛点：电路研究通常只验证必要性（消融后行为消失）但不验证充分性（电路独立运行后行为是否仍存在）。验证充分性需要在某个已理解的最小电路之上叠加目标电路，但这个最小电路是什么，一直不清楚。

核心矛盾：Transformer已知会在预训练早期过拟合bigram分布，但即使模型后来偏离了bigram预测，bigram信息是否仍编码在模型参数中？以什么形式存在？

切入角度：Bigram预测 \(P(w_i|w_{i-1})\) 是最简单的非平凡next-token预测，在整个输入空间上有定义。如果能找到实现bigram预测的子网络，它就是研究更复杂电路的理想基础。

核心 idea：用连续稀疏化在冻结的LLM中搜索mask，找到仅占0.17%参数但能达到r=0.96 bigram相关的子网络，主要集中在第一个MLP层。

冻结LLM参数 → 用连续稀疏化学习参数mask \(M\) → 最小化masked模型输出与bigram分布的交叉熵 + L1稀疏惩罚 → 得到二值mask定义的子网络。在Pythia (70M-1B) 和 GPT-2 (small-large) 上实验。

连续稀疏化找子网络：
- 每个模型参数对应一个可学习的mask值 \(m \in (-\infty, +\infty)\)，通过sigmoid映射到(0,1)
- 训练过程中逐渐降低sigmoid温度，使mask趋向二值
- 损失：\(\text{CE}(P(x), \text{MaskedModel}_M(x)) + \lambda \|M\|_1/|M|\)
- \(\lambda\) 控制稀疏度，从0到1000变化以观察不同稀疏度下的表现
关键发现：~10M参数的普适性：
- 不管模型大小（70M到1B），bigram子网络在~10M活跃参数处达到性能平台
- Pythia 1B中仅0.17%的非embedding参数就能达到 \(r=0.959\) 的bigram相关
- 说明bigram预测所需的"电路容量"与模型规模无关
结构分析：第一个MLP层的统治地位：
- 在所有模型和预训练检查点中，bigram子网络的大部分参数集中在第一个Transformer MLP层
- 甚至在随机初始化的模型中也是如此——说明这是架构+损失函数的固有偏置
- 机制解释：第一个MLP层负责将激活从"当前token表征"旋转到"next-token预测空间"

模型	子网络参数占比	Bigram \(r\)	全模型 \(r\)
Pythia 70M	~15%	0.961	0.737
Pythia 410M	~2.5%	0.983	0.650
Pythia 1B	0.17%	0.959	0.632
GPT-2 medium	~1%	0.985	0.582
GPT-2 large	~1%	0.986	0.583