Rational Sparse Autoencoder

# Rational Sparse Autoencoder
Source: https://arxiv.org/html/2606.14990

Naiyu Yin  
Department of Mathematics  
Lehigh University  
Bethlehem, PA 18015  
nay224@lehigh\.edu

&Yue Yu  
Department of Mathematics  
Lehigh University  
Bethlehem, PA 18015  
yuy214@lehigh\.edu

###### 摘要

稀疏自编码器（SAE）是机械可解释性的标准工具，但当前 SAE 系列受限于固定的编码器非线性映射，例如 ReLU、JumpReLU 和 TopK。这会将特定的稀疏性机制硬编码到模型中，并可能扭曲重建与稀疏性之间的权衡。我们引入了**理性稀疏自编码器（RSAE）**，它将固定的编码器激活函数替换为可训练的有理函数。有理激活具有足够的灵活性，可以在紧支集上均匀逼近现有 SAE 系列所使用的激活原语（对于 TopK，提供了通过分离 top-k 阈值获得的门控），同时提供更丰富的函数类别以适应观察到的预激活几何形态。我们通过一个两阶段流程实现这一想法：首先是初始化过程，复制预训练的基线 SAE 权重，插入通过松弛 Remez 交换在合成数据上获得的有理系数，并连同有理系数一起校准尺度参数；随后是在标准稀疏正则化重建目标下的微调步骤。在实验上，针对三个开源权重语言模型的残差流激活以及所有三种基线激活系列，RSAE 在经过微调步骤后**严格优于**基线，无论是在重建侧指标（MSE、ℓ₀、活跃特征占比）还是下游行为指标（交叉熵退化、损失恢复）上，且不会在稀疏探测下牺牲特征级别的可解释性。这些改进在宿主语言模型、基线激活系列以及我们测试的整个基线稀疏性范围内保持一致，而升级本身仅为每个自编码器增加少量标量参数，并在单个消费级 GPU 上运行数分钟即可完成。

## 1 引言

稀疏自编码器（SAE）已成为大型语言模型机械可解释性的核心工具，将 transformer 的内部激活分解为过完备字典中单语义特征方向的稀疏线性组合（Bricken 等人，2023；Huben 等人，2024；Gao 等人，2025；Rajamanoharan 等人，2024a，b）。尽管在训练流程和评估套件方面取得了快速进展，但本文考虑的广泛使用的已发布基线将相同的仿射编码器与三种非光滑激活原语之一配对：通过 ℓ₁ 惩罚正则化的 ReLU（Bricken 等人，2023；Bloom，2024）、TopK（Gao 等人，2025）或 JumpReLU（Rajamanoharan 等人，2024b）。这些原语各自存在有据可查的病理问题。ℓ₁ 正则化 ReLU SAE 存在活跃特征幅值收缩以及持久存在的死潜在特征问题（Taggart，2024；Rajamanoharan 等人，2024a；Gao 等人，2025）。TopK 用硬基数约束替代软惩罚，这切断了非活跃特征的梯度流，并依赖辅助复兴损失来缓解死特征。JumpReLU 插入了一个可学习的每特征阈值，但需要通过连续松弛代理来实现其指示门控的反向传播。

在本文中，我们考虑 SAE 使用的浅层编码器架构：一个仿射预激活层后接一个稀疏激活块。我们证明，可训练的有理激活可以高效表示当前 SAE 系列使用的 ReLU、JumpReLU 和提供阈值的 TopK 门控。对于不连续门控，近似在紧支集上成立，且与跳跃点之间留有裕度 δ；一个大小为 polylog(1/ε) 和逆裕度的直接有理门控就足够了，并且同一个标量门控也具有恒定宽度的深层有理实现。反之，存在 O(1) 参数的有理目标映射，对于这些映射，任何标量输出的单层 SAE 编码器（使用分段仿射 ReLU/JumpReLU/提供阈值的 TopK 门控）需要 Ω(ε^{−1/2}) 个激活坐标才能达到 ε 精度。虽然本文的理论分析聚焦于浅层 SAE 设置，但我们指出，类似的效率优势也适用于深层网络。具体来说，在深层设置中，恒定宽度的有理网络实现了 O(log log(1/ε) + log log(1/δ)) 的深度上界，而分段仿射网络遵循 Ω(log(1/ε)) 的参数下界。这种分离表明，将固定的 SAE 门控替换为可训练的有理激活可以在匹配稀疏度下改善重建保真度；深层结果作为 SAE 编码器架构之外的补充扩展包含在内。

