PairSAE: 从蛋白质共折叠中的配对表示实现机制可解释性

# PairSAE：蛋白质共折叠中配对表示的机制可解释性

来源：https://arxiv.org/html/2606.27440

Giosue Migliorini  
加州大学尔湾分校  
&Aristofanis Rontogiannis  
Flagship Pioneering  
&Grigori Guitchounts  
Flagship Pioneering  
&Nicholas Franklin  
Flagship Pioneering  
&Axel Elaldi  
Flagship Pioneering  
&Olivia Viessmann  
Flagship Pioneering  

###### 摘要

结构生物学基础模型在预测生物分子结构方面取得了显著性能，并显示出在蛋白质和小分子设计方面的潜力。然而，理解*哪些*内部特征驱动其输出仍然具有挑战性。标准稀疏自编码器（SAE）在Transformer风格的*序列嵌入*上效果显著，但无法直接迁移到类似pairformer的架构：简单地对配对表示进行操作会导致特征数量呈二次爆炸，并且会模糊分布在对序列和配对表示中的概念。我们提出**PairSAE**，它通过**N-模式SVD**将配对张量总结为词符级别的相互作用角色，然后使用稀疏自编码器学习一组*共享的*词符级别特征，这些特征可以解码为序列和配对表示。在针对PLINDER蛋白质-配体复合物的Boltz-2激活上进行评估，PairSAE产生了与UniProt注释对齐的可解释特征，并预测了Boltz-2的亲和力值。这些结果表明，PairSAE将结构生物学基础模型的潜在空间与可解释的结构概念联系起来，阐明了模型“知道”的内容，同时避免了限制传统SAE的pairformer诱导的陷阱。

## 1 引言

蛋白质结构预测的最新进展，尤其是AlphaFold和Boltz模型（Jumper等人，2021（https://arxiv.org/html/2606.27440#bib.bib18）；Abramson等人，2024（https://arxiv.org/html/2606.27440#bib.bib19）；Wohlwend等人，2025（https://arxiv.org/html/2606.27440#bib.bib20）；Passaro等人，2025（https://arxiv.org/html/2606.27440#bib.bib21）），已经改变了结构生物学，并扩展到DNA、RNA、小分子及其复合物。这些模型在未明确编码物理或化学定律的情况下取得了高精度，而是利用实验确定结构的大型数据集中的统计规律性。这提出了一个核心问题：它们在多大程度上隐含地捕捉了控制折叠的物理和化学原理？回答这个问题需要可解释性——即允许研究人员评估预测可靠性和生物物理合理性的解释——然而深度模型的高度非线性架构使得此类分析具有挑战性。

机制可解释性方面的工作表明，当网络表示的稀疏特征超过可用维度时，可能会出现明显的“多语义”神经元，迫使特征叠加而非清晰定位（Elhage等人，2022（https://arxiv.org/html/2606.27440#bib.bib29）；Scherlis等人，2022（https://arxiv.org/html/2606.27440#bib.bib31））。稀疏自编码器（SAE）通过将叠加特征通过稀疏的过完备字典分解为更单语义的组件来解决这一问题（Olshausen and Field，1997（https://arxiv.org/html/2606.27440#bib.bib24）；Ng等人，2011（https://arxiv.org/html/2606.27440#bib.bib23）；Elhage等人，2022（https://arxiv.org/html/2606.27440#bib.bib29））。在大语言模型（Cunningham等人，2024（https://arxiv.org/html/2606.27440#bib.bib30）；Bricken等人，2023（https://arxiv.org/html/2606.27440#bib.bib28））和蛋白质语言模型（pLM）（Simon and Zou，2024（https://arxiv.org/html/2606.27440#bib.bib36）；Adams等人，2025（https://arxiv.org/html/2606.27440#bib.bib3））中的结果表明，SAE可以恢复可解释的潜在特征，使其成为探索蛋白质结构预测模型所学表示的良好候选。

然而，将SAE直接应用于结构预测模型并非易事。与大多数基于Transformer的LLM和pLM不同，这些模型使用带有配对表示的pairformer块以及序列嵌入（Abramson等人，2024（https://arxiv.org/html/2606.27440#bib.bib19））。为每对学习独立的字典会导致二次扩展并混淆分析，而重复的配对-序列交互表明特征可能同时叠加在这两个空间中。为了解决这个问题，我们提出了**PairSAE**，它从一组共享特征中重构序列级别和配对级别的嵌入。应用于Boltz-2，PairSAE将潜在表示与可解释的结构概念联系起来，阐明了模型对折叠的“理解”，并使得结构生物学中更原则性的评估和发现成为可能。

