通过特征相关性更高效地训练稀疏自编码器：让人工智能模型更透明

在人工智能领域，尤其是大型语言模型的研究中，我们一直面临着一个重要挑战：如何理解这些复杂模型内部的工作机制？这就像是拥有了一台神奇的机器，它能完成各种惊人的任务，但我们却不完全明白它内部的齿轮是如何运转的。2025年5月28日，来自T-Tech和莫斯科物理技术学院的研究团队——Vadim Kurochkin、Yaroslav Aksenov、Daniil Laptev、Daniil Gavrilov和Nikita Balagansky发表了一篇题为《通过利用特征相关性高效训练稀疏自编码器》的研究论文（arXiv:2505.22255v1），为解决这一难题提供了创新方法。

稀疏自编码器（Sparse Autoencoders，简称SAE）是近年来备受关注的工具，它们能够帮助我们解释语言模型的隐藏状态，将复杂的神经激活分解成人类可理解的潜在方向。想象一下，这就像是把一束混合的光线通过棱镜分解成不同颜色的光谱，让我们能够看清每种颜色的特性。然而，使用SAE面临一个主要障碍：当我们需要处理现代大型语言模型时，它们的计算成本非常高，尤其是当我们使用大型字典（即更多的"棱镜面"）来获得更细致的分解时。

虽然在SAE的解码器部分已经有了一些高效的方法，但编码器仍然需要进行计算密集型的线性操作，这严重限制了SAE在大规模模型上的应用。这就像是我们优化了信息输出的管道，但信息输入的管道仍然是个瓶颈，大量的数据在这里堵塞，影响了整个系统的效率。

针对这一问题，研究团队提出了一个名为KronSAE的创新架构。这个架构的核心思想是利用克罗内克积分解（Kronecker product decomposition）来分解潜在表示，大幅减少内存和计算开销。此外，他们还引入了一个名为mAND的可微分激活函数，它近似模拟二进制AND操作，在他们的分解框架中提高了可解释性和性能。

一、KronSAE：如何打破编码器瓶颈？

传统的稀疏自编码器面临着一个明显的瓶颈：编码器投影。想象一下，这就像是通过一个巨大的漏斗将信息压缩——漏斗越大，需要的力气就越大。在技术术语中，这意味着对于输入维度为d、字典大小为F的SAE，每次前向传递都需要O(Fd)的计算复杂度，这在现代变换器模型上是非常昂贵的。

KronSAE通过一个巧妙的架构设计解决了这个问题。它将潜在空间分解为h个独立的组件（可以想象为不同的"头部"），每个头部k由两个薄矩阵参数化：一个是"组合基础"Pk∈R^(m×d)，另一个是"组合扩展"Qk∈R^(n×d)，其中维度m < n << d，总字典大小F = h×m×n。

这种分解方法的工作原理类似于将一个复杂的拼图分解成多个小块，每个小块都更容易处理。具体来说，系统首先计算两个预潜在表示：

pk = ReLU(Pkx) qk = ReLU(Qkx)

这些预潜在表示通过一个元素级的交互核心（mAND）在每个头部中独立组合：

