可训练子网络揭示蛋白质语言模型中结构知识的组织机制

《PLOS Computational Biology》：Trainable subnetworks reveal insights into structure knowledge organization in protein language models

本文创新性地运用可训练子网络技术，系统性地探究了蛋白质语言模型（PLM）内部如何编码和分离不同层级的结构信息。研究通过掩蔽语言建模（MLM）目标训练的子网络，成功定位了与特定蛋白质结构类别（如CATH分类和二级结构）相关的模型参数，揭示了PLM对序列层面特征的高度敏感性及其在极粗和极细粒度信息上的解纠缠能力。尤为重要的是，研究发现即使语言建模性能的微小变化也会显著削弱基于PLM的结构预测能力，为理解预训练PLM中特征纠缠提供了新框架，并有助于改进所学表征与已知生物学概念的对齐。

　　

可训练子网络揭示蛋白质语言模型中结构知识的组织机制

摘要

蛋白质语言模型（PLM）通过掩蔽语言建模（MLM）目标进行预训练，已被证明在一系列结构相关任务中有效，包括高分辨率结构预测。然而，目前尚不清楚这些模型在多大程度上将其学习的参数中蛋白质结构类别进行分解。本研究引入可训练子网络，该网络可掩蔽负责特定蛋白质结构类别语言建模性能的PLM权重。利用CATH分类法和二级结构元素定义的注释，系统性地训练了39个PLM子网络，针对不同分辨率的序列水平和残基水平特征。通过这些PLM子网络，评估了PLM中的结构分解如何影响下游结构预测。结果表明，PLM对序列水平特征高度敏感，且能够主要解纠缠极粗或极细粒度的信息。此外，结构预测对分解后的PLM表征高度敏感，语言建模性能的微小变化可显著损害基于PLM的结构预测能力。本研究提出了一个研究预训练PLM内特征纠缠的框架，并可用于改进所学PLM表征与已知生物学概念的对齐。

作者摘要

蛋白质主导细胞过程，其功能源于氨基酸序列编码的三维结构。因此，从序列预测蛋白质结构已成为现代生物序列模型的核心能力。仅通过通用语言建模目标在序列上训练的蛋白质语言模型，在结构预测方面非常准确，并广泛应用于蛋白质设计和工程工作流程中。然而，关于这些模型的权重如何编码不同蛋白质结构特征之间的关系，目前知之甚少。随着蛋白质语言模型在数据、计算和模型规模上的扩展，这一研究方向日益重要。本研究证明，可以分离出对应于特定定义结构类别的模型权重子集（即子网络）。结果表明，使用蛋白质语言模型的结构预测准确性对这些子网络高度敏感，即使语言建模性能的变化很小。当应用于不同的结构类别时，该方法表明结构知识的分布方式反映了蛋白质结构多样性的连续谱。本研究揭示了生物学相关信息如何在蛋白质语言模型权重中组织，并为未来训练模型提供了更知情和可解释的基础。

引言

