SoftSAE: Dynamic Top-K Selection for Adaptive Sparse Autoencoders

📄 arXiv: 2605.06610v2 📥 PDF

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

分类: cs.LG, cs.CV

发布日期: 2026-05-07 (更新: 2026-05-08)

🔗 代码/项目: GITHUB

💡 一句话要点

提出SoftSAE：一种基于可微Soft Top-K机制的自适应稀疏自编码器，实现输入依赖的动态特征选择。

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

关键词: 机械可解释性 稀疏自编码器 动态稀疏度 特征解耦 可微算子 神经网络分析

📋 核心要点

  1. 现有Top-K SAE强制所有输入使用固定数量的激活特征，无法适应现实数据中随样本变化的局部内在维度，导致简单输入产生噪声或复杂输入信息丢失。
  2. 提出SoftSAE架构，引入可微的Soft Top-K算子，使模型能够根据输入内容的复杂度动态调整激活特征的数量，实现输入依赖的稀疏度控制。
  3. 实验证明SoftSAE能有效提取语义特征，并根据输入概念的复杂程度自动选择最优特征数量，提升了模型对神经网络内部表示的解释精度。

📝 摘要（中文）

稀疏自编码器（SAE）已成为机械可解释性领域的重要工具，用于将大语言模型（LLM）和视觉Transformer（ViT）中的多义激活分解为单义特征。然而，传统的Top-K SAE架构依赖于固定的稀疏度水平，即强制所有输入具有相同的激活特征数量（K），这忽略了现实数据中局部内在维度随样本变化的复杂性。固定K值可能导致简单输入引入噪声，或复杂输入丢失关键结构。为此，本文提出了SoftSAE，引入了一种动态Top-K选择机制。该方法利用可微的Soft Top-K算子，使模型能够根据输入复杂度自适应地学习稀疏度水平k。实验结果表明，SoftSAE不仅能提取有意义的特征，还能为不同概念自动选择最优的特征数量，从而使表示更贴合数据结构。

🔬 方法详解

问题定义：现有稀疏自编码器（SAE）在机械可解释性任务中普遍采用固定的Top-K稀疏约束。这种设计假设所有输入样本的特征复杂度一致，但在实际神经网络激活中，不同输入对应的流形结构差异巨大，固定K值无法平衡稀疏性与表示完整性。

核心思路：引入动态稀疏度概念，将特征选择过程建模为输入依赖的优化问题。通过引入可微的Soft Top-K算子，使模型能够端到端地学习每个样本所需的最佳激活特征数量，从而实现表示长度与信息量的自适应对齐。

技术框架：SoftSAE在标准SAE架构基础上，将传统的硬阈值选择替换为可微的Soft Top-K模块。该模块接收编码器输出的激活值，通过参数化的方式计算出每个样本对应的动态k值，并输出稀疏化的特征表示，随后通过解码器重构输入。

关键创新：核心创新在于将不可导的Top-K操作转化为可微的Soft Top-K算子，打破了稀疏度必须预设为常数的限制，实现了特征选择机制的动态化与自适应性。

关键设计：采用了基于可微排序或平滑近似的Soft Top-K算子，结合输入依赖的门控机制来预测k值。损失函数设计上，除了重构损失外，还包含对动态k值的约束，以防止模型在极端情况下过度激活或完全关闭特征。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，SoftSAE在保持高重构质量的同时，显著提升了特征提取的语义相关性。相比于固定K值的基线模型，SoftSAE在处理不同复杂度输入时展现出更强的鲁棒性，能够自动识别出不同概念所需的特征数量，有效减少了简单样本中的冗余激活，并捕捉到了复杂样本中的细粒度结构。

🎯 应用场景

该研究主要应用于大语言模型（LLM）与视觉Transformer（ViT）的机械可解释性分析。通过更精准地提取单义特征，SoftSAE可辅助研究人员识别模型内部的逻辑推理路径、偏见来源及知识存储方式，在模型审计、安全性评估及可解释AI（XAI）领域具有重要应用价值。

📄 摘要（原文）

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://github.com/St0pien/SoftSAE.