[论文解读] BatchTopK Sparse Autoencoders

Bart Bussmann, Patrick Leask+1 more|arXiv (Cornell University)|Dec 9, 2024

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

本文提出批量TopK稀疏自编码器（BatchTopK Sparse Autoencoders），通过在整批样本中选择Top-k激活而非每个样本独立选择，放宽了标准TopK SAEs中每个样本的Top-k激活约束，从而实现每个样本的可变潜在激活。该方法在不牺牲平均稀疏性的情况下提升了重建性能，并可直接指定稀疏性水平，优于TopK SAEs，且在重建性能上与JumpReLU SAEs相当，同时避免了昂贵的超参数调优。

ABSTRACT

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting language model activations by decomposing them into sparse, interpretable features. A popular approach is the TopK SAE, that uses a fixed number of the most active latents per sample to reconstruct the model activations. We introduce BatchTopK SAEs, a training method that improves upon TopK SAEs by relaxing the top-k constraint to the batch-level, allowing for a variable number of latents to be active per sample. As a result, BatchTopK adaptively allocates more or fewer latents depending on the sample, improving reconstruction without sacrificing average sparsity. We show that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, and achieve comparable performance to state-of-the-art JumpReLU SAEs. However, an advantage of BatchTopK is that the average number of latents can be directly specified, rather than approximately tuned through a costly hyperparameter sweep. We provide code for training and evaluating BatchTopK SAEs at https://github.com/bartbussmann/BatchTopK

研究动机与目标

• 为解决TopK SAEs中固定每个样本稀疏性的局限性，该局限性假设所有输入均需相同数量的活跃潜在变量，而不论其复杂度如何。
• 通过根据输入复杂度允许每个样本的激活数量可变，在不增加平均稀疏性的情况下提升重建性能。
• 实现对平均稀疏性水平的直接指定，消除其他SAC变体中常见的昂贵超参数搜索需求。
• 在保持可解释性的同时，通过利用基于批次的激活选择来增强重建性能。
• 在GPT-2 Small和Gemma 2 2B上评估该方法，使用标准指标与TopK和JumpReLU SAEs进行性能比较。

提出的方法

• BatchTopK用基于批次的Top-k选择替代了每个样本的Top-k选择，即从整个批次的所有样本中选择Top-k激活用于重建。
• 该方法使用Top-k操作的可微分松弛，以在训练期间允许梯度流动，从而实现潜在字典的端到端优化。
• 推理时，模型将全局阈值θ估计为多个批次中最小正激活值的平均值，用基于阈值的激活函数替代依赖批次的选择机制。
• 阈值θ被计算为多个批次中最小正激活值的期望值，以确保推理行为的一致性。
• 该架构采用标准的SAE编码器-解码器结构，使用ReLU激活和L2重建损失，并增加稀疏性惩罚项和辅助损失以防止潜在神经元死亡。
• BatchTopK中的辅助损失与TopK SAEs完全相同，通过小规模重建死亡潜在神经元来维持其有效性。

Figure 1: On GPT-2 Small activations, BatchTopK largely achieves better NMSE and CE than standard TopK across different dictionary sizes, for a fixed number of active latents of 32 (Left). JumpReLU SAEs are omitted from this comparison as their L0 cannot be fixed to a value. For fixed dictionary siz

实验结果

研究问题

• RQ1 将每个样本的Top-k约束放宽为基于批次的约束，是否能提升稀疏自编码器的重建性能？
• RQ2 允许每个样本的活跃潜在变量数量可变，是否能在不增加平均稀疏性的前提下带来更好的重建效果？
• RQ3 BatchTopK能否在实现与最先进JumpReLU SAEs相当的性能的同时，实现对稀疏性水平的直接指定？
• RQ4 BatchTopK与固定Top-k方法在每个样本的活跃潜在变量分布上存在何种差异？
• RQ5 BatchTopK的增强灵活性是否能带来在交叉熵退化指标上的更好下游性能？

主要发现

• 在GPT-2 Small和Gemma 2 2B上，BatchTopK SAEs在所有测试字典大小和稀疏性水平下，均显著低于TopK SAEs的归一化均方误差（NMSE）。
• 在固定活跃潜在变量数（L0=32）下，BatchTopK在两个模型上的NMSE和交叉熵（CE）退化指标上均优于TopK。
• 在GPT-2 Small上，BatchTopK在所有稀疏性水平下均优于JumpReLU SAEs的NMSE和CE退化指标，尽管JumpReLU在高激活区域具有更优的下游鲁棒性。
• 在Gemma 2 2B上，BatchTopK在重建性能（NMSE）上与JumpReLU相当或更优，但在CE退化指标上仅在最低稀疏性水平（k=16）下优于JumpReLU。
• BatchTopK中每个样本的活跃潜在变量分布表现出显著差异，部分样本仅使用一个潜在变量，而其他样本则使用超过80个，证实了该方法的自适应激活策略。
• 与JumpReLU SAEs不同，BatchTopK可直接指定平均稀疏性水平而无需超参数搜索，后者需调优多个参数才能达到目标稀疏性。

Figure 2: On Gemma 2 2B activations, BatchTopK outperforms TopK SAEs across different values of k. Although BatchTopK has a better reconstruction performance (left), it only outperforms JumpReLU in terms of downstream CE degradation in the setting where k=16 (right).