理解蛋白质结构对于破译生物功能至关重要，因为蛋白质的结构决定其分子稳定性、相互作用和活性。最近，仅基于序列数据训练的蛋白质语言模型（PLM）已被证明可以学习隐式编码结构信息的表征。这些模型在广泛的蛋白质工程任务中被证明有效，包括结构预测、功能注释、突变效应估计甚至新蛋白质的设计。许多PLM通过自监督的掩蔽语言建模（MLM）目标进行预训练，其中模型的任务是预测序列中随机掩蔽令牌的氨基酸身份。由于蛋白质结构根本上由氨基酸序列决定，PLM可以在其权重中隐式编码结构信息。ESM-2是一个在进化尺度数据上训练的蛋白质语言模型家族，表明语言建模任务的性能可以预测PLM预测结构的质量。PLM表征现在广泛用作预测结构坐标模型的输入，缩放分析表明MLM性能的改进提高了单序列结构预测的准确性。预训练语言模型在蛋白质序列上的简单性和有效性导致人们对其内部机制的理解日益增长，使可解释性成为一个重要的研究方向。关于PLM学习的结构特征，先前的工作表明蛋白质接触信息存储在注意力图中，PLM学习共同进化基序，模拟对基本生物物理学的理解，并且PLM的稀疏潜在特征捕获了已知的功能生物物理特性和基序。然而，迄今为止，仍然不清楚（i）结构信息是否以及如何在学习的PLM权重中被分解和存储，以及（ii）这种分解是否影响下游结构预测任务的性能。揭示预训练语言模型中概念是否存在此类分解的一种方法是通过子网络发现。子网络是预训练模型权重的稀疏计算子图，负责特定任务或输入类的性能。在自然语言领域，子网络发现已被广泛用于通过定义概念和定位编码它们的权重、神经元或层来发现预训练期间学习的语言属性（如语义、句法和关系实体）。本工作重点是利用子网络来探究预训练PLM ESM-2中蛋白质结构类别的分解。目标是找到PLM子网络——原始模型权重的稀疏子图——当被隔离时，抑制模型对一类输入做出正确预测的能力，同时保留对所有其他类别输入的MLM性能。作为主要分析的一部分，系统性地训练了39个这样的PLM子网络，以在不同尺度的CATH层次结构上抑制残基或序列水平的结构信息。结果表明，结构类别确实以分解的方式编码在PLM权重中，并且尽管非抑制输入达到ESM-2水平的困惑度，但结构预测仍然受到统计学上的显著扰动。本研究的子网络方法和执行的一系列分析共同提供了关于PLM如何在学习参数中组织结构特征的见解。

方法

抑制和维护输入由结构注释定义

蛋白质二级结构是指氨基酸残基骨架原子的局部空间排列。蛋白质序列中的每个残基采用一种二级结构，可分类为α螺旋、β折叠或环。由此产生的结构排列使得能够在序列水平上对蛋白质进行进一步分类，根据共享的架构和进化特征将它们分组。本工作中，使用DSSP在残基水平（α螺旋、β折叠、环）定义结构类别，并使用CATH分类法在序列水平（类、架构、拓扑、同源超家族）定义结构类别。

符号

将预训练的PLM表示为f(x; θ)，其中x代表输入序列，θ是模型权重。训练过程的目标是预测一个二元掩码m ∈ {0, 1}^K，其中K表示要学习掩码的PLM参数数量。可以通过取二元掩码和预训练PLM权重的Hadamard积来获得子网络，即f(x; m ⊙ θ)。子网络被独立训练以抑制序列水平或残基水平输入的结构类别（即所有抑制输入属于由结构注释定义的同一类别）。将抑制输入序列定义为x^sup，所有其他输入序列（即维护输入）定义为x^maint。对于残基水平抑制，让A是从DSSP获得的注释，其中包含一组J个位置，对应于{α螺旋、β折叠、环}中的残基。将此DSSP注释的J个抑制输入定义为x_J^sup。维护输入是具有注释A的残基的所有剩余互补位置，然后定义为x_{not J}^maint。

子网络训练目标

为了获得子网络f(x; m ⊙ θ)，使用加权损失学习二元掩码m，该损失包括以下组成部分：

1. 1.
   抑制目标。 如果掩码工作正常，子网络应难以准确重建其抑制输入。换句话说，关于抑制输入令牌的经过良好校准的预测分布应相对于词汇表V大致均匀分布。因此，将抑制损失定义为最小化（i）子网络在对应于抑制输入的令牌上的预测分布与（ii）词汇表中令牌的均匀参考分布（表示为U(V)）之间的Kullback–Leibler（KL）散度。对于序列水平抑制，这对应于L_sup^seq= D_KL( U(V) || f(x^sup; m ⊙ θ) )。类似地，对于残基水平抑制，L_sup^res= D_KL( U(V) || f(x_J^sup; m ⊙ θ) )。
2. 2.
   维护-KL目标。 即使在存在掩码的情况下，子网络也应保留完整PLM在维护输入上的预测行为。因此，作为维护目标，还旨在最小化（i）子网络在对应于维护输入的令牌上的预测分布与（ii）预训练PLM在相同元素上的预测分布之间的KL散度。对于序列水平抑制，这表示为L_maint-KL^seq= D_KL( f(x^maint; θ) || f(x^maint; m ⊙ θ) )。类似地，对于残基水平抑制，L_maint-KL^res= D_KL( f(x_{not J}^maint; θ) || f(x_{not J}^maint; m ⊙ θ) )。
