Attribution-Guided Distillation of Matryoshka Sparse Autoencoders

📄 arXiv: 2512.24975v1 📥 PDF

作者: Cristina P. Martin-Linares, Jonathan P. Ling

分类: cs.LG

发布日期: 2025-12-31

💡 一句话要点

提出DMSAE，通过归因引导蒸馏Matryoshka稀疏自编码器，提升特征一致性和可迁移性。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 稀疏自编码器 蒸馏学习 模型解释性 特征提取 归因分析

📋 核心要点

  1. 现有稀疏自编码器学习到的特征冗余且不稳定，导致解释性差，难以跨任务迁移。
  2. DMSAE通过迭代蒸馏，提取并重用一致且重要的特征核心，提升特征的稳定性和可迁移性。
  3. 实验表明，DMSAE在SAEBench指标上有所提升，验证了该方法在不同稀疏度水平上的有效性。

📝 摘要（中文）

稀疏自编码器(SAE)旨在将模型激活解耦为单义的、人类可解释的特征。然而，实践中学习到的特征通常是冗余的，并且在不同的训练运行和稀疏度水平上有所不同，这使得解释难以转移和重用。我们提出Distilled Matryoshka Sparse Autoencoders (DMSAEs)，一种训练流程，它提炼出一个紧凑的、始终有用的特征核心，并重用它来训练新的SAE。DMSAEs运行一个迭代蒸馏循环：训练一个具有共享核心的Matryoshka SAE，使用梯度X激活来测量每个特征对最嵌套重建中下一个token损失的贡献，并且只保留解释固定比例归因的最小子集。只有核心编码器权重向量在循环中转移；核心解码器和所有非核心潜在变量每次都会重新初始化。在Gemma-2-2B第12层残差流激活上，七个循环的蒸馏（500M tokens，65k宽度）产生了一个重复选择的197个特征的蒸馏核心。使用这个蒸馏核心进行训练提高了几个SAEBench指标，并证明了一致的潜在特征集可以在不同的稀疏度水平上转移。

🔬 方法详解

问题定义：论文旨在解决稀疏自编码器(SAE)学习到的特征冗余、不稳定，难以解释和迁移的问题。现有方法在不同训练轮次和稀疏度下，学习到的特征差异较大，导致模型的可解释性和泛化能力受限。

核心思路：论文的核心思路是通过迭代蒸馏，提取一个紧凑且一致的特征核心，并将其用于后续SAE的训练。这样可以保证学习到的特征更加稳定和具有代表性，从而提高模型的可解释性和可迁移性。

技术框架：DMSAE的整体框架是一个迭代蒸馏循环。首先，训练一个Matryoshka SAE，该SAE具有一个共享的核心编码器。然后，使用梯度X激活来衡量每个特征对下一个token损失的贡献。接下来，只保留解释固定比例归因的最小特征子集，作为蒸馏后的核心。最后，将核心编码器的权重向量转移到下一个循环，并重新初始化核心解码器和所有非核心潜在变量。

关键创新：DMSAE的关键创新在于使用归因引导的蒸馏方法来提取特征核心。通过梯度X激活来衡量特征的重要性，可以有效地选择对模型性能贡献最大的特征。此外，迭代蒸馏的过程可以进一步提高特征的稳定性和一致性。

关键设计：DMSAE的关键设计包括：1) 使用Matryoshka SAE作为基础模型，允许在不同稀疏度下进行训练；2) 使用梯度X激活作为归因方法，衡量特征的重要性；3) 设置固定的归因比例，控制蒸馏的强度；4) 迭代蒸馏循环，逐步提取更稳定的特征核心。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在Gemma-2-2B模型上的实验表明，经过七个循环的蒸馏，DMSAE提取了一个包含197个特征的稳定核心。使用该核心进行训练，可以提高SAEBench的各项指标，证明了DMSAE在不同稀疏度水平上的有效性。该结果表明，一致的潜在特征集可以在不同的稀疏度水平上进行转移。

🎯 应用场景

DMSAE可以应用于自然语言处理领域，例如语言模型的解释性分析、知识发现和迁移学习。通过提取稳定的特征核心，可以更好地理解语言模型的内部机制，并将其知识迁移到其他任务中。此外，该方法还可以应用于其他领域，例如图像识别和语音识别，以提高模型的可解释性和泛化能力。

📄 摘要（原文）

Sparse autoencoders (SAEs) aim to disentangle model activations into monosemantic, human-interpretable features. In practice, learned features are often redundant and vary across training runs and sparsity levels, which makes interpretations difficult to transfer and reuse. We introduce Distilled Matryoshka Sparse Autoencoders (DMSAEs), a training pipeline that distills a compact core of consistently useful features and reuses it to train new SAEs. DMSAEs run an iterative distillation cycle: train a Matryoshka SAE with a shared core, use gradient X activation to measure each feature's contribution to next-token loss in the most nested reconstruction, and keep only the smallest subset that explains a fixed fraction of the attribution. Only the core encoder weight vectors are transferred across cycles; the core decoder and all non-core latents are reinitialized each time. On Gemma-2-2B layer 12 residual stream activations, seven cycles of distillation (500M tokens, 65k width) yielded a distilled core of 197 features that were repeatedly selected. Training using this distilled core improves several SAEBench metrics and demonstrates that consistent sets of latent features can be transferred across sparsity levels