Learning to Read Out: Unembedding Dynamics in Language Model Pretraining¶

会议: NeurIPS 2026
arXiv: 2610.21025
代码: https://github.com/hematteo/learning-to-read-out
领域: 预训练
关键词: 语言模型预训练、解嵌入动态、参数轨迹跨编码器、表征可用性与表达、发展性可解释性

一句话总结¶

本文揭示了语言模型预训练早期最后一层隐状态已编码任务信息但自身解嵌入矩阵（Readout/Unembedding）尚未学会表达的滞后脱节现象，并提出基于纯权重轨迹的参数轨迹跨编码器（Parameter-Trajectory Crosscoding）定量刻画了 Readout 结构的早期快速重组窗口与动态演进规律。

研究背景与动机¶

在大语言模型的发展性可解释性（Developmental Interpretability）研究中，一个长期存在的基础问题是：模型究竟在预训练的哪个时间点真正习得了某种特定的语言或推理能力？以往的研究通常采用行为学指标或对数概率差异（Logit Lens / Surprisal）来为能力的习得“打时间戳”，例如观察模型何时开始在特定上下文下稳定输出正确的目标词元、何时从表面词汇复制转向真实跨语言翻译，或者何时突现特定的上下文学习回路（如 Induction Heads）。然而，所有这些直接基于预测输出或 Logit 投影的评估方式，都不可避免地依赖于模型最后一层的解嵌入矩阵（Unembedding Matrix，即 Readout $W_U$）。

这引出了一个被普遍忽视的关键矛盾：Readout 自身也是在整个预训练过程中随梯度逐步学习并演进的参数矩阵。对词元 $v$ 而言，其预测对数几率由隐状态 $h$ 与解嵌入行向量 $w_v$ 的内积给出：$\ell_v = h \cdot w_v$。在两个训练检查点之间，Logit 的变化量 $\Delta \ell_v$ 实际上同时包含了隐状态表征的改变 $\Delta h \cdot w_v$、解嵌入向量的改变 $h \cdot \Delta w_v$ 以及二者联合演进的交叉项 $\Delta h \cdot \Delta w_v$。若仅根据输出端是否生成正确词元来判定能力习得时刻，就会将“内部隐状态是否已经表征了任务所需信息”与“解嵌入矩阵是否已经学会将该方向正确映射至对应词元”彻底混淆。特别是在预训练最初期（例如前 1% 步数内），Readout 仍在经历剧烈的几何重构与尺度漂移，此时模型的隐状态可能早已编码了关键语法或语义属性，却由于未成熟的 Readout 无法在输出端体现。

为了打破这种测量混淆，必须将 Readout 的发展轨迹作为一个独立的研究对象进行直接解耦与追踪。以往基于激活值的跨编码器需要海量前向传播语料且受采样词元分布偏置限制，无法完整刻画全词表的参数权重演化。本文的核心 idea 是：将跨编码器直接拟合于全词表解嵌入矩阵在预训练所有检查点间的参数权重轨迹，无需激活样本即可解构 Readout 演化模式，并结合隐状态线性探测（Availability）与跨检查点 Readout 替换实验（Native Expression），严格解耦表征习得与输出表达的时序差。

方法详解¶

整体框架¶

本文构建了一套用于解耦与分析解嵌入矩阵动态演化、以及分离隐状态可用性与输出端表达能力的完整研究体系。其核心工作流分为两个主要阶段：第一阶段是对预训练全过程中保存的各检查点解嵌入矩阵进行标准化与“参数轨迹跨编码器（Parameter-Trajectory Crosscoding）”拟合，将高维行向量的演化轨迹分解为可跨时间步追踪的稀疏特征方向；第二阶段则是在固定检查点隐状态下，对比任务线性探针（衡量表征可用性 Availability）、原生解嵌入头（Native Expression）以及跨检查点解嵌入替换（Readout Swap），精确测量表达滞后差距。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多检查点解嵌入矩阵<br/>{W_U^(t)} 及全词表行向量"] --> B["行中心化与尺度归一化<br/>消除全局均值偏移与范数增长"]
    B --> C["参数轨迹跨编码器<br/>共享编码与时间步特异解码器"]
    C --> D["重组窗口与生命周期分析<br/>几何均值突变分数与六大轨迹画像"]
    D --> E["双轨测量与解耦诊断<br/>线性探针可用性 vs Readout替换表达"]
    E --> F["学习率因果干预实验<br/>调节Readout独立学习率验证时序收敛"]

关键设计¶

1. 参数轨迹跨编码器：基于纯权重的全词表跨时间步稀疏字典学习

