蛋白质功能预测这件事做的人不少但真正能在“速度”和“准确率”之间同时站住的方案并不多。大多数同行都有过类似的体验拿一条新测序出来的蛋白序列跑一遍BLAST做同源比对运气好能命中几个已知功能的同源物运气不好就是一片空白只能靠结构域扫描或者GO注释工具硬猜。而深度学习路线呢AlphaFold那一套精度确实高但推理成本摆在那里动辄几分钟到几十分钟一条面对动辄上万条的宏基因组或转录组数据根本跑不动。这篇要聊的思路是把预训练语言模型和基于同源的标签扩散这两条线捏在一起用序列本身的语言学特征做初步判断再用同源关系做标签传播最终在保持秒级推理速度的前提下把功能预测的准确率拉到接近结构方法的水平。适合做蛋白组学、宏基因组注释、酶功能挖掘的同行参考也适合刚接触蛋白质功能预测、想找一个能快速上手的baseline的朋友。1. 为什么单纯靠同源比对或单纯靠语言模型都不够用1.1 同源比对的“覆盖率天花板”问题做蛋白功能注释的人对BLAST和DIAMOND都不陌生。给定一条查询序列在UniProt或NCBI的nr库里搜一圈取top hit的注释直接迁移过来这是最经典的做法。它的逻辑很硬序列相似功能大概率相似。但问题在于这个逻辑有一个隐含前提——数据库里得有足够近的同源物。我实测过一批来自深海沉积物的宏基因组ORF长度在150到400个氨基酸之间。用DIAMOND跑nr库e-value卡在1e-5结果有将近四成的序列连一个像样的hit都找不到。这不是参数没调好而是这些序列代表的蛋白家族在现有数据库里根本没有被表征过。同源比对的覆盖率是有天花板的天花板的高度取决于数据库的完整度而数据库的完整度永远追不上自然界蛋白多样性的增长速度。更麻烦的是即使命中了同源物注释迁移也有风险。一条序列和已知蛋白有40%的相似度但功能可能完全不同。酶学里这种例子太多了一个点突变就能让底物特异性翻转而序列相似度只掉了两三个百分点。所以单纯靠同源比对既解决不了“找不到同源物”的问题也解决不了“找到了但功能不一定对”的问题。1.2 预训练语言模型能补上什么预训练语言模型在蛋白序列上的应用核心思路是把氨基酸序列当成一种“语言”。一条蛋白序列就是由20种氨基酸“单词”组成的句子只不过这个句子的语法规则是亿万年进化写出来的。ProtTrans系列模型比如ProtBert、ProtT5在大规模序列库上做过掩码语言建模训练模型内部已经学到了残基之间的长程依赖、保守motif的上下文模式、甚至一些隐含的结构倾向。这意味着什么意味着即使一条序列在数据库里找不到近亲语言模型也能从它的“遣词造句”里提取出有判别力的特征。比如某个催化位点周围的残基组合模式语言模型见过成千上万次它知道这种模式通常出现在水解酶里。这种能力是同源比对不具备的因为它不依赖显式的数据库命中而是依赖模型在预训练阶段积累的“语感”。但语言模型也不是万能的。它的输出是一个高维向量本身不直接告诉你“这条蛋白是激酶”。你需要一个下游的分类头或者映射机制把向量翻译成GO term或者EC号。而且语言模型的训练目标毕竟是语言建模不是功能分类它学到的表征对功能预测任务来说是有信息冗余和噪声的。所以单独拿ProtTrans的embedding去做功能预测效果通常比同源比对好但离“准确”还有距离。1.3 两条线合并的逻辑起点既然同源比对擅长“有近亲时精确迁移”语言模型擅长“没近亲时靠语感推断”那自然的想法就是能不能让两者互补具体来说先用语言模型给每条序列生成一个功能倾向的初始分布然后利用序列之间的同源关系构建一个图在这个图上做标签扩散让有明确注释的序列把标签“传”给同源但未注释的序列。这个思路的关键在于标签扩散不是简单地把注释复制过去而是在图结构上做迭代传播。一个未注释的节点如果它的多个同源邻居都指向同一个功能那这个功能的置信度就会被放大如果邻居们的功能五花八门那传播后的分布就会很平相当于告诉下游“这条序列的功能不确定”。这种机制比硬迁移要稳健得多。2. ProtTrans表征提取与自注意力池化的实操细节2.1 模型选型ProtBert还是ProtT5ProtTrans家族里常用的两个基座是ProtBert和ProtT5。ProtBert是encoder-only结构输出的是每个残基的上下文嵌入维度通常是1024。ProtT5是encoder-decoder结构encoder部分也能提特征维度同样是1024但它的预训练数据量和训练步数通常更大。我自己的经验是如果你的下游任务需要残基级别的细粒度信息比如预测活性位点ProtBert的逐残基嵌入更直接。如果只需要序列级别的功能标签两者差别不大但ProtT5的encoder输出在某些GO term上会略好一两个百分点。代价是ProtT5的推理速度比ProtBert慢大概30%到40%显存占用也更高。对于标签扩散这个框架我们最终需要的是序列级别的向量表示所以两者都能用。我一般默认用ProtBert因为它在社区里的教程和预训练权重更全踩坑成本低。如果你手头有T5的权重且GPU显存充裕可以试试ProtT5但别指望有质的飞跃。2.2 逐残基嵌入的提取流程提取ProtTrans嵌入的标准流程不复杂但有几个细节容易翻车。以ProtBert为例用HuggingFace的transformers库加载模型from transformers import BertModel, BertTokenizer import torch model_name Rostlab/prot_bert tokenizer BertTokenizer.from_pretrained(model_name, do_lower_caseFalse) model BertModel.from_pretrained(model_name) model.eval() sequence MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG # ProtBert要求残基之间加空格 spaced_seq .join(list(sequence)) encoded tokenizer(spaced_seq, return_tensorspt, paddingTrue, truncationTrue, max_length1024) with torch.no_grad(): output model(**encoded) residue_embeddings output.last_hidden_state # shape: [1, seq_len, 1024]这里有几个坑。第一ProtBert的tokenizer要求输入序列的残基之间用空格隔开如果你直接传连续字符串tokenizer会把多个残基当成一个未知token处理输出就废了。第二序列长度超过1024的蛋白会被截断对于长蛋白你需要么分段处理再拼接么只用前1024个残基的嵌入。我一般建议对超长蛋白做分段每段512残基重叠128残基分别提嵌入后再做平均池化。第三batch推理时要注意padding token的位置后续池化时要mask掉。2.3 自注意力池化为什么比平均池化更合适拿到逐残基嵌入之后下一步是把它压成一个序列级别的向量。最朴素的做法是平均池化把所有残基的嵌入求平均。这个做法快但有个明显问题它把所有残基一视同仁而蛋白里真正决定功能的往往只是少数几个关键残基。平均池化会把这些关键信号稀释掉。自注意力池化attention pooling的思路是让模型自己学一组权重给每个残基分配一个重要性分数然后做加权平均。具体实现可以是一个单层的前馈网络加softmaximport torch.nn as nn class AttentionPooling(nn.Module): def __init__(self, hidden_dim): super().__init__() self.attention nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.Tanh(), nn.Linear(hidden_dim // 2, 1) ) def forward(self, residue_embeddings, maskNone): # residue_embeddings: [batch, seq_len, hidden_dim] scores self.attention(residue_embeddings).squeeze(-1) # [batch, seq_len] if mask is not None: scores scores.masked_fill(mask 0, -1e9) weights torch.softmax(scores, dim-1) pooled torch.bmm(weights.unsqueeze(1), residue_embeddings).squeeze(1) return pooled, weights我对比过平均池化和自注意力池化在GO term预测上的表现。在MF分子功能这一类上自注意力池化能把Fmax提升大概3到5个百分点BP生物过程和CC细胞组分的提升小一些大概1到2个百分点。原因不难理解MF更依赖局部残基的化学性质注意力机制能聚焦到这些关键位置而BP和CC更多是全局属性平均池化已经能抓住大部分信息。注意自注意力池化的权重网络是需要训练的。如果你直接拿预训练好的池化层去用效果可能还不如平均池化。所以这个模块必须和下游分类器一起fine-tune。3. 同源图构建与标签扩散的工程实现3.1 同源关系的获取DIAMOND还是MMseqs2标签扩散的前提是有一个同源关系图。图的节点是蛋白序列边表示两条序列同源。获取同源关系最常用的工具是DIAMOND和MMseqs2。两者都是基于k-mer的快速比对工具速度比BLAST快几个数量级。DIAMOND的blastp模式在灵敏度上略好但内存占用高。MMseqs2的easy-search模式在速度和内存之间平衡得更好而且支持--split-memory-limit参数适合在内存有限的机器上跑大规模库。我自己的选择是如果参考库在百万条序列以内用DIAMOND超过百万条用MMseqs2。构建图的时候e-value阈值建议设在1e-3到1e-5之间。太松会引入大量假阳性边标签扩散会被噪声带偏太严会导致图太稀疏很多节点变成孤岛扩散传不过去。我一般先用1e-3跑一遍看看图的连通分量分布如果最大的连通分量覆盖了80%以上的节点就说明阈值合适如果覆盖不到50%就放宽到1e-2。3.2 标签扩散的数学形式与迭代过程标签扩散的标准形式是给定一个图G(V, E)每个节点有一个标签分布向量。已注释节点的分布是one-hot或者从语言模型得到的软分布未注释节点的初始分布是均匀分布或者语言模型给的先验。然后迭代执行Y^{(t1)} α * S * Y^{(t)} (1-α) * Y^{(0)}其中S是行归一化的邻接矩阵α是扩散系数通常取0.8到0.9。Y^{(0)}是初始标签分布。这个迭代会收敛到一个闭式解Y^* (I - αS)^{-1} Y^{(0)}但在实际工程里我们一般不直接求逆因为图可能有几十万个节点求逆的复杂度是O(n^3)不可接受。用迭代法跑个20到50轮就够了收敛判据是Y的变化量小于1e-6。这里有一个关键设计Y^{(0)}怎么设。如果只用one-hot那未注释节点的初始分布就是均匀的扩散完全靠邻居传。但如果把语言模型输出的功能倾向作为Y^{(0)}的一部分相当于给扩散过程一个“先验引导”收敛会更快最终结果也更好。我的做法是已注释节点用one-hot未注释节点用语言模型分类头输出的softmax分布。这样扩散既利用了图结构又利用了序列本身的语感信息。3.3 大规模图上的内存优化技巧当节点数超过10万时邻接矩阵的存储就成了问题。一个10万节点的图如果平均度是10稀疏矩阵的非零元素有100万个用scipy的sparse矩阵存大概占几十MB还能接受。但如果平均度到50非零元素500万内存就上去了。我的优化策略是不显式构建邻接矩阵而是用边列表加稀疏矩阵乘法。具体来说用scipy.sparse.csr_matrix从边列表构建S然后每次迭代做一次稀疏矩阵向量乘。这个操作在scipy里是高度优化的10万节点、500万边的图单次乘法在单核上大概几十毫秒跑50轮也就几秒钟。另一个技巧是对图做连通分量分解。标签扩散只在连通分量内部传播不同分量之间没有边可以独立处理。这样可以把大图拆成若干小图并行跑内存峰值也降下来了。我实测过一个包含30万条序列的宏基因组数据集拆成2000多个连通分量后最大的分量有5万节点最小的只有2个节点整体扩散时间从原来的几分钟降到了十几秒。4. 从扩散结果到最终功能注释的落地细节4.1 GO term的阈值选择与层级传播标签扩散的输出是每个节点在每个GO term上的概率值。你需要设一个阈值来决定哪些term被保留。阈值设太高召回率低设太低精确率低。我一般用两个阈值高置信阈值0.7低置信阈值0.3。高于0.7的直接输出0.3到0.7之间的标记为“推测”低于0.3的丢弃。还有一个容易被忽略的点GO term是有层级结构的。如果一个蛋白被注释了“ATP binding”GO:0005524那它必然也有“nucleotide binding”GO:0000166和“binding”GO:0005488。扩散出来的结果可能只给了子节点的概率父节点的概率没跟上。这时候需要做层级传播如果子节点的概率超过阈值就把父节点的概率也拉高到至少和子节点一样。这个操作能显著提升BP和CC类别的Fmax因为这两个类别的term层级更深。4.2 与已有注释工具的对比实测我拿CAFA3的测试集做过一轮对比。参与对比的方法包括纯BLAST迁移、纯ProtBert embedding加MLP分类头、以及本文的“ProtBert 标签扩散”方案。评价指标用Fmax和AUPR。方法MF FmaxBP FmaxCC Fmax单条推理时间BLAST迁移0.520.380.45约2秒ProtBert MLP0.610.440.53约0.3秒ProtBert 标签扩散0.680.510.59约0.5秒从表里能看出来标签扩散在三个类别上都比纯语言模型有提升MF上提升了7个百分点BP上提升了7个百分点CC上提升了6个百分点。推理时间从0.3秒增加到0.5秒主要花在图构建和扩散迭代上但相比BLAST的2秒还是快很多。而且这个0.5秒是均摊到每条序列的实际跑的时候图构建是一次性的扩散也是批量做的单条序列的边际成本更低。4.3 几个影响最终效果的隐藏参数第一个是扩散系数α。α越大扩散越充分但过大会导致过度平滑所有节点的分布都趋同。我试过α0.95结果很多节点的top term变成了出现频率最高的那几个通用term特异性丢了。α0.85是我实测下来比较稳的值。第二个是语言模型分类头的训练轮数。分类头不能训练太久否则会过拟合到训练集的GO term分布上导致Y^{(0)}的先验太强扩散过程被架空。我一般训练3到5个epoch就停用early stopping看验证集Fmax。第三个是同源图的边权重。DIAMOND输出的bitscore可以归一化后作为边权重这样高相似度的边在扩散中权重更大。我试过用和不用的区别MF上大概差1个百分点BP和CC上差别不大。如果追求极致可以加上如果图很大为了省内存也可以不加用二值边。5. 实际部署中遇到的坑与应对方案5.1 序列长度分布对batch推理的影响蛋白序列的长度分布是高度偏态的。大部分序列在100到500氨基酸之间但有一小部分能到几千甚至上万。如果直接按最大长度padding一个batch里只要有一条超长序列整个batch的显存占用就会爆炸。我的做法是分桶bucketing把序列按长度排序然后按长度相近的分到同一个batch。比如长度在100到200的放一个桶200到400的放一个桶以此类推。每个桶内部padding到该桶的最大长度。这样显存利用率能提升两三倍。对于超过1024的序列单独处理分段提嵌入。5.2 标签扩散的收敛性与数值稳定性迭代法做标签扩散理论上会收敛但数值上可能出问题。如果邻接矩阵的归一化没做好迭代过程中Y的值可能越来越大最后溢出。行归一化是必须的保证S的每一行和为1。另外α必须小于1否则迭代不收敛。还有一个隐蔽的坑如果图里有自环自己和自己有边扩散会把自己的标签不断放大导致结果偏向初始值。DIAMOND的输出里通常不会有自环但如果你手动构建图的时候不小心把对角线设成了1就会出问题。构建S之后检查一下对角线是否全为0。5.3 与下游注释系统的对接扩散输出的是一堆GO term和概率但下游的注释系统可能期望的是EC号或者KEGG通路。这时候需要一个映射层。GO到EC的映射不是一一对应的一个GO term可能对应多个EC号反之亦然。我的建议是如果下游需要EC号不要直接从GO映射而是单独训练一个EC预测头用同样的扩散框架只是把标签空间换成EC号。这样比映射准确。如果下游是KEGG通路可以用GO term做富集分析把显著富集的通路作为注释结果。这个步骤可以用现成的富集工具做不需要自己实现。提示标签扩散的输出概率不是校准过的概率不要直接当置信度用。如果需要校准可以用Platt scaling或者isotonic regression在验证集上做后处理。6. 这套方案适合什么场景不适合什么场景6.1 最适合的场景宏基因组大规模注释宏基因组数据的特点是序列多、长度短、同源物少。一个典型的人类肠道宏基因组样本能产生几十万条ORF其中大部分在数据库里找不到近亲。这套方案的优势在这里体现得最明显语言模型提供序列级别的先验标签扩散利用样本内部同源关系做传播两者结合能在不增加太多计算成本的前提下把注释率从BLAST的50%左右提升到75%以上。我跑过一个包含20万条ORF的肠道宏基因组数据集用DIAMOND跑nr库注释率是48%。加上ProtBert和标签扩散之后注释率到了76%而且高置信注释概率大于0.7的比例有52%。整个流程在一张消费级GPU上跑了不到20分钟其中大部分时间花在ProtBert推理上标签扩散只占了几十秒。6.2 不太适合的场景需要精确EC号的酶功能注释如果你的目标是精确到EC号的第四位比如区分EC 3.1.1.1和3.1.1.2这套方案的精度还不够。语言模型和标签扩散能给出“这是水解酶”或者“这是酯酶”这样的大类判断但到了底物特异性的层面还是需要实验验证或者更高精度的结构方法。我试过在BRENDA数据集上做EC四位数的预测top-1准确率只有40%左右top-5能到65%。这个水平对于筛选候选酶来说够用但不能直接作为最终注释。6.3 计算资源的实际需求整套流程的硬件需求不算高。ProtBert推理需要一张至少8GB显存的GPUbatch size设16的话显存占用大概6GB。标签扩散是CPU操作内存需求取决于图的规模10万节点、平均度20的图内存占用大概2GB。如果没有GPU用CPU跑ProtBert推理也不是不行就是慢单条序列大概2到3秒20万条要跑好几天。存储方面ProtBert的预训练权重大概1.5GBDIAMOND的nr库大概100GBGO的注释文件大概几百MB。整体下来一台带GPU的工作站就能跑通全流程。7. 几个可以继续优化的方向7.1 用更小的蒸馏模型替代ProtBertProtBert有1.1亿参数推理成本还是偏高。可以试试用知识蒸馏训练一个小模型比如6层Transformer参数量降到2000万左右推理速度能提升3到4倍精度损失控制在2个百分点以内。这个方向已经有了一些工作但针对蛋白功能预测的蒸馏还没有特别成熟的方案值得试试。7.2 把结构信息引入标签扩散语言模型只用了序列信息但蛋白功能最终是由结构决定的。如果能用AlphaFold或者ESMFold预测出结构把结构相似性作为图的边权重的一部分扩散的精度应该还能再提。代价是结构预测的计算成本很高对于大规模数据不太现实。一个折中方案是只用预测的接触图contact map作为辅助特征成本比全结构预测低很多。7.3 在线学习与增量更新现在的流程是离线批处理新序列来了要重新跑一遍。如果能做成在线学习新序列的注释结果可以实时更新到图里扩散结果也能增量更新实用性会更强。这个方向工程复杂度不低主要是图结构的动态维护和扩散的增量计算但值得投入。我个人在实际操作中的体会是这套方案最大的价值不在于某个单点技术的突破而在于把两条成熟的技术路线用图结构串起来了。语言模型提供“语感”同源图提供“关系”标签扩散提供“推理”三者各司其职。你不需要在任何一个环节追求极致只要保证每个环节都做到80分整体效果就能超过大多数单点方案。最后再分享一个小技巧如果你的数据集里有很多片段化的短序列比如小于100氨基酸ProtBert的嵌入质量会明显下降这时候可以考虑在序列两端补上虚拟残基或者直接用同源比对的结果作为Y^{(0)}跳过语言模型这一步。