大规模解读语言模型隐藏状态

# 大规模语言模型隐藏状态解读

来源：https://arxiv.org/html/2608.10260

Jordan Pettyjohn1, Mansi Sakarvadia1, Nathaniel Hudson2,3, Daniel McKenzie4, Kyle Chard1,3, Ian Foster1,3

###### 摘要

透镜方法通过将中间激活映射到输出词表来解读大型语言模型（LLMs），揭示下一个词元预测如何在网络中逐步形成。经过训练的透镜仍然代价高昂：仿射变换器的参数随模型宽度呈二次方增长，而精确的全词表 Kullback–Leibler（KL）训练则主导内存消耗。因此，此前的训练透镜仅应用于不超过 20B 参数的模型，并且仍局限于特定的组件类型。我们提出 OmniLens，它将统一的透镜族应用于任意模型宽度的激活——无论是残差流、注意力还是 MLP——并结合了两种独立的扩展技术。首先，低秩变换器使每个透镜的参数增长随模型宽度呈线性，并将可训练参数最多减少 98.4%。其次，Subset-KL 仅物化选中的词表 logits：其 Top-k 模式将峰值训练内存最多削减 70%，而其重要性采样变体为完整 KL 保留无偏随机梯度。这些节省使得我们能够在 LLaMA-3.3-70B 上训练一个包含 482 个透镜的稠密集成，其覆盖范围是相同深度下残差流设计的 6 倍。模型范围的覆盖揭示了单组件透镜无法看到的现象：行为最可见的组件不一定是最佳干预位置，而最有效的干预位于先前透镜研究考察的注意力头之外。在三个案例研究（提示注入检测、多跳记忆注入和毒性定位）中，OmniLens 以显著更低的成本复现了关键已发表结果。

参见图注 图 1：全秩透镜参数随 O(Layer×d²) 增长，在 LLaMA-3-405B 上超过 2000 亿。

参见图注 图 2：OmniLens 概览。钩子被放置在模型中任意用户定义的位置；每个透镜应用公式 (3)（https://arxiv.org/html/2608.10260#S3.E3）中的低秩变换器，并根据第 4 节（https://arxiv.org/html/2608.10260#S4）的 Subset-KL 目标针对模型最终分布进行训练。

## 1 引言

理解语言模型如何形成预测、行为在其计算中的何处出现，对于解释和控制语言模型至关重要（Orgad 等人 2026（https://arxiv.org/html/2608.10260#bib.bib75）；Shapira 等人 2026（https://arxiv.org/html/2608.10260#bib.bib76））。*透镜*通过将中间激活解码为词表分布提供直接视角，揭示下一个词元预测如何在网络中演化，以及行为在何处可以被检测或影响。透镜可以解码多种中间激活；我们将为透镜提供输入的模型组件称为其*钩点*。例如，Tuned Lens 读取残差流（Belrose 等人 2023（https://arxiv.org/html/2608.10260#bib.bib6）），而 Attention Lens 读取单个注意力头（Sakarvadia 等人 2023（https://arxiv.org/html/2608.10260#bib.bib7））。与需要为每个目标提供标注样本的概念特定分类探针（Alain and Bengio 2018（https://arxiv.org/html/2608.10260#bib.bib35）；Hewitt and Manning 2019（https://arxiv.org/html/2608.10260#bib.bib84）；Belinkov 等人 2017（https://arxiv.org/html/2608.10260#bib.bib83））不同，透镜可以以自监督方式训练：模型自身的最终分布提供目标，因此无需人工标注。一个训练好的透镜因此可以支持多种下游分析。当以稠密方式应用于整个模型时，即每一层有多个钩点，这种能力最为宝贵。然而，训练这样一个透镜集成至今仍代价高昂。具体而言，训练好的透镜在规模应用时面临三个挑战。

*参数。*每个钩点的全秩变换器成本为 O(d²)，其中 d 是模型的隐藏维度，因此对大型模型的稠密覆盖可能接近模型本身（图 1（https://arxiv.org/html/2608.10260#S0.F1））。

*内存。*透镜训练使用模型自身最终分布与透镜分布之间的 KL 散度作为训练目标，每个词元需要物化两个词表大小的分布。这会迅速耗尽可用显存。

*专用性。*现有透镜各自读取一种钩点类型，每种新钩点都需要新的透镜族。