针对传统稀疏自编码器（SAE）在不同检查点独立训练会导致特征索引随机乱序、无法长期追踪单一特征轨迹的缺陷，以及激活跨编码器必须依赖激活采样语料且难以覆盖生僻词的问题，本文提出了参数轨迹跨编码器。方法将预训练各个检查点 $t \in \mathcal{T}$ 的解嵌入行向量 $W_U^{(t)}[v, :]$ 统一进行词表级均值中心化与均方根缩放，得到归一化向量 $x_v^{(t)}$。跨编码器对每个词元 $v$ 学习一个在所有检查点间完全共享的高维预激活代码向量 $z_v \in \mathbb{R}^F$（$F = d_{\mathrm{SAE}}$ 为字典维度），并在每个具体检查点 $t$ 采用时间步特有的解码器矩阵 $D^{(t)} \in \mathbb{R}^{F \times d_{\mathrm{model}}}$ 与门控阈值进行重建： $$z_v = \sum_{t \in \mathcal{T}} \left(x_v^{(t)} E^{(t)} + b_E^{(t)}\right), \quad a_{v, f}^{(t)} = z_{v, f} \cdot \mathbf{1}\left\{z_{v, f} \|D_f^{(t)}\|_2 > \theta_f\right\}, \quad \widehat{x}_v^{(t)} = a_v^{(t)} D^{(t)} + b_D^{(t)}$$ 门控函数引入了 JumpReLU 机制与解码器行向量范数 $\|D_f^{(t)}\|_2$。这一精妙设计保证了特征能够随训练推进自发涌现或衰减，同时对特征代码与解码器之间的正标量缩放保持完全不变性。在损失函数上，结合了均方误差重建项 $\mathcal{L}_{\mathrm{rec}}$ 与饱和激活率惩罚项 $\mathcal{L}_s = \sum_f \bar{q}_f (1 + \bar{q}_f / \alpha)$（其中 $\bar{q}_f$ 为包含 $\tanh$ 软计数的平均使用频率），既惩罚过度泛化的冗余特征，又避免了标准 $L_1$ 惩罚过度压低强特征尺度的收缩偏差。

2. 几何重组分数与特征生命周期剖面：定量界定 Readout 演化突变窗口

为了精准定位解嵌入空间发生最剧烈质变的时期，本文定义了综合重组分数（Reorganization Score）。在相邻检查点之间同时测量六种互补的几何形变统计量（包括解码器范数变化幅度、解码器方向旋转角度、以及无字典的顶层子空间漂移等），取其窗口内归一化和的几何平均值。只有当多种维度的形变同时发生时，该综合得分才会达到极值。通过构建时间维度独立循环平移的原假设检验（Circular-shift Nulls），算法可以严格检验重组峰值的统计显著性。在此基础上，论文通过对归一化解码器范数轨迹 $\|D_f^{(t)}\|_2 / \max_t \|D_f^{(t)}\|_2$ 进行生命周期分类，归纳出早期衰减型（Early-decaying）、持久活跃型（Persistent）以及晚期涌现型（Late-emerging）等典型轨迹模式。

3. 双轨解耦诊断体系：线性探针可用性对比跨检查点 Readout 替换

为了严谨区分隐状态中“信息是否存在”与输出端“信息是否被表达”，本文构建了双轨测试系统： - 表征可用性（Availability, $A$）：在固定检查点 $t$ 的最后一层经过最终归一化后的隐状态 $h_p^{(t)}$ 上训练线性分类探针（如收缩线性判别分析 Shrinkage-LDA 或 Logistic 回归），直接预测底层任务标签（如主谓一致中的单复数标记）。该指标反映了表征空间是否已呈线性可分，完全不经过 Readout。 - 原生表达（Native Expression, $E$）：计算模型在自身检查点解嵌入矩阵 $W_U^{(t)}$ 下对于正确答案 $y^+$ 与错误干扰词 $y^-$ 的输出 Margin，即 $m^{(t, t)} = h_p^{(t)} \cdot (W_U^{(t)}[y^+, :] - W_U^{(t)}[y^-, :])$，并以二元对比准确率评估输出偏好。 - Readout 替换与 Oracle 扫描（Readout Swap & Oracle Sweep）：保持隐状态 $h_p^{(t)}$ 严格冻结，用其他检查点 $s \neq t$ 的解嵌入矩阵 $W_U^{(s)}$ 替代原有 Readout 进行解码：$\ell_p^{(t, s)} = h_p^{(t)} (W_U^{(s)})^\top$。通过遍历中间所有中间检查点 $s$，测量 Oracle 最佳 Readout 能否在不改变表征的情况下唤醒未被原生表达的任务性能。