3. 3.
   维护掩蔽语言建模（MLM）目标。 在实践中，维护-KL目标是不够的，因为它只强制子网络和原始PLM输出分布之间的相似性，而不确保子网络保留将正确概率分配给预测令牌的能力。先前的工作通过一组消融实验证明所有3个损失分量对于实现子网络在抑制和维护输入上的期望行为是必要的，并且发现省略维护-KL或维护-MLM都会增加维护输入上的困惑度。因此，引入额外的维护-MLM损失，确保子网络仍然可以将适当的概率质量分配给正确的对应令牌，保留其整体语言建模行为。为了在维护输入上包含MLM目标，随机选择15%的序列位置，在这些位置上计算MLM损失。在这些M个位置中，80%被替换为掩码令牌，10%被随机替换，10%保持不变。即，L_maint-MLM^seq= -Σ_i∈Mlog f(x_i^maint; m ⊙ θ)。为了执行残基水平抑制，再次随机选择一组M个掩码位置，采用相同的掩码和突变方案；然而，这次计算MLM损失是关于与特定残基水平注释A的位置重叠的掩码索引，将其表示为M ∩ not J。这可以表示为L_maint-MLM^res= -Σ_{i∈M ∩ not J}log f(x_i^maint; m ⊙ θ)。

总体而言，最终的加权训练目标然后表示为和L = λ₁L_sup+ λ₂L_maint-KL+ λ₃L_maint-MLM，其中λ₁, λ₂, λ₃是超参数。

可微分权重掩码用于子网络

遵循先前的工作，采用可微分权重掩码方案来学习m。这里，每个二元掩码参数m_i从Gumbel分布中采样。为每个第i个参数学习一个logit l_i，并通过以下Gumbel softmax变换获得单位区间上的连续掩码分数s_i= σ( (l_i+ g_i) / τ )，其中σ是sigmoid函数，τ是温度缩放超参数，g_i是从标准均匀分布中抽取的随机变量。这种Gumbel噪声在logit采样过程中引入了随机性，导致在训练期间探索不同的二元掩码配置。通过连续掩码分数集合s反向传播，并使用以下公式获得二值化掩码值m：m_i= 1_{si> T}，其中1是指示函数，T是掩码分数阈值，stopgrad防止通过m中的离散值反向传播。这种阈值操作允许通过连续掩码分数流动梯度，同时仍然关于二元预测计算损失，从而实现可微分训练。稀疏性定义为学习的二元掩码中零的比例——因此在子网络中——计算为Sparsity = (1/K) Σ_i=1^K(1 - m_i)。

模型架构和数据集

在主要分析集中，在ESM-2 6.5亿参数模型中学习子网络。虽然其他二元掩码方法通常关注最后几层以捕获细粒度概念或属性，但选择学习完整模型（即所有层）上的掩码，这防止了对弱早期层信号和虚假后期层相关性的任何依赖。为了进行额外验证，还将子网络方法应用于三个具有不同大小、架构和预训练任务的最先进PLM：ProtBERT-UR100、CARP-640M和Dayhoff-170M-UR90。

为了训练和评估，使用了CATH S20版本4.3.0发布。该数据集由CATH结构域组成，这些序列在最大20%的成对序列同一性下聚类，且至少具有60%的比对重叠。这确保了低冗余性，同时保持了结构域之间的结构和功能多样性。每个结构域在CATH层次结构的连续级别上都有注释：类（C）、架构（A）、拓扑（T）和同源超家族（H）。使用DSSP获得简化的三向残基水平二级结构类别。在训练期间，筛选出长度在64到1024个残基（ESM-2的最大上下文窗口）之间且具有可用PDB结构的CATH结构域集合。总共有8886个CATH结构域。对于每个子网络，将这些数据随机分成70%用于训练，20%用于验证，10%用于保留集。根据训练分割中存在的抑制和维护输入学习子网络。在评估时，对所有分割执行MLM评估，这允许汇总通过训练分割在训练期间看到的抑制和维护输入的性能，以及验证和测试分割中未见输入的性能。在验证和测试集上使用ESMFold折叠主干进行结构预测评估，以限制计算成本。