这些挑战共同限制了训练透镜仅能应用于具有稀疏覆盖的小型模型。我们用 *OmniLens* 解决全部三个挑战，为（接近）前沿规模的模型解锁稠密透镜覆盖。OmniLens 是一个与钩点无关的框架，将统一的透镜族应用于残差、注意力和 MLP 激活。OmniLens 通过两种并行技术使稠密覆盖变得可行。首先，低秩变换器将每个钩点的参数数量从 O(d²) 降至 O(rd)，其中 r 是目标秩。其次，一类灵活的 KL 散度近似——我们称之为 Subset-KL 目标——每个词元仅物化透镜分布的一小部分，大幅降低内存需求。我们将这些贡献结合在一个开源框架中¹¹¹代码：OmniLens（训练框架）https://github.com/pettyjohnjn/OmniLens；Hookbox（钩点插桩）https://github.com/pettyjohnjn/hookbox；IndexedLogits（融合 CUDA 核）https://github.com/pettyjohnjn/indexed_logits；SubsetKL（目标与估计器）https://github.com/pettyjohnjn/subset-klo，并以显著更低的成本复现先前工作的关键指标（Belrose 等人 2023（https://arxiv.org/html/2608.10260#bib.bib6）；Sakarvadia 等人 2024（https://arxiv.org/html/2608.10260#bib.bib32）；Pettyjohn 2025（https://arxiv.org/html/2608.10260#bib.bib109））。OmniLens 覆盖的透镜数量是 tuned lens 基线的六倍，同时成本更低，在 LLaMA-3-70B 上可训练参数减少 90.5%，在 GPT-2 上峰值训练内存最高降低 70%。在三个可解释性案例研究中，它匹配现有透镜框架的应用效果，而模型范围的覆盖进一步揭示了单组件透镜无法看到的内容：行为最可见的钩点不一定是最佳干预位置。为了测试该方法的极限，我们在 LLaMA-3.1-405B 上运行了八个优化步骤，据我们所知，这是在此规模上首次在现有硬件上对训练透镜优化的实证演示，确立了训练路径在前沿规模上的可执行性。

## 2 背景与相关工作

我们研究预训练自回归 Transformer 语言模型（Vaswani 等人 2017（https://arxiv.org/html/2608.10260#bib.bib1））。令 x 为来自词表 V 的长度为 T 的词元序列，P(⋅|x) 表示模型在给定位置的下一个词元分布。模型具有隐藏维度 d 和 L 层，钩点暴露残差流状态、注意力和 MLP 输出，或单头输出，具体取决于所需分辨率。我们记 H_{ℓ,u}∈R^{T×d_u} 为第 ℓ 层组件 u 的输出，h_{ℓ,u}∈R^{d_u} 为单个位置的激活；透镜逐位置应用，我们省略位置索引。我们依据决定透镜方法可扩展性的三个选择来组织相关工作：读取位置、激活变换方式以及训练目标的计算方式。

##### 透镜公式与组件专用性。
透镜是一个辅助解码器，将中间激活映射为模型词表上的分布。我们将注意力限制在重用模型冻结的最终归一化 η（LayerNorm 或 RMSNorm，与模型读出中的使用方式完全相同）和去嵌入 W_U∈R^{|V|×d} 的线性透镜上：

Q_{ℓ,u}(⋅|x)=softmax(W_U η(L_{ℓ,u} h_{ℓ,u}+b_{ℓ,u}))    (1)

其中 L_{ℓ,u}∈R^{d×d_u} 和 b_{ℓ,u}∈R^d 是学得的。logit lens（Nostalgebraist 2020（https://arxiv.org/html/2608.10260#bib.bib5））直接将模型冻结的读出应用于模型宽度的残差流状态，对应 L_{ℓ,u}=I 和 b_{ℓ,u}=0。由于这种直接读出可能无法很好地代理最终预测，后续方法学习仿射变换器，使 Q_{ℓ,u}(⋅|x) 逼近 P(⋅|x)（Belrose 等人 2023（https://arxiv.org/html/2608.10260#bib.bib6）；Din 等人 2024（https://arxiv.org/html/2608.10260#bib.bib62）；Pal 等人 2023（https://arxiv.org/html/2608.10260#bib.bib59））。这些方法针对残差流，而 Attention Lens（Sakarvadia 等人 2023（https://arxiv.org/html/2608.10260#bib.bib7））为单个注意力头学习解码器。每种构造都绑定于一个组件族，因此读取不同组件需要另一种透镜设计；这种变换器与组件之间的耦合正是 OmniLens 所解决的专用性瓶颈。Backward Lens（Katz 等人 2024（https://arxiv.org/html/2608.10260#bib.bib120））将梯度而非激活投影到词表空间，证明了此类投影具有低秩结构。

##### 参数高效变换器。
稠密仿射变换器具有 d_u d + d 个学得参数，且每个组件在所有层上拟合一个变换器会很快变得不切实际：在我们研究的 (6L+2) 钩点密度下，透镜集接近基础模型参数数量的一半（图 1（https://arxiv.org/html/2608.10260#S0.F1）），而更精细的构造如 Attention Lens 的逐头解码器可能直接超过研究对象模型。低秩适应（LoRA）（Hu 等人 2021（https://arxiv.org/html/2608.10260#bib.bib27））最初用于参数高效微调，冻结模型权重并学习每个权重的低秩更新，将可训练参数减少多个数量级。低秩参数化已与透镜结合，但现有实现仍保留其全秩前身的组件专用性。LoRA Lens（Pettyjohn 2025（https://arxiv.org/html/2608.10260#bib.bib109））将逐头 Attention Lens 解码器分解为对冻结去嵌入的秩 r 更新，应用于最高 8B 的模型；同期工作（Trimigno 等人 2026（https://arxiv.org/html/2608.10260#bib.bib117））在最高 32B 的模型上训练低秩残差流透镜。两者都仅在单个组件族内应用低秩，且均未提供覆盖残差、注意力和 MLP 组件的统一变换器架构，也未解决限制稠密覆盖规模的词表侧内存成本。完整的透镜分类见附录 A（https://arxiv.org/html/2608.10260#A1）。