4. 稀疏局部归因干预与学习率因果验证

为验证特定 Readout 特征方向与输出 Margin 之间的因果联系，方法通过跨编码器将最终检查点重构解嵌入矩阵分解为各个特征贡献的线性和，在数据子集上计算各特征对对比 Margin 的均值归因量： $$\operatorname{attr}_f = \mathbb{E}_{n \in A} \left[ \left(a_{y^+, f}^{(143\text{k})} - a_{y^-, f}^{(143\text{k})}\right) \left(h_n^{(1000)} \cdot D_f^{(143\text{k})}\right) \sigma^{(143\text{k})} \right]$$ 通过对贡献最高的前 8 个特征执行消融（Ablation）与孤立保留（Preserve-only），并与匹配范数、激活率及归因符号的随机特征对照组相比，严格建立特定解嵌入方向对任务表现的因果支撑。进一步地，为了验证 Readout 自身优化速率是否为导致时序差的决定性因素，设计了独立调节 Readout 学习率倍率（$0.25\times, 1\times, 4\times$）的受控预训练实验。

实验关键数据¶

主实验¶

论文在 Pythia 家族（160M、1B、6.9B）以及 OLMo-2-7B 等开源大模型上进行了全面评估。字典在解释全词表权重轨迹方差上达到了 80.8%–92.0%，且每行平均仅激活 1.17%–2.26% 的稀疏特征。主谓一致（SVA）与间接宾语识别（IOI）等核心任务展现了隐状态可用性与原生表达之间的显著解耦。

在 Pythia-6.9B 及 Pythia-1B 上的代表性任务对比如下（表 1）：

模型	任务族	样本数 $n$	初始可用性 $A_1@0$	Step 512 可用性 $A_1$	Step 512 原生表达 $E$	Step 512 Oracle最佳 $B$	最终步原生表达 $E$
Pythia-6.9B	SVA 基础（Simple）	320	0.60	0.93	0.50	0.49	1.00
Pythia-6.9B	SVA 跨介词短语（Across-PP）	2,000	0.43	0.95	0.49	0.48	0.97
Pythia-6.9B	BLiMP 照应语一致性	1,455	0.60	0.63	0.48	0.48	0.99
Pythia-1B	SVA 基础（Simple）	320	0.45	0.97	0.50	0.49	1.00
Pythia-1B	SVA 跨介词短语（Across-PP）	2,000	0.53	0.98	0.48	0.48	0.97
Pythia-1B	SVA 主语关系从句（Subject-RC）	1,200	0.56	1.00	0.22	0.39	0.79

注：$A_1$ 为 Shrinkage-LDA 线性探针准确率（表征可用性），$E$ 为原生 Readout 二元对比准确率（随机猜测基线为 0.50），$B$ 为中间检查点 Oracle 最佳解嵌入替换准确率。

在多语言维基百科语料上测试 Pythia-1B 早期隐状态的解嵌入替换（Readout Swap），利用 Step 1000 的 Readout 解码 Step 256 与 Step 512 的隐状态，总体负对数似然（NLL）分别显著下降 1.503 nats 与 0.147 nats，9 种语言无一例外均偏好更成熟的中间 Readout。

消融与因果干预实验¶

1. 稀疏方向局部干预实验（Pythia-1B Step-1000 隐状态 + 最终重构 Readout）

任务族	原始终端Readout	重构终端Readout	Top-8特征消融	符号匹配对照消融	仅保留Top-8特征	符号匹配对照保留
主谓一致（SVA）	0.925	0.900	0.487	0.931	0.938	0.463
大于关系补全（Greater-than）	0.529	0.512	0.084	0.500	0.894	0.507
间接宾语识别（IOI）	0.455	0.429	0.394	0.424	0.481	0.420
关系事实（Relational Facts）	0.269	0.231	0.154	0.231	0.923	0.308

2. 31M 受控预训练消融（调节 Readout 独立学习率与 Warmup）

实验配置	字典解释方差 EV	平均激活数 $L_0$	最终验证损失 (nats)	中位特征峰值步数	最大范数更替发生步数	表达差距弥合步数
$0.25\times$ Readout LR	0.905	193.9	3.405	128	2048	2048
$1.0\times$ 基线基准	0.891	177.2	3.381	64	1024	1024
$1.0\times$ 减半预热 (Warmup 715步)	0.890	175.5	3.350	64	1024	1024
$4.0\times$ Readout LR	0.874	168.5	3.369	32	512	512

关键发现¶

