Scaling and evaluating sparse autoencoders
缩放与评估稀疏自编码器
这是Anthropic放出Scaling Monosemanticity那篇后OpenAI马上就放出来的一篇文章,相比去年上半年尝试用语言模型解释更小的语言模型中的neurons,利用SAE搭配activation看起来是一个做模型可解释性更加solid的方法(Anthropic的文章对此进行了比较说明)。这篇文章针对SAE本身的一些特性进行了探索,看之前我强烈推荐读者去读Scaling Monosemanticity!从中了解一些基础概念,而且内容非常好(之前想解读这篇但是实在太长了)。
Sparse Autoencoder(SAE,稀疏自编码器)为从LM中提取可解释特征提供了一种很有前景的无监督方法。由于LM学习了许多概念,自编码器本身参数规模非常大才能恢复所有相关特征。但研究SAE的缩放性质很难,主要是因为需要在重构目标和稀疏性目标之间保持平衡,以及存在无效潜变量(dead latent)的问题。作者在本文提出使用k-SAE来直接控制稀疏性,简化调参并改进重构-稀疏性边界。此外,作者还找到一种能减少无效潜变量的修改,在他们所尝试的最大规模上的SAE上也有效。使用这些技术,作者找到了关于SAE大小和稀疏性的清晰的缩放律。他们还引入了几个新的评估特征质量的指标,它们基于假设特征的恢复、激活模式的可解释性和下游效果的稀疏性,这些指标通常随SEA大小的增加而提高。为展示方法的可扩展性,作者在40亿token的GPT-4的激活上训练了一个16M潜变量的自编码器。另外,作者还发布了针对开源模型的代码和自编码器,以及一个可视化工具。
首先第一部分,介绍了一个SOTA的训SAE的方法。作者选取最小为GPT-2 small到最大为GPT-4模型的残差流激活训了自编码器,选取的层较深,使得这些激活既能捕捉丰富的特征而不只专注于预测下一个token。训练涉及对输入进行标准化处理,并在实验中使用了固定的上下文长度。训练完成后,自编码器的性能通过评估稀疏度和重构误差(MSE)来衡量。此外,学习率的选择是通过在小规模上做实验并推断大规模的趋势来完成的。
他们选取的基线是一个ReLU自编码器(公式见上),Loss形式比较典型(见下),结合了重建误差与对稀疏性的惩罚项,而作者做的改变就是利用k-稀疏自编码器,它通过使用一个激活函数(TopK)来直接控制活跃潜变量的数量,该函数只保留k个最大的潜变量,其余置零。其和基线的区别在于编码器中的ReLU替换成了TopK,训练损失用简单的MSE,没有第二项。这种方法简化了模型训练的损失计算,省去了对L1正则化的需求,使得模型调优更直接、更简单。此外,k-稀疏自编码器在实验中显示出比传统ReLU自编码器更好的性能,特别是在处理稀疏性与数据重构之间的平衡时。此技术还能提高模型处理数据的一致性,即通过将小激活置零来增加模型输出的单义性。
另外关于防止无效潜变量,他们采用了两种策略:1.通过将编码器初始化为解码器的转置来提高潜变量的活性;2.引入一种辅助损失函数,专门针对无效潜变量的重构误差进行优化。这些措施显著降低了无效潜变量的比例,即使在规模较大的自编码器模型中也能维持较低的无效潜变量比率,提高了训练效率和模型性能。
其次第二部分,SAE的缩放定律。作者假设说必须要很大数量的稀疏特征来表述前沿模型的状态,他们考虑两种主要方法来选择自编码器的规模与token量:1.训练自编码器达到给定可用计算资源的最优MSE,不考虑收敛性。他们发现MSE和算力遵循幂律L(C),尽管最小的模型偏离了趋势,见figure1左边(figure1右边则是作者发现的MSE随着latents数量n和稀疏水平k的变化曲线);2.另一种是训练到模型收敛,这更注重模型的最终性能,适用于对模型精确度要求较高的情况。他们发现,能够收敛的最大学习率与1/sqrt(n)成比例(figure3)。而且,对于L(N)的最优学习率大约是L(C)最优学习率的四分之一。
另外,figure4展示了当固定稀疏程度时,对于更大的目标模型,达到相同的MSE需要更多的潜变量。
最后第三部分,作者介绍了评估潜变量质量的指标,而且根据这些指标能找到更大规模的SAE,毕竟自编码器的最终目标不是为了改善稀疏性-重构边界,而是为了找到对应用有用的特征,例如机制解释性(mechanistic interpretability)。他们使用以下指标来衡量自编码器的质量,细节很多我就不多余展开了,直接在问题后面给出他们文章中的部分结论:
下游损失:如果将残差流潜变量替换为自编码器重构后的潜变量,语言模型损失表现如何?作者发现,k-稀疏自编码器在下游损失上的改进比在MSE上的改进更明显(figure5a)。而且当保持稀疏性L0不变且仅改变自编码器大小时,MSE与KL散度和交叉熵损失的差值呈现出明显的幂律关系(figure5b)。当然作者认为更自然的度量标准是考虑训练具有可比下游损失的语言模型所需的相对预训练计算量,例如,将16M潜变量的自编码器被替换到GPT-4中时,所获得的语言建模损失相当于GPT-4预训练计算量的10%,这意味着自编码器能够在利用10%的算力时,保持与原始模型相似的性能。
探针损失:自编码器是否能够恢复我们认为它们可能具有的特征?作者构建了一个二分类集并训了一维探针,发现TopK方法优于ReLU,且均优于直接使用残差流。
可解释性:是否存在既必要又充分的简单解释来解释自编码器潜变量的激活?一般来说,总潜变量更多且活跃潜变量较少的自编码器最容易用N2G(一种输出包含通配符的n-gram集合来解释神经元激活模式的方法)建模。证据还表明,TopK模型比其ReLU模型具有更少的伪阳性激活。
消融稀疏性:对单个潜变量进行消融是否对下游logits产生稀疏影响?是的,实验结果表明在适当的k值范围内,自编码器潜变量的消融确实对下游logits产生了稀疏影响,这意味着每个潜变量对输出的影响是有限的和可解释的。
最后作者还对TopK激活函数做了更多讲解,略过。
撰文:戴剑波;编辑:戴剑波
未经本公众号授权不得转载,欢迎转发。