SoftSAE: Dynamic Top-K Selection for Adaptive Sparse Autoencoders

📄 arXiv: 2605.06610v1 📥 PDF

作者: Jakub Stępień, Marcin Mazur, Jacek Tabor, Przemysław Spurek

分类: cs.LG, cs.CV

发布日期: 2026-05-07

💡 一句话要点

提出SoftSAE，通过动态Top-K选择机制实现自适应稀疏自编码器，提升表征能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 稀疏自编码器 机制可解释性 动态Top-K选择 自适应稀疏性 特征选择

📋 核心要点

  1. 传统Top-K SAEs采用固定稀疏度，无法适应真实世界数据复杂性变化，导致次优表示。
  2. SoftSAE通过可微Soft Top-K算子学习输入相关的稀疏度，动态调整激活特征数量。
  3. 实验表明SoftSAE能找到有意义的特征，并为每个概念选择合适数量的特征，提升表征能力。

📝 摘要（中文）

稀疏自编码器(SAEs)已成为机制可解释性的重要工具，有助于分析大型语言模型(LLMs)和视觉Transformer(ViTs)中的内部表示。通过将多义激活分解为单义特征的稀疏集合，SAEs旨在将神经网络计算转化为人类可理解的概念。然而，诸如TopK SAEs之类的常见架构依赖于固定的稀疏度。它们对所有输入强制执行相同数量的活动特征(K)，忽略了真实世界数据变化着的复杂性。自然数据通常位于具有不同局部固有维数的流形上，这意味着相关因素的数量可能因样本而异。因此，固定的稀疏度不是最优的。简单的输入可能只需要几个特征，而更复杂的输入则需要更具表现力的表示。因此，使用恒定的K可能会在简单情况下引入噪声，或者在更复杂的情况下遗漏重要的结构。为了解决这个问题，我们提出了一种具有动态Top-K选择机制的稀疏自编码器SoftSAE。我们的方法使用可微的Soft Top-K算子来学习输入相关的稀疏度k。这允许模型根据每个输入的复杂性来调整活动特征的数量。因此，表示更好地匹配数据的结构，并且解释长度反映了输入中的信息量。实验结果证实，SoftSAE不仅找到了有意义的特征，而且还为每个概念选择了正确数量的特征。

🔬 方法详解

问题定义：现有Top-K稀疏自编码器(SAEs)依赖于固定的稀疏度，即对所有输入强制执行相同数量的活动特征(K)。然而，真实世界的数据具有不同的复杂性，简单的输入可能只需要少数特征，而复杂的输入则需要更多的特征。固定K值会导致简单输入引入噪声，复杂输入丢失重要结构，无法有效捕捉数据内在的复杂性。

核心思路：SoftSAE的核心思路是引入一个动态的Top-K选择机制，使得稀疏自编码器能够根据输入数据的复杂程度自适应地调整激活特征的数量。通过学习一个输入相关的稀疏度k，模型可以更好地匹配数据的结构，并提供更准确的解释。这种自适应性使得模型能够更好地处理各种复杂程度的数据，从而提高表征能力。

技术框架：SoftSAE的整体框架与传统的稀疏自编码器类似，包括编码器和解码器两个主要部分。不同之处在于，SoftSAE在编码器和解码器之间引入了一个动态Top-K选择模块。该模块使用一个可微的Soft Top-K算子，根据输入的特征激活值，动态地选择前k个最重要的特征。选择后的特征被传递到解码器进行重构。整个框架通过最小化重构误差和稀疏性约束进行训练。

关键创新：SoftSAE最关键的创新点在于其动态Top-K选择机制。传统的Top-K方法是不可微的，无法直接进行端到端训练。SoftSAE通过使用可微的Soft Top-K算子，使得模型可以学习一个输入相关的稀疏度k，从而实现自适应的特征选择。与固定K值的Top-K SAEs相比，SoftSAE能够更好地适应不同复杂程度的数据，并提供更准确的解释。

关键设计：SoftSAE的关键设计在于Soft Top-K算子的实现。该算子需要满足可微性，以便进行端到端训练。一种常见的实现方式是使用Gumbel-Softmax技巧，将离散的Top-K选择转化为连续的概率分布。此外，损失函数通常包括重构误差项和稀疏性约束项。稀疏性约束项可以采用L1正则化或KL散度等形式，鼓励模型学习稀疏的特征表示。具体参数设置需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SoftSAE不仅能够找到有意义的特征，而且能够为每个概念选择合适数量的特征。与固定K值的Top-K SAEs相比，SoftSAE在多个数据集上取得了更好的性能，验证了其自适应特征选择机制的有效性。具体的性能提升幅度取决于数据集和任务，但总体趋势是SoftSAE能够提供更准确和更具解释性的特征表示。

🎯 应用场景

SoftSAE在机制可解释性领域具有广泛的应用前景，可用于分析大型语言模型和视觉Transformer的内部表示。通过将复杂的神经网络计算分解为人类可理解的概念，SoftSAE有助于理解模型的决策过程，并提高模型的可信度。此外，SoftSAE还可以应用于特征选择、异常检测等领域，具有重要的实际价值和未来影响。

📄 摘要（原文）

Sparse Autoencoders (SAEs) have become an important tool in mechanistic interpretability, helping to analyze internal representations in both Large Language Models (LLMs) and Vision Transformers (ViTs). By decomposing polysemantic activations into sparse sets of monosemantic features, SAEs aim to translate neural network computations into human-understandable concepts. However, common architectures such as TopK SAEs rely on a fixed sparsity level. They enforce the same number of active features (K) across all inputs, ignoring the varying complexity of real-world data. Natural data often lies on manifolds with varying local intrinsic dimensionality, meaning the number of relevant factors can change significantly across samples. This suggests that a fixed sparsity level is not optimal. Simple inputs may require only a few features, while more complex ones need more expressive representations. Using a constant K can therefore introduce noise in simple cases or miss important structure in more complex ones. To address this issue, we propose SoftSAE, a sparse autoencoder with a Dynamic Top-K selection mechanism. Our method uses a differentiable Soft Top-K operator to learn an input-dependent sparsity level k. This allows the model to adjust the number of active features based on the complexity of each input. As a result, the representation better matches the structure of the data, and the explanation length reflects the amount of information in the input. Experimental results confirm that SoftSAE not only finds meaningful features, but also selects the right number of features for each concept. The source code is available at: https://anonymous.4open.science/r/SoftSAE-8F71/.