结果

稀疏子网络使得能够在ESM-2权重中成功分解

如果能够识别一个子网络，该网络选择性地降低抑制输入上的MLM性能，同时相对于完整PLM保持维护输入上的性能，则认为蛋白质的结构类别在PLM权重中被分解。使用困惑度（即负对数似然的指数）测量MLM性能，分别计算子网络和预训练ESM-2模型在抑制和维护输入上的困惑度。

通过训练过程，子网络可以通过识别PLM权重的稀疏子图来实现这种差异性能，其中编码抑制输入类别信息的参数被置零。由于子网络训练过程没有明确促进任何稀疏正则化，学习的稀疏性百分比是任何发现的隐式分解所固有的。量化了跨结构级别的平均学习稀疏性。随着抑制类别在CATH层次结构中变得更加细粒度，学习的稀疏性降低，表明较小的结构类别可以通过置零更少的参数来分解。这一趋势表明子网络的学习百分比稀疏性与数据中结构注释的频率（即注释粒度）之间存在相关性。

训练后的子网络在抑制输入上 consistently 产生更高的困惑度，暗示结构相关的表征可以在PLM权重空间中被选择性地分解。针对属于二级结构（即α螺旋或β链）的残基的抑制导致这些残基的困惑度大约增加两倍。对更广泛的序列水平注释（如CATH类，其中抑制集大小最大）的抑制产生类似于随机噪声的预测分布，困惑度接近20——这是对20种氨基酸进行均匀猜测的期望值。当抑制同源超家族时，尽管抑制集小得多（<100个序列），也观察到类似的困惑度增加，表明分解在最粗和最细的注释粒度水平上最强。相比之下，结构类别在中间CATH级别（如架构和拓扑）上分解较弱，证据是抑制输入上的困惑度低于其他级别。一个可能的解释是，分解在类和同源超家族上最有效，这些类别代表了蛋白质序列空间中不同的进化边界。粗粒度的类捕获全局的、基本的二级结构组成（即主要是α、β或“混合”α-β类），而细粒度的超家族反映了局部关系，这两者都被证明在PLM嵌入空间中形成良好分离的区域。相比之下，中间类别如架构和拓扑在结构上是异质的，结合了仅共享部分基序但在几何形状上不同的折叠，使它们 inherently 更难以分离，因此子网络更难将它们隔离成不同的组。当在相同的抑制目标下跨多个ESM-2种子重复训练时，产生的子网络表现出高度相似的稀疏模式，并在抑制和维护输入上实现几乎相同的性能，表明优化收敛到稳定且功能等效的解决方案。

作为对照，训练了两种类型的子网络。第一个对照是抑制N个随机选择的序列，其中N值的选择模仿CATH序列类别的大小。第二个对照是抑制随机选择的残基。评估残基对照子网络分别在α螺旋和β折叠上。随机序列和随机残基子网络都未能对抑制输入产生差异性能，唯一的例外是N=2000的序列对照，导致抑制输入上的困惑度相对于ESM-2（8.8）增加了2.6点。相比之下，训练用于抑制CATH类级别序列类别的子网络能够获得大于30的困惑度，同时实现维护目标。推断序列对照中困惑度的任何增加是因为掩码学习过程删除了编码跨序列广泛共享特征的参数。由于随机选择的序列不共享有意义的结构特征，子网络无法识别特定于抑制结构信息的参数，而是必须置零对一般MLM性能重要的权重，这可能导致所有输入上的困惑度均匀增加。

还将子网络方法应用于三个额外的具有不同大小、架构和预训练任务的最先进PLM，每个都在UniRef数据上训练过：ProtBERT-UR100（4.2亿参数）、CARP-640M（6.4亿参数）和Dayhoff-170M-UR90（1.7亿参数）。这些额外PLM中的子网络显示出与ESM-2相似的趋势，因为它们通过修剪少于3%的模型参数被发现，在类和同源超家族级别上分解最强，并且与残基相比，在更大程度上分解序列类别。