见图注 见图注 见图注 见图注 见图注 见图注

图1：在Boltz-2的第33层和第64层（第三次循环步骤）训练的PairSAE恢复的可解释特征。
*左上*：*大肠杆菌*复合物II（PDB 1NEK）上的特征889；在**跨膜蛋白质**上激活（词符F1=0.58，复合物F1=0.65）。
*中上*：流感N2神经氨酸酶链B（PDB 4K1I）上的特征55；在**二硫键**上激活（词符F1=0.79，复合物F1=0.81）。
*右上*：HIV-1蛋白酶抑制剂（PDB 1HPS）上的特征885；在**蛋白酶活性位点**（与二聚体伙伴共享）上激活（词符F1=0.93，复合物F1=0.93）。
*底部*：带有UniProt真实值（蓝色）的每个残基激活。

## 2 背景

### 2.1 稀疏自编码器

稀疏字典学习（Olshausen and Field，1997（https://arxiv.org/html/2606.27440#bib.bib24）；Lee等人，2006（https://arxiv.org/html/2606.27440#bib.bib16）；Ng等人，2011（https://arxiv.org/html/2606.27440#bib.bib23））是一种历史悠久的表示学习技术，旨在寻找一个字典D={dj}j=1D，通过其列的线性组合与一组从x本身提取的稀疏潜在特征h∈RD来近似输入向量x∈Rn。最近的研究表明，这种公式可以在基于变换器的模型（如大语言模型（Elhage等人，2022（https://arxiv.org/html/2606.27440#bib.bib29）；Templeton等人，2024（https://arxiv.org/html/2606.27440#bib.bib35）；Cunningham等人，2024（https://arxiv.org/html/2606.27440#bib.bib30））、蛋白质语言模型如ESM2（Simon and Zou，2024（https://arxiv.org/html/2606.27440#bib.bib36）；Garcia and Ansuini，2025（https://arxiv.org/html/2606.27440#bib.bib2）；Adams等人，2025（https://arxiv.org/html/2606.27440#bib.bib3）；Gujral等人，2025（https://arxiv.org/html/2606.27440#bib.bib4）；Parsan等人，2025（https://arxiv.org/html/2606.27440#bib.bib17））和基因组学基础模型如Evo 2（Brixi等人，2025（https://arxiv.org/html/2606.27440#bib.bib41）））中产生可解释的特征。这些特征可以通过使用带有稀疏约束的重构损失训练稀疏自编码器来获得。在机制可解释性中，自编码器通常是线性的，对潜在特征有稀疏性和正性约束。一个常见的构造是

h=σ(Ex+benc)，x̂=Dh+bdec，(1)

其中E∈RD×n，benc∈RD是编码器的权重和偏置项，D∈Rn×D，bdec∈Rn是解码器的权重和偏置项，σ是一个诱导稀疏性的非线性函数，如ReLU（Cunningham等人，2024（https://arxiv.org/html/2606.27440#bib.bib30）；Bricken等人，2023（https://arxiv.org/html/2606.27440#bib.bib28））、TopK（Gao等人，2025（https://arxiv.org/html/2606.27440#bib.bib27））、BatchTopK（Bussmann等人，2024（https://arxiv.org/html/2606.27440#bib.bib37））等。

### 2.2 Pairformer表示

在pairformer模块（Abramson等人，2024（https://arxiv.org/html/2606.27440#bib.bib19））中，长度为Ntok的输入序列x被编码为序列嵌入矩阵S∈RNtok×ns，其第i行是词符嵌入向量si∈Rns，以及一个三维配对表示张量Z∈RNtok×Ntok×nz，其(i,j)切片Zi,j∈Rnz表示词符对(i,j)的嵌入。这里，ns和nz分别是序列级别和配对级别的嵌入维度。配对表示通过在pairformer每层中偏置注意力logits来控制序列级别嵌入更新的信息流（Abramson等人，2024（https://arxiv.org/html/2606.27440#bib.bib19））。配对表示最初随Evoformer引入，后者是AlphaFold 2架构的一部分（Jumper等人，2021（https://arxiv.org/html/2606.27440#bib.bib18））。

## 3 PairSAE

#### 动机。
我们的目标是获得词符级别的稀疏特征，这些特征能够联合重构序列和配对表示。配对嵌入没有施加结构，因此我们首先通过保留行/列交互角色的方式将它们压缩为一个词符级别的摘要。然后，一组共享的特征解码为两种表示类型。

