使用奇偶瓶颈层扩展可解释的Transformer

# 通过奇偶瓶颈层实现可扩展的可解释 Transformer 来源：https://arxiv.org/html/2607.20652 Andrew Mack\* Principles of Intelligence &Kraig Yuheng Tou 独立研究员 &Mark Henry 独立研究员 &Zhengxun Wu 独立研究员 &Lauren Greenspan Principles of Intelligence ###### 摘要 语言模型被认为表现出叠加现象，即在残差流中表征的特征数量远多于维度数。稀疏自编码器（SAE）旨在事后恢复此类特征，但训练*通过构造*即可解释的模型在实践中仍然不可行，因为每层的过完备瓶颈在内存和计算上都过于昂贵。为了解决这一问题，我们引入了 *ParityTransformer*，一个 GPT-2 规模的架构，其中间表示在设计上既高效又宽/稀疏。在每一层，深度奇偶瓶颈（DPB）用无参数代数字典替代了学习得到的过完备基，提供了确定性的非相干性保证，并消除了阻碍每层可解释瓶颈大规模应用的内存需求。DPB 是一种层次化结构的稀疏瓶颈，通过多级混合专家方法高效地强制稀疏性：一种硬件感知的实现，将激活稀疏训练与密集训练之间的成本差距缩小到可管理的可解释性税。实验上，ParityTransformer 在稀疏探针任务上至少与事后 SAE 表现相当，而在特征吸收、操控有效性和细粒度因果干预等指标上则更优。由于后续计算仅作用于通过稀疏瓶颈的特征，ParityTransformer 的特征在构造上就原生存在于模型的前向传播中，这回应了 SAE 探针特征是否是模型在计算中实际使用特征的问题。我们认为这是朝着训练内部表示通过设计即可解释而非事后恢复的模型迈出的一步。††footnotetext:\*通讯作者，andrew@princint\.ai。## 1 引言 在本文中，我们介绍了奇偶变换器（parity transformer），一种旨在既通过设计可解释又高效的 GPT-2 规模架构。该模型采用层次化的混合专家（MoE）方法来诱导激活稀疏性，使用“奇偶哈希函数”按需在芯片上选择性计算特征方向，而不是从高带宽内存（HBM）中加载。其名称反映了特征方向的构建方式：每个坐标是根据特征索引位的子集的*奇偶性*（奇数或偶数比特数）计算得出的，这是一种廉价的计算原语，如下文所述，它能够高效生成近乎正交的符号模式。这种设计消除了传统宽稀疏瓶颈所需的逐层字典内存存储和带宽需求（Tamkin 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib139)）。它产生的稀疏编码直接集成到模型的计算中，这与事后训练的稀疏自编码器（SAE）不同（Bricken 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib138)）。瓶颈架构的草图如图 1 (https://arxiv.org/html/2607.20652#S1.F1) 所示。参见图注 (a) 我们使用片上“哈希特征”构造将稀疏表示的激活“解码”为密集形式。参见图注 (b) 我们通过对代数结构的层次进行束搜索来编码特征。图 1：我们高效的深度奇偶瓶颈（DPB）示意图。我们的工作解决了 AI 可解释性研究中的两个障碍。第一个是叠加：网络似乎表征的特征数量远多于残差流方向数量，并诉诸于非正交地打包这些特征（Elhage 等人，2022 (https://arxiv.org/html/2607.20652#bib.bib53)；Bricken 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib138)）。主要的应对方法是通过稀疏自编码器（SAE）事后解开特征，这些编码器在稀疏性惩罚下训练以重建冻结的激活（Cunningham 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib91)；Templeton 等人，2024 (https://arxiv.org/html/2607.20652#bib.bib84)）。SAE 可扩展到数百万个特征，并产生看似可解释的方向，但它们留下了一个重要问题：恢复的特征是否与模型在计算中使用的特征相对应，还是代表了训练分布或自编码器的归纳偏差？已有研究表明，不同的 SAE 方法在探针基准测试上得分相似的特征，在因果干预下表现不同（例如，Karvonen 等人，2025b (https://arxiv.org/html/2607.20652#bib.bib90)），这支持了对更优的定义和度量来探针诸如忠实性等可解释属性的需求。这个问题是第二个障碍的下游问题：可解释性文献中“特征”一词至少有三种用法。 (i) *表征特征* 是在激活上训练的 SAE 返回的方向；(ii) *计算特征* 是下游计算因果依赖的活动组件，取决于任务和输入；(iii) *概念特征* 是数据本身的真实属性（“是名词”、“与科学家相关”）。事后可解释性工具通过旨在重建 (ii) 来使用 (i) 作为 (iii) 的代理。隐含的主张是 (i)∼\\sim(iii) 意味着 (i)∼\\sim(ii)，但当前的评估技术要么探针前者（例如稀疏探针），要么探针后者（例如因果消融）。相反，我们的构造将 (ii) 作为原语，通过设计一个具有合意计算属性的基。所得特征与 (i) 相比如何，以及它们是否也与 (iii) 对齐，是我们在第 3 节 (https://arxiv.org/html/2607.20652#S3) 中处理的实证问题。 ##### 贡献。 - • 我们引入了深度奇偶瓶颈（DPB），它通过一个固定的、层次化结构的稀疏字典路由特征，该字典使用代数伪随机构造在芯片上计算“哈希”特征方向，而不是从内存加载。 - • 我们训练了 200M 和 1.3B 参数的 ParityTransformer，并量化了这种架构选择带来的可解释性税。我们估计，在 1.3B 参数规模上，将 ParityTransformer 训练到类似能力水平比密集变压器昂贵 8−14×8\-14\\times 倍。这使得 ParityTransformer 成为全权重稀疏模型（据估计，训练成本是密集模型的 100−1,000×100\-1,000\\times 倍，见 Gao 等人，2025 (https://arxiv.org/html/2607.20652#bib.bib140)）与事后可解释性方法之间的实用折中方案。 - • 我们通过实验表明，ParityTransformer 在多个可解释性指标上（包括操控、细粒度因果干预和特征吸收）达到或优于配有事后 SAE 的密集变压器。 本文的其余部分组织如下。第 2 节 (https://arxiv.org/html/2607.20652#S2) 详细介绍了奇偶瓶颈层及其在 ParityTransformer 中的实现。我们在第 3 节 (https://arxiv.org/html/2607.20652#S3) 中讨论实验设计选择并展示实证结果。 ### 1\.1 相关工作 ##### 稀疏性、叠加与电路。 机制可解释性领域很大程度上建立在*电路假说*（Elhage 等人，2022 (https://arxiv.org/html/2607.20652#bib.bib53)；Bricken 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib138) 和*叠加假说*之上，后者将神经元方向视为高维稀疏激活特征空间的低维压缩。SOTA 模型该空间的维度很大，据报道 Claude Sonnet 3 中多达 3400 万个可解释特征（Templeton 等人，2024 (https://arxiv.org/html/2607.20652#bib.bib84)）。除了这种表征层面的叠加概念外，hänni2024mathematicalmodelscomputationsuperpositions 的理论模型显示了叠加如何在模型计算中主动有用。实际上，对叠加中计算的分析得出的匹配上界和下界表明，最优训练的密集 MLP 可以在高达 d2d^{2} 个特征上进行计算，最多对数因子（Adler 和 Shavit，2026 (https://arxiv.org/html/2607.20652#bib.bib148)）。另一条独立的工作线关注的是权重稀疏而非激活稀疏的模型。稀疏训练方法包括 RigL 使用的布尔掩码（Evci 等人，2021 (https://arxiv.org/html/2607.20652#bib.bib144)）以及 Zhang 等人（2026 (https://arxiv.org/html/2607.20652#bib.bib142)）在训练过程中使用 Cannistraci-Hebbian 学习迭代剪枝权重的方法。通过在权重稀疏 Transformer 中识别电路，Gao 等人（2025 (https://arxiv.org/html/2607.20652#bib.bib140)）为稀疏性有利于可解释性提供了更多证据。Tamkin 等人（2023 (https://arxiv.org/html/2607.20652#bib.bib139)）通过在每一层插入扁平的 TopK 查找到一个过完备“码本”来强制激活稀疏性。他们报告了小模型（410M）上的可解释特征和操控，但没有解决将宽码本层插入原本狭窄的 Transformer 所导致的低效问题。尽管与此架构类似，我们的工作针对的是这种效率差距，对中间稀疏编码施加了层次化约束，并提供了与密集 Transformer 和事后 SAE 的比较点。 ##### 细粒度混合专家（MoE） 为降低计算成本，我们采用了 MoE 策略。细粒度混合专家的相关工作展示了专家粒度的缩放规律，并发现 MoE 相对于密集模型的 FLOP 效率优势随着规模扩大而持续扩大（Krajewski 等人，2024 (https://arxiv.org/html/2607.20652#bib.bib161)）。这支持了以下观点：对于任何给定的词元，只有总前馈参数的一小部分需要激活。PEER（He，2024 (https://arxiv.org/html/2607.20652#bib.bib162)）应用了一种高效的路由机制来利用巨大的专家池。两项工作均未涉及可解释性或硬件感知的实现。 ## 2 深度奇偶瓶颈描述 受现有工作（第 1.1 节 (https://arxiv.org/html/2607.20652#S1.SS1)）的指导，我们优先考虑以下设计原则。 ##### 固定基中的稀疏性。 文献中强制稀疏性带来的可解释性收益众多（Tamkin 等人，2023 (https://arxiv.org/html/2607.20652#bib.bib139)；Gao 等人，2025 (https://arxiv.org/html/2607.20652#bib.bib140)；Kosowski 等人，2025 (https://arxiv.org/html/2607.20652#bib.bib141)），但在训练期间学习稀疏的过完备基很快就会变得笨重。为了解决这一挑战，我们依赖以下假设：*任何*近似的正交过完备基，当用作稀疏瓶颈时，都会诱导出可解释的特征。因此，我们可以*设计*这个基，使其便于在 GPU 上计算。 ##### 具有“哈希特征方向”的多级 MoE。 一种自然降低对 m≫dm\\gg d 个候选进行评分成本的方法是条件计算：将特征组织成层次结构，仅对活动父节点的子节点进行评分，如 MoE 架构中所示（Shazeer 等人，2017 (https://arxiv.org/html/2607.20652#bib.bib232)），这种方法已成功应用于 SAE（Mudide 等人，2025 (https://arxiv.org/html/2607.20652#bib.bib233)）。然而，标准 MoE 并不适合我们的场景。每个专家都携带位于 HBM 中的学习权重矩阵。将这些权重加载到计算单元是一个重大瓶颈，因此只有当足够的批次元素被路由到该专家以将这一单次昂贵的加载分摊到大量 FLOP 上时，专家才具有成本效益。对于可解释性，我们预期存在长尾的稀有特征，每个批次可能仅触发一次——这正是分摊失败、每个专家的内存访问主导挂钟时间的场景。这使得在标准范式下深度层次结构不切实际：增加层次会创建指数级更稀有的叶专家，每个都必须从 HBM 加载。我们的关键技术思想是通过用*计算*得到的特征方向替换学习到的特征方向来消除这一内存瓶颈。我们不是存储一个 m×dm\\times d 的字典矩阵并在需要某个特征时加载一行，而是通过一个轻量级哈希函数从其特征整数索引确定性地推导出每个特征的方向。该哈希通过一个小种子矩阵将索引映射到 dd 维符号模式（元素为 ±1/d\\pm 1/\\sqrt{d}），该种子矩阵完全适合 GPU 寄存器——内存层次中最小的但最快的层级。不同的索引在构造上产生近乎正交的方向，从而最小化特征之间的干扰。因为实现一个特征方向只需要少数寄存器级别的按位操作，每个特征的成本主要由对其进行评分的矩阵乘法决定。无论方向是如何获得的，这都是所需的相同算术运算。使得稀有专家在标准 MoE 中变得不可行的内存负载被完全消除，这使得具有数百万个潜在特征的多级层次结构变得实用。我们现在描述深度奇偶瓶颈（DPB）的架构。类似于学习的 SAE，它有两个组件：一个*编码器*，将密集的残差流向量映射到 m≫dm\\gg d 个特征上的稀疏编码；以及一个*解码器*，将稀疏编码映射回 Rd\\mathbb{R}^{d}。我们首先描述解码器，因为它定义了编码器必须从中选择的特征基。 ### 2\.1 奇偶解码器 我们的解码器将一组活动特征 SS（具有系数 \{ai\}i∈S\\{a_{i}\\}_{i\\in S}）转换为稠密向量：x^=∑i∈Saiφi\\hat{x}=\\sum_{i\\in S}a_{i}\\,\\phi_{i}。挑战在于如何在不从 HBM 存储或加载 m×dm\\times d 字典矩阵的情况下获取每个 φi∈Rd\\phi_{i}\\in\\mathbb{R}^{d}。我们的基本原语是一族哈希符号向量，每个整数索引对应一个，在芯片上计算： φ~i=1dR(Mbits(i))，\\tilde{\\phi}_{i}=\\frac{1}{\\sqrt{d}}\\mathcal{R}(M\\textrm{bits}(i))，(1) 其中 MM 和 bits(i)\\textrm{bits}(i) 的元素属于有限域 F2\\mathbb{F}_{2}，而 R\\mathcal{R} 将 F2\\mathbb{F}_{2} 映射到实数（将 0→10\\rightarrow 1 和 1→−11\\rightarrow -1）。矩阵 M∈F2d×rM\\in\\mathbb{F}_{2}^{d\\times r} 的行由稠密坐标 k∈\{1,...,d\}k\\in\\{1,\\ldots,d\\} 索引，列由位位置 p∈\{1,...,r\}p\\in\\{1,\\ldots,r\\} 索引，其中 r=log2⁡(m)r=\\log_{2}(m)；MM 和 R\\mathcal{R} 共同将一个 rr 位索引转换为 dd 个稠密坐标上的符号模式。φ~i\\tilde{\\phi}_{i} 的每个坐标因此由 ii 的二进制表示中一个选定子集位的奇偶性（奇数或偶数计数）决定，并映射到 ±1/d\\pm 1/\\sqrt{d}。这就是“奇偶瓶颈”名称的由来。字典由该族构建，但有一个修改：前 dd 个解码器方向被覆盖为标准基， φi=ei (i≤d)，φi=φ~i (i>d)，\\phi_{i}=e_{i}\\;\;(i\\leq d)，\\qquad\\phi_{i}=\\tilde{\\phi}_{i}\\;\;(i>d)，(2) 这使得下一节描述的层次编码器的第一阶段基本上免费，因为对基对齐特征进行评分是坐标读取而非点积。该覆盖不能