SAE助力蛋白质模型毒力特征可解释性审计

当前生成式人工智能模型在蛋白质设计领域取得了显著进展，但这项强大的技术也带来了潜在的安全风险——可能被用于不透明地生成有害蛋白质。为了应对这一问题，研究团队推出了VFUSE（基于稀疏自编码器的毒力特征理解工具），这是一种机制可解释性方法，通过在扩散Transformer的激活空间中训练稀疏自编码器（SAE），来审计蛋白质生成模型中的有害特征。

研究背景与前期探索

此前，大语言模型的机制可解释性研究已经十分丰富，从SAELens、Neuronpedia到Golden Gate Claude，针对卷积神经网络和视觉Transformer的相关探索也有不少成果。但蛋白质模型的可解释性研究仍然相对不足，目前仅有少数早期探索方向，仍有大量未解决的问题等待解答。团队此次的研究核心目标是验证：在RFDiffusion3和RoseTTAFold3这两款热门的开源蛋白质模型上训练的稀疏自编码器，能否以可解释的方式完成有害与无害蛋白质的分类任务。

实验设计与数据准备

为开展本次研究，团队针对RFDiffusion3（一款生成式蛋白质模型，简称RFD3）和RoseTTAFold3（一款类似AlphaFold的蛋白质结构预测模型，简称RF3）的扩散Transformer激活，训练了Matryoshka Batch TopK稀疏自编码器。实验数据集共包含1475对长度匹配的良性与有害蛋白质样本，样本来源涵盖多个公开数据库。在模拟基于RFD3的有害基序生成时，团队会将原始蛋白质坐标添加5埃的噪声，通过部分扩散后再还原为原始蛋白质结构，研究中还以蝰蛇毒ammodytoxin A蛋白作为典型案例展示。

核心实验结果

在分别基于原始模型激活和SAE隐空间激活训练逻辑回归分类探针后，团队发现针对部分模型层，SAE探针的分类效果要优于原始激活探针，其中在RF3的第12层，针对特定数据集的分类AUROC达到了0.848。为了避免模型记忆折叠家族的特征，团队使用专业工具基于同源性对样本进行了聚类处理。

更值得关注的是，团队还从SAE中识别出了单义特征，这些特征仅会在有害蛋白质样本上被激活，最高的特征分类AUROC可达0.84。此外，研究还发现特征的判别能力会随着模型层数的加深而提升，尤其是在RFD3模型中，这说明深层的模型已经学习到了更为复杂的蛋白质结构概念。

研究意义与未来展望

据团队介绍，此次研究是首个在全原子扩散模型上训练的稀疏自编码器，同时也是首次针对蛋白质设计模型开展的毒力级特征级审计工作，为安全且可解释的蛋白质生成模型发展铺平了道路。

目前这项研究仅仅是可解释性技术与蛋白质模型结合的初步探索，还有大量值得挖掘的方向：比如蛋白质设计与折叠模型还学到了哪些其他结构特征？包括蛋白质特异性、结合强度、热稳定性、免疫原性等；能否利用这些模型区分真实蛋白质与AI生成的蛋白质？除了训练时内置的条件信号外，能否通过这些可解释性特征来引导蛋白质的生成过程？团队对后续的研究进展充满期待。