zki,j := mAND(pki, qkj) := { √(pki*qkj), 如果pki > 0且qkj > 0 0, 否则

这里的mAND核心平滑地近似布尔AND门，确保只有当两个输入都为正时才产生非零激活，同时保持梯度流和激活幅度以实现稳定的重建。

最后，我们扁平化并连接所有头部的结果，得到后潜在表示f∈R^F，然后应用通常的TopK操作（即保留k个最大激活值）。

通过这种方法，每个token的编码器成本从O(Fd)下降到O(h(m+n)d)，显著减少了FLOPs和参数数量，而不需要像其他方法那样引入路由开销。更重要的是，KronSAE与现有的稀疏解码器核心兼容，因此可以与它们结合使用，实现端到端的加速。

二、实验验证：KronSAE的效果如何？

研究团队在Qwen-2.5-1.5B-Base、Pythia-1.4B-deduped和Gemma-2-2B等语言模型上进行了广泛的实验。他们使用了FINEWEB-EDU（FineWeb语料库中经过过滤的教育网页子集）收集激活值，并通过解释方差（EV）来衡量重建质量，其中1.0为最佳。

在等效计算预算下，研究团队比较了KronSAE和传统的TopK SAE的性能。实验结果令人印象深刻：

在100M token预算下，所有KronSAE变体在重建质量上超过了TopK基线，解释方差提高了高达4.3%，同时参数数量减少了约54.7%。这就像是用更少的零件构建了一个更高效的机器。

在500M token预算下，KronSAE在大多数模型大小上保持了0.8%的优势，参数减少了43.8%。较小的构成基础维度（m=2）在紧张的计算约束下表现尤为出色。

在1000M token预算下，尽管TopK SAE缩小了差距，KronSAE仍然能够匹配基线的重建质量，同时参数数量减少了约46.1%。

研究团队还进行了详细的消融实验，以了解不同设计选择对性能的影响：

首先，他们比较了mAND操作与两个更简单的交互核心：ReLU(u)·ReLU(v)和原始乘积u·v。结果显示，在1B token训练预算下，mAND变体始终实现最高的解释方差，明显优于替代方案。

其次，他们系统地改变头部数量h和每个头部的基础维度m（同时保持n = F/(mh)）。结果表明，对于大型训练预算（500M-1B token），较小的m（因此较大的n）产生更高的重建质量，因为较小的基础维度释放了容量用于更具表现力的扩展特征。在更紧张的100M token预算下，m=4的配置优于m=2或m=8，表明每个头部的表示丰富性与数据效率之间存在权衡。此外，固定m并增加h几乎线性地提高了解释方差。

最后，他们评估了KronSAE在不同稀疏性水平和层深度上的稳健性。在所有情况下，KronSAE在相同FLOPs预算下匹配或超过了TopK基线的重建质量，证明了其克罗内克分解编码器无论稀疏性水平或深度如何都能保持其优势。

三、特征吸收：KronSAE如何改善特征解释性？

在可解释性研究中，一个主要挑战是"特征吸收"，即一个学习特征成为另一个特征的严格子集（例如，"狮子"特征完全被"以L开头"特征包含），因此无法在满足更广泛概念但不满足其超集表示的实例上激活。

研究团队报告了三个吸收指标：平均吸收分数（部分被吸收的特征比例）、平均完全吸收分数（量化完全包含事件）和平均特征分裂数（单个概念特征分裂成多个激活的频率）。

实验结果表明，在所有稀疏性水平l0∈{16, 32, 64, 128, 256}上，KronSAE变体始终减少了相对于TopK SAE基线的吸收分数和完全吸收分数，同时保持类似的特征分裂率。

研究团队将KronSAE改进的解耦归因于两个互补的设计选择：

1. 平滑mAND激活：通过仅在两个预潜在值都为正时才输出非零值，它引入了一个可微分AND门，防止广泛的多语义原语完全包含更具体的原语。因此，复合后潜在主要在其构成概念的交集处触发，鼓励每个预潜在专注于单一语义模式，而不是继承其"父"激活区域。

2. 头部式笛卡尔分解：将潜在空间划分为h个独立的子空间（每个子空间都有自己的m×n原语交互网格）确保专门的概念（如"大象"）被限制在单个头部中，不能完全吸收其他头部中的更一般概念（如"以E开头"）。

这些机制共同产生了更多的单语义特征，简化了下游因果干预和目标探测。值得注意的是，平均特征分裂数与TopK基线保持相当，因为笛卡尔分解本身不会固有地改变单个原语的分裂。

四、深入分析：KronSAE如何模拟特征关联？

为了评估不同稀疏自编码器架构如何恢复已知的相关模式，研究团队构建了一个受控实验，使用合成的、块结构化的协方差模型。他们生成带有不同块结构的输入向量，然后训练自编码器重建它们，检查所学习的特征相关性。

结果显示，KronSAE的解码器权重协方差Cdec = WdecW?dec比TopK SAE更忠实地再现了地面真相分组。特别是，在第三种协方差模式（其中一些块非常小）上，TopK的学习相关性几乎消失，而KronSAE仍然发现了正确的块结构。

这些观察通过RV系数和排列测试得到了量化。即使在最佳匹配TopK原子到密集AE参考后，TopK SAE也只能实现微弱的相关性对齐（RV≈0.05-0.08），p值不显著或边缘显著。相比之下，KronSAE配置实现了0.11到0.35之间的RV值（所有p<0.001），在相关性恢复方面提高了3-6倍。

此外，研究团队分析了在语言数据上训练的SAE中的特征相关性，发现KronSAE中一个头部内的特征相关性确实显著高于不同头部之间的特征，这表明他们的设计成功地在SAE潜在空间中施加了所需的相关结构。

五、学习特征的详细分析：KronSAE如何提高可解释性？

研究团队对KronSAE和TopK架构进行了深入的可解释性分析，使用自动化管道解释激活模式，并通过检测分数和模糊分数评估所获得的解释。

KronSAE学习的特征更加具体，体现在计算指标的较低值和较高的可解释性分数上。因为后潜在比相应的预潜在更具可解释性，研究团队推测了编码和检索所需语义的隐藏机制。

通过检查激活示例和潜在解释，他们观察到预潜在可能携带多个不同的、可识别的激活模式。例如，头部23中的组合基础元素3展示了与比较描述符、地理区域和精神概念相关的子语义。预潜在的多语义性是减少"工作"编码器潜在数量的预期结果，因为它们分解了完整的字典大小并减少了编码器容量。

研究团队假设特定语义的编码可能通过幅度来完成，这通过检查激活示例得到了验证。例如，上述预潜在中，"比较"部分编码在前75%分位数，而"精神"部分主要在前25%分位数中找到，"地理"部分主要编码在四分位范围内。

头部通常包含语义相关的预潜在组，例如，头部136包含三个基础元素和一个与数字和序数相关的扩展，两个与地理和空间事物相关的扩展元素，一个与问题相关的基础和一个与增长相关的扩展。有趣的是，该头部的大多数后潜在具有比其父预潜在更高的可解释性分数，这是不寻常的。

检索主要通过类似逻辑AND电路的机制发生，其中一些预潜在作为多个语义的载体，相应的预潜在（基础或扩展）作为指定器。例如，在基础包含三个可检测的子语义的情况下，每个扩展然后检索特定的语义。

其他类型的交互也可能发生，例如完全新的语义的出现，如头部23中基础3和扩展1之间的组合，其中出现了医学术语，不能简单地解释为两个预潜在语义的交集。

经常出现的一种情况是后潜在只继承一个父语义，或者另一个父语义的影响无法检测到，这通常发生在父语义具有非常广泛的解释和低分数时。然而，需要更复杂的技术来正确识别交互的细粒度结构。

在后潜在的几何方面，每个后潜在向量在残差流中都有一个向量表示，即Wdec中的相应列，这是我们在训练SAE时寻找的过完备基向量的近似。研究团队没有观察到TopK和KronSAE之间的特征几何有任何显著差异，除了KronSAE的架构设计导致聚类，使得由相同头部、基础或扩展元素产生的后潜在被分组在一个紧密的簇中，几何结构取决于我们选择的超参数h、m、n，这是预期的，可能对进一步应用如引导很有用。

六、结论与未来展望

KronSAE代表了稀疏自编码器设计的重要进步，通过头部式克罗内克分解和mAND门控直接解决了长期存在的编码器效率瓶颈。相比标准的TopK SAE，它显著降低了参数数量，同时提高了重建保真度，并通过利用特征相关性产生了更具解释性的特征。

研究团队的分析将这些收益归因于组合潜在结构和逻辑AND风格交互的互补效应，为稀疏性和分解如何在表示学习中协同作用提供了新的视角。

尽管有这些优点，KronSAE也有一些限制。它的收益取决于对(m, n, h)和mAND激活的仔细调整——配置不当的设置可能会抵消效率和质量的改进。该评估仅限于中型变换器模型和单一网络语料库，因此它在更大的模型、其他领域或语言上的适用性仍有待测试。

研究团队确定了三个扩展这项工作的方向：

1. 转码：将转码器视为信息的隐式路由器，研究替代逻辑门控函数（例如XOR或复合门）以提高可解释性和电路分析。

2. 交叉编码：将KronSAE推广到交叉编码器设置，通过逻辑操作揭示可解释的、跨层次的组合性。

3. 动态组合：探索对注意力头数量及其维度进行可学习调整，实现在不同尺度上对相关特征组进行细粒度分解。

总之，KronSAE为大型语言模型的解释提供了一个强大而高效的工具，有望推动我们对这些复杂系统内部工作原理的理解。通过减少计算负担并提高特征质量，它使研究人员能够更深入地探索模型内部，最终可能导致更透明、更可控的AI系统。