##### 内存高效蒸馏。
大多数透镜框架使用词元级 Kullback-Leibler（KL）散度来衡量 P(⋅|x) 与 Q_{ℓ,u}(⋅|x) 之间的差异：

D_KL(P‖Q_{ℓ,u}) = Σ_{v∈V} P(v|x) log(P(v|x)/Q_{ℓ,u}(v|x))   (2)
               = E_{v∼P(⋅|x)}[log(P(v|x)/Q_{ℓ,u}(v|x))],

我们将 P(⋅|x) 称为*教师*，Q_{ℓ,u}(⋅|x) 称为*学生*。对于单个输入 x，物化全部 T|V| 个学生 logits 代价高昂，许多近似方案可以避免这一点。我们使用 KL 作为蒸馏损失（Hinton 等人 2015（https://arxiv.org/html/2608.10260#bib.bib92）；Sanh 等人 2020（https://arxiv.org/html/2608.10260#bib.bib93））。由于可训练的学生是 D_KL(⋅‖⋅) 的第二个参数，而期望遍历固定的教师，从 P 中抽取词元可以得到行为良好且无偏的蒙特卡洛估计器。当从另一个提议分布 R 中抽取词元时，重要性采样以 P(v|x)/R(v|x) 对每个采样贡献进行加权，从而使期望仍能恢复原始 KL（Amini 等人 2025（https://arxiv.org/html/2608.10260#bib.bib72））。确定性 Top-k 截断则仅对教师最可能的 k 个词元评分，这是一种内存廉价但有偏的缩减（Shao 等人 2024（https://arxiv.org/html/2608.10260#bib.bib70））。附录 D（https://arxiv.org/html/2608.10260#A4）将此设置与可训练采样分布下的 KL 正则化区分开来（Tang and Munos 2025（https://arxiv.org/html/2608.10260#bib.bib108））。最接近的稀疏蒸馏基线是随机采样知识蒸馏（RS-KD）（Anshuman 等人 2025（https://arxiv.org/html/2608.10260#bib.bib115）），在第 4.2 节（https://arxiv.org/html/2608.10260#S4.SS2）中进行比较；我们的采样尾部建立在用于大型输出空间的重要性采样之上（Katharopoulos and Fleuret 2019（https://arxiv.org/html/2608.10260#bib.bib94）；Blanc and Rendle 2018（https://arxiv.org/html/2608.10260#bib.bib95））。这些动机引出了 *Subset-KL*（第 4 节（https://arxiv.org/html/2608.10260#S4））：有偏的 Top-k 截断，以及具有无偏梯度的精确头部、重要性采样尾部变体。

##### 免训练与互补读出。
Jacobian lens（Gurnee 等人 2026（https://arxiv.org/html/2608.10260#bib.bib119））通过解码钩点处平均局部输出敏感性来避免变换器训练。它询问状态在一阶扰动下局部*表示*什么，而训练好的预测透镜则优化以恢复模型最终的输出分布；其作者观察到 tuned lens 因此可以“跳过”中间表示。PatchScopes（Ghandeharioun 等人 2024（https://arxiv.org/html/2608.10260#bib.bib121））同样免训练，将隐藏状态修补到解释性提示中，使模型自身的生成作为读出。这些视角是互补的，尽管每一种都仍是逐钩点的读出。更广泛地说，透镜是观察性解码器，本身并不建立因果机制。它们与因果技术互补，如电路发现和激活修补（Elhage 等人 2021（https://arxiv.org/html/2608.10260#bib.bib38）；Wang 等人 2022（https://arxiv.org/html/2608.10260#bib.bib40）；Conmy 等人 2023（https://arxiv.org/html/2608.10260#bib.bib8）；Meng 等人 2022（https://arxiv.org/html/2608.10260#bib.bib82））；观察性注意力分析（Clark 等人 2019b（https://arxiv.org/html/2608.10260#bib.bib79）；Voita 等人 2019（https://arxiv.org/html/2608.10260#bib.bib80）；Vig and Belinkov 2019（https://arxiv.org/html/2608.10260#bib.bib81））；以及面向特征的方法，包括特征可视化、将激活分解为叠加特征的自编码器，以及针对权重而非激活的随机参数分解（Olah 等人 2017（https://arxiv.org/html/2608.10260#bib.bib77）；Cammarata 等人 2021（https://arxiv.org/html/2608.10260#bib.bib78）；Cunningham 等人 2023（https://arxiv.org/html/2608.10260#bib.bib87）；Templeton 等人 2024（https://arxiv.org/html/2608.10260#bib.bib88）；Sharkey 等人 2022（https://arxiv.org/html/2608.10260#bib.bib89）；Bushnaq 等人 2025（https://arxiv.org/html/2608.1