训练方法
论文在多种LLM上都尝试了训练TopK自动编码器,包括GPT-2 small、GPT-4以及一系列有GPT-4类似架构的不同大小的模型。 TopK,或者k-稀疏方法是只保留k个最大的隐变量,将其余变量置零。TopK编码器定义如下: 所有实验中都使用64 tokens长度的上下文,选取网络末尾的包含较多特征的一层(但不涉及next token预测任务)的残差流进行训练。 GPT-2 small中使用了第8层,也就是整体架构的3/4部分;对于GPT-4类架构,这个比例则是5/6。 训练后,分别使用L0和均方差(MSE)两个指标分别评估稀疏性和重建误差,并与基线模型ReLU自动编码器进行对比。 在规定隐变量数量为32768时,相比ReLU和ProLU,TopK方法在稀疏性和重建性能上实现了更好的平衡,与门控机制相近。 在固定稀疏性时,随着隐变量数目的增加,TopK方法的MSE下降速度也是所有方法中最高的。 变量数量逐渐增长时,也会有越来越多的隐变量在训练过程中的某个时刻完全停止激活,成为「死亡隐变量」。 消融实验中发现,如果不采取任何缓解措施,死亡隐变量的占比可以高达90%,这会导致MSE大幅恶化,且浪费训练算力。 为了防止隐变量「死亡」,论文提出了2个训练技巧—— - 将编码器初始化为解码器的转置 - 添加辅助损失AuxK,即使用k个「死亡隐变量」计算的重构误差(论文戏称为「幽灵梯度」) 使用这两个方法后,即使是在GPT-4上训练出的有1600万个隐变量的SAE中,也仅有7%的「死亡隐变量」。
可扩展性
由于论文的主要目的就是提升SAE的可扩展性,因此作者通过改变各种超参数设置,着重探讨了模型这方面的表现。 首先,使用与训练LLM相似的策略,不考虑收敛,在可用计算资源的条件下训练不同规模的SAE,直至达到MSE最优。 下面左图中可以看到,MSE的下降趋势基本符合计算量的幂律分布L(C),尽管其中最小的模型有些偏离。 右图则表示,选取不同的k值时,随着隐变量数量N的增大,MSE的下降趋势基本类似,遵循N和k的联合scaling法则。这体现出模型较好的优化性能和可靠性。 随着k值增加,MSE的下降曲线逐渐陡峭,意味着不可约损失的减少。 但L(N,k)的这种趋势仅在k值较小时成立,若k值过大、接近模型维数,SAE的重建就失去意义了。 下面这张图展示的结果也比较符合直觉——随着语言模型逐渐增大,SAE也需要更多的隐变量才能实现相同的MSE。
评估结果
虽然使用L0和MSE两个指标已经可以很好地评估稀疏性和重建表现,但这并不是模型的最终目标。 我们需要评估SAE提取出特征的质量和可解释性,因此论文提出了以下几种指标: 1. 下游损失:如果将原语言模型残差流中的隐变量替换为自动编码器重建的变量,语言模型的损失表现如何 2. 探测损失:SAE能否恢复我们预期的特征 3. 可解释性:对于SAE提取的隐变量在LM中的激活,是否存在充要的简单解释 4. 消融稀疏性:消除单个隐变量对LM下游得到的logits是否有稀疏影响? 通过这些指标的评估,论文得出的总体结论是,隐变量总数增加会使得SAE表现更好,然而激活的隐变量数量L0产生的影响会比较负责。
下游损失
论文考虑使用KL散度和交叉熵损失衡量语言模型的下游表现,TopK方法在稀疏性和交叉熵两方面可以取得更好的均衡。 对比其他方法,在达到相同的MSE值时,TopK方法带来的交叉熵最低,意味着下游表现最好。探测损失
为评估探测损失,团队训练了61个1维的logistic探测器,用于检测相应特征是否存在并计算最佳交叉熵(在所有隐变量中)作为探测损失。 结果显示,k值越高,探测损失越低,且TopK模型的性能在各个k值上的表现都优于ReLU。可解释性
可解释性的量化比较困难。为此,论文借鉴了2023年一篇论文提出的可解释性自动评分方法,打出的分数与召回率相关但不成比例。 由于在这种方法中使用GPT-4计算准确率和召回率过于昂贵,因此评估时作者找了另一种更便宜的模型N2G。 下面这个例子可以更直观地展现N2G的打分机制。 相比ReLU,TopK模型的F1分数更高,召回率总体有1.5×的提升,但准确率稍逊一筹,约为ReLU方法的0.9×。消融稀疏性
为了评估消融稀疏性,团队独创了一个指标(L1/L2)^2计算向量的稀疏度,来表征每删除一个隐变量时,有多少token受到了稀疏影响。 图6b可以看到,对于有不同隐变量数量的模型,随着L0逐渐提升(模型逐渐稀疏),消融稀疏性也逐渐提升(受到影响的token数量更少)。理解TopK激活函数
TopK防止激活收缩
L1正则化的一个主要缺点是,它会使所有激活值趋向于零。 而论文提出的TopK激活函数解决了这个问题,因为它不需要L1。 团队使用投影梯度下降法来求解带有正约束的最优激活值。这种优化过程在ReLU模型中通常会增加激活值,但在TopK模型中不会(见图9a),这表明TopK不受激活收缩的影响。TopK模型的优化幅度也比ReLU模型小。 在ReLU和TopK模型中,优化过程显著改善了重建的MSE(见图9b)和下游的下一个token预测的交叉熵(见图9c)。 然而,这种优化只部分缩小了ReLU和TopK模型之间的差距。渐进恢复
在渐进编码中,即使只传输部分数据,也可以以合理的精度重建信号。 对于自动编码器来说,学习渐进编码意味着通过激活幅度对潜变量进行排序,可以逐步恢复原始向量。 为了研究这一特性,团队在训练后将自动编码器的激活函数替换为TopK(k′)激活函数,其中k′与训练时的k值不同。然后,通过将每个k′值放在L0-MSE平面上进行评估(见图10)。 结果显示,用TopK训练的模型只能在训练时使用的k值范围内提供渐进编码。MSE在稍高于k的值时有所改善,但随着k′的进一步增加,MSE显著变差。 为了缓解这个问题,我们将多个不同k值的TopK损失相加,称为Multi-TopK。比如,使用L(k) + L(4k)/8就可以在所有k′上实现渐进编码。 不过需要注意的是,使用Multi-TopK训练在k值处的表现略逊于单一TopK。 对于用TopK训练的自动编码器,在测试时,TopK和JumpReLU曲线仅在L0低于训练时的L0值时重叠,否则JumpReLU激活效果不如TopK激活。 然而,这种差异在使用Multi-TopK时消失,两条曲线几乎重叠,这意味着模型可以在每个token上使用固定或动态数量的潜变量而不会影响重建效果。对于用ReLU训练的自动编码器,两条曲线也重叠。 有趣的是,有时用低L1惩罚训练ReLU模型并在测试时使用TopK或JumpReLU,比用较高L1惩罚以获得类似稀疏度更有效。局限性和未来方向