因此，我们提出了**理性稀疏自编码器（RSAE）**，这是一种 SAE，其编码器激活函数是一个可学习的有理函数，逐元素应用于仿射预激活，并具有可学习的输入/输出尺度（Cᵢₙ, Cₒᵤₜ），将每特征预激活分布映射到一个有界区间。然后我们提出一个两步 RSAE 训练算法。在初始化过程中，我们逐字复制预训练的基线 SAE 权重，插入通过松弛 Remez 交换（Chen 等人，2018）在合成数据上获得的有理系数，并将尺度参数和系数校准到基线的预激活分布。在微调过程中，我们解冻所有参数，并最小化标准的 ℓ₁ 正则化重建目标。在实验上，有理函数足够表达，以致于可以在合成数据上以低阶逼近每个基线激活。在 SAE 层面，我们在三个不同尺度的开源权重语言模型的残差流激活上，针对所有三种基线激活系列评估 RSAE，支持我们的核心主张：RSAE 在可比稀疏度下实现了更好的保真度，并严格改进了基线的重建侧指标（MSE、ℓ₀、活跃特征占比）和下游行为指标（交叉熵退化、损失恢复），且在所有宿主语言模型和基线激活系列上保持一致。这些改进在我们测试的整个基线稀疏性范围内保持一致，且不会在稀疏探测下以牺牲特征级别的可解释性为代价。所有这一切仅通过为每个自编码器增加少量标量参数，并在单个消费级 GPU 上运行数分钟即可实现。

**贡献。** 我们引入了 **RSAE**，这是一种基于可训练激活函数的新型稀疏自编码器。我们的模型以针对 SAE 编码器量身定制的逼近理论为基础：可训练的标量有理激活可以用 polylog 大小的规模模拟浅层 SAE 编码器中使用的固定 ReLU、JumpReLU 和提供阈值的 TopK 门控，而相反的下界表明，对于某些有理目标，标量输出的单层分段仿射编码器可能需要 Ω(ε^{−1/2}) 个激活坐标。为了实现这种升级策略，我们提出了一个两步 RSAE 训练算法：初始化过程复制预训练的基线 SAE 权重，随后是在标准 ℓ₁ 正则化重建目标下解冻所有参数的微调过程。我们通过实验验证，RSAE 在可比稀疏度下实现了更好的保真度，并在重建侧指标（MSE、ℓ₀、活跃特征占比）和下游行为指标（交叉熵退化、损失恢复）上均改进了基线，且在所有宿主语言模型、基线激活系列和基线稀疏级别上保持一致，同时在稀疏探测下保持特征级别的可解释性，仅增加了可忽略的参数和运行时间开销。

## 2 预备知识与相关工作

**稀疏自编码器（SAE）** 将语言模型的内部激活 \(\bm{x}\in\mathbb{R}^{d_{\mathrm{in}}}\) 分解为过完备字典中 \(d_{\mathrm{sae}}\gg d_{\mathrm{in}}\) 个特征方向 \(\bm{z}\in\mathbb{R}^{d_{\mathrm{sae}}}\) 的稀疏线性组合。它们遵循一个由编码器和解码器函数 \((f,g)\) 组成的框架，定义如下：

编码器：\(\bm{z}=f(\bm{x}):=\phi(\bm{W}_{\text{enc}}\,(\bm{x}-\bm{b}_{\text{dec}})+\bm{b}_{\text{enc}})\)，解码器：\(\hat{\bm{x}}=g(\bm{z}):=\bm{W}_{\text{dec}}\,\bm{z}+\bm{b}_{\text{dec}}\)。 (1)

我们记 \(\bm{h}\coloneqq\bm{W}_{\text{enc}}(\bm{x}-\bm{b}_{\text{dec}})+\bm{b}_{\text{enc}}\) 为预激活，因此 \(\bm{z}=\phi(\bm{h})\)。这里，\(\bm{W}_{\text{dec}}\) 的列表示用于从稀疏编码 \(\bm{z}\) 重建 \(\bm{x}\) 的解码器字典方向，且具有单位 \(\ell_2\) 范数。编码器/解码器函数中的权重通过如下形式的损失函数进行优化：

\[
\mathcal{L}(\bm{W})=\mathbb{E}_{\bm{x}\sim\mathcal{D}}\Bigl[\,\bigl\|\bm{x}-\hat{\bm{x}}(\bm{x};\,\bm{W})\bigr\|_2^2\;+\;\lambda\,S\bigl(\bm{z}(\bm{x};\,\bm{W})\bigr)\,\Bigr],\quad \bm{W}:=\{\bm{W}_{\text{enc}},\bm{W}_{\text{dec}},\bm{b}_{\text{enc}},\bm{b}_{\text{dec}}\},
\]
(2)