总之，百分比稀疏性和MLM评估表明，子网络可以通过识别PLM中结构信息的稀疏分解来有针对性地增加困惑度，与蛋白质结构多样性的连续性质一致。

子网络扰动抑制和维护输入上的结构预测准确性

在语言建模任务上评估子网络之后，评估了通过子网络抑制如何影响结构预测准确性。ESMFold引入了一个折叠主干——AlphaFold2的Evoformer的简化版本——它将语言模型表征转换为3D结构。使用这个冻结的主干允许隔离子网络诱导的序列表征变化，并能够研究它们如何影响结构预测准确性。由于折叠主干充当固定解码器，模板建模（TM）分数或预测局部距离差异测试（pLDDT）的任何退化，或增加的均方根偏差（RMSD），直接反映了序列表征中相关结构信息的丢失。对于每个输入序列，从子网络和ESM-2提取序列表征，并将它们作为单独输入传递给ESMFold折叠主干进行结构预测。两个预测结构都独立地与真实PDB结构对齐，以获得每个模型预测的TM分数、RMSD和pLDDT。对验证数据集中的所有抑制输入x^sup和所有维护输入x^maint重复此过程。

首先，评估子网络是否导致结构预测准确性相对于ESM-2在这些输入上的性能下降。对于每个输入，计算RMSD的变化，其中ΔRMSD = RMSD_subnetwork- RMSD_ESM-2。这些差异在抑制输入上 consistently 更大，在维护输入上较小。这表明子网络能够消融PLM权重中与抑制输入相关的结构信息。然后分别对抑制集和维护集中的输入执行配对t检验，比较RMSD_subnetwork和RMSD_ESM-2。两个检验都产生显著的p值，表明尽管仅针对抑制输入，子网络也显著影响维护输入上的结构预测。也就是说，抑制和维护输入上的ΔRMSD都是统计学上显著的。使用TM-score和pLDDT重复此过程，并观察到相似的趋势。推断抑制输入上的显著效应和维护输入上结构预测准确性的最小变化——加上显著的p值——表明冻结的ESMFold主干对细微的表征变化敏感，而不是反映序列表征中缺乏模块化。子网络掩码学习过程并非为保持结构预测准确性而优化；相反，它旨在基于序列水平表征和MLM性能划分PLM的权重空间。因此，下游结构指标的小变化是预期的，该分析主要用于揭示序列水平分解如何影响结构预测能力。

接下来，评估这些差异的幅度是否对于抑制输入大于维护输入跨CATH和二级结构类别。也就是说，旨在确认|ΔRMSD_sup| > |ΔRMSD_maint|。这些差异 consistently 为正，表明子网络总是导致抑制输入上的RMSD增加大于维护输入。还通过应用双样本Kolmogorov–Smirnov（KS）检验 between |ΔRMSD_sup| and |ΔRMSD_maint|，拒绝抑制和维护输入表现出相同RMSD增加幅度分布的原假设，确认幅度差异是 consistently 统计学上显著的。对TM-score和pLDDT重复相同的分析，并观察到相同的趋势。还观察到每个指标变化幅度的差异平均而言对于Alpha-Beta蛋白质更大，其次是Mainly Beta，然后是Mainly Alpha蛋白质。推断这可能是因为混合二级结构组成的Alpha-Beta折叠施加了更严格的结构约束，这些折叠将α螺旋和β折叠耦合到相互依赖的架构中。与此解释一致，使用ProteinMPNN困惑度评估序列水平结构约束，并发现支持趋势：Alpha-Beta蛋白质表现出最低的困惑度，其次是Mainly Beta，然后是Mainly Alpha蛋白质，表明混合Alpha-Beta折叠具有最受限的序列可变性和最强的结构-序列耦合。

结构预测对抑制的序列特征更敏感

α螺旋和β折叠是局部二级结构元素，当在序列中占主导时，定义蛋白质结构域的总体CATH类。由主要α螺旋组成的序列形成“Mainly Alpha”结构