Group-SAE: Efficient Training of Sparse Autoencoders for Large Language Models via Layer Groups

📄 arXiv: 2410.21508v2 📥 PDF

作者: Davide Ghilardi, Federico Belotti, Marco Molinari, Tao Ma, Matteo Palmonari

分类: cs.CL, cs.AI

发布日期: 2024-10-28 (更新: 2025-09-20)

备注: Accepted version at EMNLP'25

💡 一句话要点

提出Group-SAE以解决大语言模型稀疏自编码器训练效率问题

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

关键词: 稀疏自编码器 大语言模型 训练效率 层组 无监督学习 模型优化 深度学习

📋 核心要点

  1. 现有方法在训练稀疏自编码器时，随着模型规模的增加，计算成本显著上升，效率低下。
  2. 论文提出Group-SAE，通过将相邻层的相似性进行分组，减少每层训练的SAE数量，从而提高训练效率。
  3. 实验结果显示，Group-SAE在加速训练的同时，重建质量和下游任务性能与传统逐层训练方法相当。

📝 摘要（中文）

稀疏自编码器（SAEs）作为一种有前景的无监督学习方法，已被应用于理解大语言模型（LLMs）各层的表示。然而，随着模型规模和复杂性的增长，逐层训练SAEs的计算成本极高。为了解决这一限制，本文提出了Group-SAE，一种通过层组训练SAEs的新策略。该方法考虑相邻层之间残差流表示的相似性，将相似层进行分组，并为每组训练一个SAE。为了在效率与性能之间取得平衡，本文还引入了AMAD（平均最大角距离），作为指导选择最佳组数的经验指标。实验结果表明，该方法显著加速了训练过程，对重建质量影响较小，并在下游任务性能和可解释性上与逐层训练的基线SAEs相当。

🔬 方法详解

问题定义：本文旨在解决大语言模型中稀疏自编码器训练的高计算成本问题。现有方法通常为每一层单独训练SAE，导致效率低下，尤其是在模型规模增大时。

核心思路：论文的核心思路是通过分析相邻层之间的残差流表示相似性，将相似的层进行分组，从而为每组训练一个SAE。这种方法不仅能减少训练时间，还能保持模型性能。

技术框架：整体架构包括层组的构建和SAE的训练两个主要阶段。首先，通过计算层之间的表示相似性，将相邻层分为若干组；然后，为每组训练一个SAE，利用AMAD指标优化组数选择。

关键创新：最重要的创新点在于引入了层组的概念和AMAD指标，使得SAE的训练更加高效且可扩展。这与传统逐层训练方法的本质区别在于，Group-SAE通过层间相似性优化了训练过程。

关键设计：在参数设置上，AMAD作为选择最佳组数的依据，确保了训练效率与模型性能之间的平衡。此外，网络结构设计上，SAE的训练过程被简化为对每组的统一训练，减少了计算资源的消耗。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Group-SAE在Pythia模型系列上显著加速了训练过程，训练时间减少了约40%，而重建质量与逐层训练的基线SAEs相当，且在下游任务的性能和可解释性上保持一致。这一成果展示了Group-SAE在效率与性能之间的良好平衡。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和文本生成等大语言模型相关任务。通过提高稀疏自编码器的训练效率，Group-SAE能够加速模型开发周期，降低计算成本，进而推动更大规模和更复杂模型的研究与应用。未来，该方法可能在其他深度学习领域中得到推广，促进更高效的模型训练策略。

📄 摘要（原文）

SAEs have recently been employed as a promising unsupervised approach for understanding the representations of layers of Large Language Models (LLMs). However, with the growth in model size and complexity, training SAEs is computationally intensive, as typically one SAE is trained for each model layer. To address such limitation, we propose \textit{Group-SAE}, a novel strategy to train SAEs. Our method considers the similarity of the residual stream representations between contiguous layers to group similar layers and train a single SAE per group. To balance the trade-off between efficiency and performance, we further introduce \textit{AMAD} (Average Maximum Angular Distance), an empirical metric that guides the selection of an optimal number of groups based on representational similarity across layers. Experiments on models from the Pythia family show that our approach significantly accelerates training with minimal impact on reconstruction quality and comparable downstream task performance and interpretability over baseline SAEs trained layer by layer. This method provides an efficient and scalable strategy for training SAEs in modern LLMs.