#### 总结配对交互。
我们提出对张量Z进行**N-模式奇异值分解（SVD）**，并利用它构建一个专门设计的序列表示，以融入每个词符在序列中如何交互的信息。我们在附录A（https://arxiv.org/html/2606.27440#A1）中简要介绍了N-模式SVD。通过检查左奇异向量矩阵U(1),U(2)∈RNtok×Ntok（De Lathauwer等人，2000a（https://arxiv.org/html/2606.27440#bib.bib33）；Vasilescu and Terzopoulos，2002（https://arxiv.org/html/2606.27440#bib.bib32））的模式1和模式2矩阵，可以获得每个词符在Z中作为行和作为列所起作用的信息。我们考虑截断到前r列U:,1:r(1)和U:,1:r(2)，按奇异值排序。然后我们将它们按列拼接成一个新的序列级别嵌入：

mi=[Ui,1:r(1)∥Ui,1:r(2)]，i=1,...,Ntok。(2)

#### 自编码器。
PairSAE的潜在特征通过将这个新的嵌入矩阵与pairformer模型原始的序列级别表示拼接后进行编码得到：

hi=σ(E[si∥mi]+benc)，i=1,...,Ntok，(3)

其中E∈RD×(ns+2r)，benc∈RD，σ是BatchTopK和ReLU的组合（Bussmann等人，2024（https://arxiv.org/html/2606.27440#bib.bib37））。我们使用共享特征解码到两个空间：

ŝi=Dshi+bdecs，ẑi,j=Dzrowhi+Dzcolhj+bdecz，i,j∈{1,...,Ntok}，(4)

其中Ds∈Rns×D，bdecs∈Rns，Dzrow，Dzcol∈Rnz×D，bdecz∈Rnz。

#### 损失函数。
选择字典大小D会显著影响学习到的特征类型，更大的字典并不一定意味着下游任务性能更好（Templeton等人，2024（https://arxiv.org/html/2606.27440#bib.bib35）；Gao等人，2025（https://arxiv.org/html/2606.27440#bib.bib27））。最近研究表明，Matryoshka SAE损失（Bussmann等人，2025（https://arxiv.org/html/2606.27440#bib.bib38））在多种字典大小下均表现出鲁棒性能。这是通过在几个嵌套级别上切片解码器矩阵和特征向量（我们使用三个嵌套宽度c1τij}。(9)

我们在PLINDER训练集中的7,680个复合物上选取最佳阈值τij，基于验证集（2,560个复合物）上的F1分数为每个概念j选择最佳特征i，最后在由5,120个随机选择的复合物组成的测试集上报告F1分数。我们将分析限制在至少出现在五个不同复合物上的概念。

### B.3 假设生成

我们通过获取PLINDER中980个系统（来自PairSAE训练集的保留集）的Boltz-2预测亲和力值，为LASSO回归构建训练集。测试集是Posebusters（Buttenschoen等人，2024（https://arxiv.org/html/2606.27440#bib.bib12））。在图B.4（https://arxiv.org/html/2606.27440#A2.F4）中，我们强调了这些在响应变量方面的差异，并注意到训练集中缺乏高亲和力复合物。LASSO回归使用`sklearn.linear_model.Lasso`（Pedregosa等人，2011（https://arxiv.org/html/2606.27440#bib.bib8））进行拟合。两个模型的详细结果见表B.2（https://arxiv.org/html/2606.27440#A2.T2）。参见图B.5（https://arxiv.org/html/2606.27440#A2.F5）中基于R3-L33 PairSAE的Boltz-2亲和力真实值与预测值。如第B节（https://arxiv.org/html/2606.27440#A2）所述，我们未在分析中使用MSA。在图B.6（https://arxiv.org/html/2606.27440#A2.F6）中，我们比较了有MSA和无MSA时Boltz-2的输出，注意到预测亲和力值和亲和力概率存在显著差异。由于基于MSA的预测基于更多信息，我们预计这些预测更准确。在未来的工作中，我们旨在使用MSA输入重复我们的分析。

表 B.2：基于最大池化PairSAE特征的LASSO回归。ρ表示Spearman秩相关系数。

见图注

图 B.4：假设生成任务中使用的训练集和测试集的Boltz-2亲和力值。

见图注

图 B.5：Boltz-2的负亲和力值（越高表示亲和力越强），以及基于R3-L33 PairSAE的LASSO回归（公式(8)）的预测值。

见图注

图 B.6：在Posebusters数据集中，有MSA和无MSA时Boltz-2亲和力值和亲和力概率的比较。

## 附录C 附加结果

在图C.7（https://arxiv.org/html/2606.27440#A3.F7）中，我们展示了LASSO系数如何随不同正则化水平变化，并突出显示特征