提出可扩展的字典学习算法,有效解耦大模型高维嵌入空间。
DB-KSVD: Scalable Alternating Optimization for Disentangling High-Dimensional Embedding Spaces
- 采用改进的双批次KSVD算法,兼顾理论基础与大规模数据可扩展性。
- 在6项SAEBench指标上达到与SAE相当性能,验证其有效性。
- 适合研究大模型可解释性、高维嵌入解耦的科研人员使用。
字典学习已成为理解大型Transformer模型机制的重要方法。解耦高维Transformer嵌入需要能处理大规模样本和高维度数据的算法。现有工作探索了稀疏自编码器(SAEs)解决该问题,但其使用简单线性编码器求解稀疏编码子问题,该问题已知为NP难。因此,有必要探究此类方法是否足以找到优质解,或更复杂的算法能否获得更好结果。本文提出双批次KSVD(DB-KSVD),一种可扩展的字典学习算法,基于经典KSVD并适配大规模场景。该算法在包含百万级样本和数千维数据的场景下仍具高效性。我们在Gemma-2-2B与Pythia-160M文本嵌入上进行解耦,并在SAEBench基准的六项指标上评估,结果与基于SAE的主流方法相当。进一步在DINOv2-S与DINOv2-B图像嵌入上也取得相似成果,验证了方法稳定性。通过以不同优化路径达到与SAE相当表现,我们的结果表明:(i) SAE确实能发现字典学习中的强解;(ii) 传统优化方法经改造后可适应大规模需求,为后续研究提供新方向。代码开源于https://github.com/romeov/ksvd.jl。
原文摘要 · Abstract (English)
Dictionary learning has recently emerged as a promising approach for mechanistic interpretability of large transformer models. Disentangling high-dimensional transformer embeddings requires algorithms that scale to high-dimensional data with large sample sizes. Recent work has explored sparse autoencoders (SAEs) for this problem. However, SAEs use a simple linear encoder to solve the sparse encoding subproblem, which is known to be NP-hard. It is therefore interesting to understand whether this approach is sufficient to find good solutions to the dictionary learning problem or if a more sophisticated algorithm could find better solutions. In this work, we propose Double-Batch KSVD (DB-KSVD), a scalable dictionary learning algorithm that adapts the classic KSVD algorithm. DB-KSVD is informed by the rich theoretical foundations of KSVD but scales to datasets with millions of samples and thousands of dimensions. We demonstrate the efficacy of DB-KSVD by disentangling text embeddings of the Gemma-2-2B and Pythia-160M models and evaluating on six metrics from the SAEBench benchmark, where we achieve competitive results when compared to established approaches based on SAEs. We further show similar results when disentangling image embeddings obtained from the DINOv2-S and DINOv2-B models, solidifying our findings. By matching SAE performance with an entirely different optimization approach, our results suggest that (i) SAEs do find strong solutions to the dictionary learning problem and (ii) traditional optimization approaches can be scaled to the required problem sizes, offering a promising avenue for further research. We make an implementation of DB-KSVD available at https://github.com/romeov/ksvd.jl.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。