其中 \(S\) 是一个对非稀疏分解进行惩罚的函数，具有可调的稀疏系数 \(\lambda\)。SAE 编码器有两个目标：稀疏性，即只需要字典中的少量元素；以及保真度，即重建的 \(\hat{\bm{x}}\) 接近原始的 \(\bm{x}\)。为了实现这两个目标之间的良好平衡，提出了三种主要的 SAE 激活函数，它们在编码器激活 \(\phi\) 和施加于 \(\bm{z}\) 的稀疏机制 \(S\) 上有所不同。

- **ReLU SAE**（Bricken 等人，2023；Bloom，2024）设置 \(\phi=\mathrm{ReLU}\)，并通过显式的 \(\ell_1\) 惩罚 \(S(\bm{z}):=\left\|\bm{z}\right\|_1\) 施加稀疏性。在原始的 ReLU SAE 中，软 \(\ell_1\) 惩罚导致活跃特征的幅值收缩，并造成重建保真度的损失（Taggart，2024；Rajamanoharan 等人，2024a；Gao 等人，2025）。
- **TopK SAE**（Gao 等人，2025）提出用硬 top-k 选择 \(\bm{z}=\mathrm{TopK}_k(\bm{h})\)（产生精确的 \(\ell_0=k\)）替代软惩罚。
- **JumpReLU SAE**（Rajamanoharan 等人，2024b）保留了 \(\ell_1\) 风格的软稀疏性，但在激活函数 \(\phi\) 中插入了一个可学习的每特征阈值 \(\theta_j>0\)，设置 \(\phi(\bm{h})=\bm{h}\odot H(\bm{h}-\bm{\theta})\)，其中 \(H\) 是 Heaviside 函数，满足若 \(z\leq 0\) 则 \(H(z)=0\)，否则 \(H(z)=1\)。

与编码器激活函数的开发正交，**Matryoshka SAE**（Bussmann 等人，2025）将解码器重新组织为嵌套前缀字典，而**无数据 SAE**（Laptev 等人，2025）无需流式激活即可直接从模型权重拟合字典。虽然相关变体如**门控 SAE**（Rajamanoharan 等人，2024a）、ProLU（Taggart，2024）、**BatchTopK SAE**（Bussmann 等人，2024）以及端到端 SAE 训练（Braun 等人，2024）修改了阈值、门控、批次级稀疏性或训练目标，我们的理论和实证比较聚焦于本文考虑的被广泛使用的已发布基线——ReLU、JumpReLU 和 TopK SAE——它们的编码器非线性是固定的函数形式，稀疏性通过惩罚系数、学习阈值或基数预算来控制，而不是通过可训练的有理激活。通过一个受逼近理论支持的可训练激活架构，我们的 RSAE 为预训练 SAE（教师模型）提供了一个即插即用的修改：在维持类似稀疏水平的同时，严格改善了模型保真度。

**评估基准与预训练基线。** **SAEBench**（Karvonen 等人，2025）提供了跨模型规模的匹配预训练 ReLU、JumpReLU 和 TopK SAE，以及统一的评估套件（涵盖重建、稀疏性、下游性能和可解释性指标）；我们使用其预训练的 ReLU、JumpReLU 和 TopK SAE 作为 Pythia-160m 和 Gemma-2-2B 的基线，并在第 5 节中复用其稀疏探测工具。对于 GPT-2 small，我们额外使用 Bloom 的 *gpt2-small-res-jb* 发布（Bloom，2024）作为 ReLU 基线，以及 OpenAI 的 v5 发布（Gao 等人，2025）作为 TopK 基线。

**有理神经网络** 建立在有理函数逼近非光滑函数的关键理论优势之上（Newman，1979；Telgarsky，2017；Beckermann and Townsend，2017；Chen 等人，2018）。在 Boullé 等人（2020）的工作中，采用有理函数作为前馈网络中 ReLU 或 tanh 的可学习替代品，用于图像分类任务。在算子学习和 PDE 代理模型中，有理逼近的谱密度加速了在光滑目标算子上的收敛（Trimmel 等人，2022）。在有理神经网络中，前馈层的标准激活函数被替换为可训练的有理函数 \( \frac{P(t)}{Q(t)} \)。一个天真的可学习分母 \( Q(t) \) 在训练过程中可能发展出发散极点，从而导致数值不稳定性。为了解决这个问题，通常采用稳定的参数化，例如将分母系数约束为正数或使用 Chebyshev 展开。在我们提出的 RSAE 中，我们选择使用一种稳定的有理函数形式，确保分母在整个感兴趣域上非零，从而避免发散极点并保证平滑的训练过程。通过这种方式，RSAE 能够以较少的参数捕获复杂的非线性行为，同时保持数值稳定性。

在后续章节中，我们将详细阐述 RSAE 的理论基础、初始化与微调算法，并通过大量实验验证其有效性。