重组窗口极度收敛于训练最初期：在 Pythia 所有尺度模型（160M、1B、6.9B）中，综合重组分数均高度一致地在 Step 512 至 1.6k（对应预训练进度的 0.4%–1.1%） 形成显著尖峰（$p \le 0.04$），与学习率预热结束点（Step 1430）高度重合。在此窗口内，WordNet 35 个超感大类在 Readout 向量中的线性可分性从随机基线剧增至成熟水平（中位平衡准确率从 0.499 跃升至 0.874）。
可用性与表达能力的严重时序脱节：在 Pythia-6.9B 的主谓一致任务中，Step 512 时线性探针准确率已高达 0.93–0.95，而模型自身的原生表达准确率仍在 0.49–0.50（完全处于随机概率）。不仅原生 Readout 无法输出，即使扫描全部中间检查点 Readout 也无法解码，表明早期的语言能力突现曲线很大程度上反映的是 Readout 对齐滞后。
解嵌入学习率对发展时序具有决定性因果控制：受控预训练表明，仅提高 Readout 权重自身的学习率乘数（从 $0.25\times$ 到 $4\times$），特征范数突变峰值步数从 128 严格单调前移至 32，任务表达滞后差距从 Step 2048 提前至 Step 512 闭合；而在训练终点（10B Tokens），不同学习率配置的模型验证损失差距在 0.06 nats 以内。

亮点与洞察¶

将字典学习的视角从激活流拓展至静态权重轨迹：传统跨编码器必须在前向传播中截获百万量级的激活向量，计算昂贵且受限于语料词频；本文提出的参数轨迹跨编码器直接以全词表解嵌入行向量在检查点网格上的演变为输入，不仅覆盖率达到 100%，而且纯靠离线权重即可训练，为研究模型权重动力学提供了全新的极轻量分析范式。
重新审视“能力突现”与“早期训练评估”的时间基准：诸如 Induction Heads 形成（约前 1%–2% 步数）或语法结构对齐等经典发展阶段，恰好与本文定位的 Readout 重组窗口高度重合。很多所谓的“模型在某一步骤突然学会了某能力”，本质上可能只是 Readout 终于建立起了输出读出通道。
模型优化与评估层面的实用启发：在预训练配方筛选或利用小模型早期步数做 Scaling Law 预测时，评测指标严重受限于未成熟的 Readout。适当调大 Readout 学习率可以在不损害最终收敛损失的前提下加速模型表达能力的成熟，更早暴露底层隐状态的真实表征潜质。

局限与展望¶

作者承认的局限：受算力限制，关于调节 Readout 学习率能前移发展时序的因果干预实验目前仅在 31M 参数模型、10B tokens 规模上完成，而 31M 模型中 Readout 参数占比（约 42%）显著高于 1B 模型（约 10%），大规模超参数规律仍待系统验证；此外，任务评测主要基于二元对比 Logit Margin，尚不能完全等同于全词表自由生成的开放式涌现。
潜在研究空白与改进方向：解嵌入矩阵与最后一层 LayerNorm 之间存在数学标量对称性，未来可结合层级联合规范化进一步厘清权重绝对范数与特征尺度的物理意义；同时，该方法可天然推广至词表动态扩展场景（如在预训练模型基础上扩增多语言 Token），针对性赋予新增词元行更高的独立学习率，有望显著缓解冷启动适应迟缓。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出参数轨迹跨编码器，并开创性地指出语言模型发展分析中 Readout 动态演进所造成的长期测量混淆。
实验严谨性: ⭐⭐⭐⭐⭐ 涵盖 4 种开源大模型与自研受控预训练，设计了循环平移零假设、正交框架对齐以及多重随机特征消融对照，论证极为扎实。
写作质量: ⭐⭐⭐⭐⭐ 概念界定清晰准确，公式推导严谨自洽，散文叙述连贯且逻辑链条层次鲜明。
实用价值: ⭐⭐⭐⭐⭐ 为预训练动态解释、Logit Lens 分析避坑以及优化器特定层超参调优（如解嵌入独立学习率）提供了极具指导意义的结论。

实验配置	字典解释方差 EV	平均激活数 \(L_0\)	最终验证损失 (nats)	中位特征峰值步数	最大范数更替发生步数	表达差距弥合步数
\(0.25\times\) Readout LR	0.905	193.9	3.405	128	2048	2048
\(1.0\times\) 基线基准	0.891	177.2	3.381	64	1024	1024
\(1.0\times\) 减半预热 (Warmup 715步)	0.890	175.5	3.350	64	1024	1024
\(4.0\times\) Readout LR	0.874	168.5	3.369	32	512	512