两年前我第一次拿mRNA序列跑语言模型的时候有个细节让我一直很别扭模型在核苷酸层面玩得挺好但一旦涉及翻译成蛋白质这个功能层效果提升就极其有限。后来我把几何深度学习里那套等变思想搬到序列建模上才发现问题不在模型容量而在表征结构本身。这篇Equi-mRNA的工作本质上是把蛋白质翻译过程当作一种天然的对称性约束用等变编码的方式嵌入mRNA语言模型让模型对读框和密码子的理解从靠数据硬学变成架构上天生带着。这篇文章我会按我的理解把这套方案拆开讲清楚它到底等变的是什么、模型怎么搭、训练和评测怎么做、消融实验里那些反直觉的教训以及如果要复现或迁移到自己项目里哪些坑需要提前绕开。内容偏干货适合已经对RNA语言模型有一定了解、但想深入理解结构设计逻辑的读者。1. 为什么mRNA语言模型总在隔靴搔痒等变编码出场的背景1.1 mRNA建模的核心矛盾序列层与功能层的脱节mRNA的生物学功能高度依赖翻译这一步核糖体读取连续的核苷酸三联体密码子决定氨基酸序列进而决定蛋白质结构与功能。但过去几年大部分mRNA/DNA语言模型做的事本质上和普通文本模型没有区别——把A、U、G、C当成字符用掩码语言建模或者自回归的方式去预测下一个token。这类模型把序列当成一条一维字符串没有把三联体→氨基酸这个生物学约束内化到网络结构里。结果就是模型可能在核苷酸层面的统计规律上学得很好比如某段序列有些保守模体但到了下游任务——比如预测翻译效率、mRNA稳定性、密码子最优性——性能常常卡在一个瓶颈。原因并不难理解这些任务本质上都和翻译过程强相关而模型如果没有对读框结构足够敏感的归纳偏置它就要花费大量参数和样本量从零去发现原来密码子是三个一组这件事。我当时做过一个小测试拿同一个mRNA序列分别用正常读框和错开一个碱基的读框去跑某个预训练RNA模型发现隐藏层表征几乎没什么变化。这其实是个信号——模型压根不关心读框那它对翻译相关下游任务的理解就只能停留在浅层模式匹配。1.2 从几何深度学习借来的思路与其学不如让它天生如此等变Equivariance这个概念在几何深度学习中已经很成熟。做分子性质预测的人会要求模型对旋转和翻转等变做点云处理的人会要求模型对平移等变。核心思想是一句话如果输入经历某种变换输出应该以可预测的方式跟着变化而不是输出一个被打码之后的固定结果。把它翻译到mRNA场景就是这篇文章的核心直觉。mRNA序列存在一种非常自然的变换——读框移动。当你把序列从第0位、第1位、第2位开始切密码子你会得到三条不同的氨基酸翻译路径生物学中称为三个读框。在正常的编码区里只有一个读框是正确的会得到有功能的蛋白另外两个读框通常会很快遇到终止密码子。但重要的是这三条读框之间不是毫无关系的三条独立序列它们之间差的就是一个碱基的平移。一个已经理解翻译过程的编码器应该能对序列平移一个碱基这种操作给出结构化的、可预期的响应而不是输出某种随机变化。Equi-mRNA的核心主张正是编码器应当对读框平移等变。如果输入序列向右移动一位相当于读框切换模型的某种内部表示也应当跟着轮转这种轮转关系必须在架构上得到保证而不是靠训练时隐式地让模型记住哦平移之后结果差不多。这个归纳偏置的好处是显而易见的模型从第一轮训练步开始就默认理解了密码子和读框的关系不需要大量数据去摸索。1.3 等变与不变的边界为什么不是对读框不敏感第一次听到对读框等变的时候容易产生一个混淆等变是不是就等于不管哪个读框编码结果都差不多如果是这样那模型对读框仍然是盲的只是换个说法。恰恰相反不变Invariance指的是输出完全不随变换改变适合判断图中物体是否存在这类任务而等变要求输出跟着输入的变换以确定性的方式同步改变。举个例子设输入序列为x把序列整体平移一个碱基得到x如果编码器f满足等变性那么f(x)应当等于某个固定置换作用在f(x)上的结果。换句话说模型内部可以区分这是0号读框和这是1号读框并且它们之间的关系是结构化的、已知的。类似CNN中卷积层的平移等变图像平移特征图也平移但特征图的模式没有改变——模型既对位置敏感又对模式稳定。这才是Equi-mRNA的精髓。它不是要让模型忽略读框而是要让模型用一种可预测的方式把三个读框组织在一起。后面我会详细讲具体的组织方式。2. 翻译等变性到底等变什么三联体密码子与三读框的对称性2.1 三读框结构mRNA里天生的循环对称性先回到生物学基础因为对称性从哪来、为什么是三不是别的数字全都由遗传密码的结构决定。任意一条RNA序列都可以从第0、第1、第2个碱基开始按三个一组切分密码子由此得到三条不同的密码子串。注意这三条串之间有着严格的关系读框1的密码子序列等价于把读框0的序列右移一个碱基后重新从起点切分。如果用记号来表示设读框0的密码子序列是C₀[c₀,c₁,...]每个c_i对应位置3i,3i1,3i2的三个碱基读框1的密码子序列则是从位置1开始的三个碱基组成第一个密码子对应[1,2,3]。这三个读框构成了一个天然的循环结构——每一次平移一个碱基就在三个读框之间轮转。这个数学结构非常像循环群Z₃的对称性。三个读框的集合在平移操作下形成一个闭环0→1→2→0。如果编码器把输入序列平移一位那么三个读框的身份就会按这个循环顺序轮换。对语言模型来说这是一个非常干净、非常确定的对称群。2.2 为什么通用Transformer做不到这个等变性一般Transformer没有这个保证原因在位置编码。以绝对位置编码为例序列中第i个位置被叠加了位置向量p_i之后位置的语义就与绝对坐标绑定了。当你输入整体平移一位位置编码不会跟着做循环移动模型在注意力层看到的位置信号就完全不同最终的表示层很难保证只产生按照预测方式移动一下的变化。就算用一个完全随机初始化的Transformer训练很久它顶多能学到某种近似的、不稳定的读框敏感性不可能达到严格等变。也有些模型试图通过把3-mer当成token来绕开这个问题——直接把三个碱基拼成一个符号。这样确实让模型天然按密码子粒度处理但坏处是丢失了碱基粒度的信息比如同义密码子之间那一个碱基的差异可能完全不可见。而且一旦序列对齐方式变了从另一个读框开始切整个token序列就完全不同了模型没有任何机制去感知这些token和之前那条序列其实是同一条RNA的不同读框。Equi-mRNA等变编码方案解决的是既保留碱基级表示又让模型结构上保证平移一个碱基读框循环轮转。2.3 形式化的等变约束编码器与置换表示用更严格一点的语言来描述这个约束。设输入序列为x∈Σ^L其中Σ{A,U,G,C}。定义一个操作g将序列向右平移一个碱基并在左端补上特殊填充。设编码器f在合适的维度上输出一个表示张量该张量可以按三读框维度分解为三个通道或三个子空间。要求存在一个固定的置换矩阵P使得f(g·x) P·f(x)其中P的作用是把三读框通道做循环轮转。也就是说平移一个碱基模型内部三个读框的身份就按预定方式交换。这个约束必须在架构中显式构造而不是加一个正则项去软性鼓励。软性约束的问题是模型总能找到某种欺骗性的trivial解比如把所有通道输出变成同一个向量那就符合了轮转但什么信息都没区分。实际实现时并不是硬套一个置换矩阵在输出上而是用三读框共享编码器循环聚合的方式来实现等变。我下一节详细拆架构。3. Equi-mRNA的模型骨架三读框通道、等变聚合与双头读出3.1 总体结构一个主干三条共享权重的读框流整条模型可以分成三块序列Embedding与相对位置编码、三读框等变编码器、输出侧的双头核苷酸语言模型头翻译表征头。需要先说一个关键决策整个编码器主体是共享权重的。输入序列会用三个偏移量0、1、2各自做一次切分形成三个输入变体。三个变体进入到同一个编码器分支里参数完全一致。绘图表示的话就像一个主干网络被复制成三份但复制份之间共享权重只是输入数据略有不同。那这么做为什么能等变假设原始序列x从第0位开始切编码器分支得到读框0的密码子表示。如果把原始序列整体右移一位那么原来读框0的切法就正好等于新序列读框1的切法。由于三个分支共享权重新序列在分支1中得到的表示和旧序列在分支0中得到的表示是完全一致的。再把各分支的表示做一次循环聚合整体就满足等变关系。这个设计的一个额外好处是参数效率极高三个读框不会增加三倍参数量只增加了激活内存。后面讲训练代价时你会看到光靠这一点就比很多三倍参数硬编码三读框的朴素方案划算。3.2 相对位置编码等变的前提条件如果要让平移一个碱基之后表示按预期轮转绝对位置编码会坏事儿。假设第i个token的编码叠加了绝对位置sinusoidal向量那么右移一位之后同一个碱基的绝对位置就变了共享权重分支之间也没有办法对应上。解决方式是用相对位置编码或者在多头注意力中只使用相对位置偏置。这样token之间的依赖基于它们在序列中的相对距离比如相隔1个碱基相隔18个碱基而不是绝对坐标。平移不改变相对距离的集合所以平移前后模型的内部激活模式可以保持一致。顺带说一个工程细节输入右移时要补什么样的填充token文章实现里与常规语言模型类似用特殊的[MASK]或者[PAD]填充并在注意力掩码里遮掉补丁区域。填充token不会参与梯度计算因为其位置对应的loss被mask掉。这其实和文本模型处理padding的方式如出一辙但这里要确保填充token无论位置在哪都产生相同的表示所以还需要让编码器对填充位置使用独立于位置的Embedding。3.3 等变聚合与读出如何把三个读框拧成统一表征三个共享权重的分支各自输出一组序列表征。但这些表征还是分别对应三个不同切法的输入要进入下游任务必须把它们融合起来。融合方式必须满足等变关系。这里采用的是循环卷积式聚合circular aggregation。具体来说对三个通道的输出向量先做一次1×1卷积相当于一组可学习的线性混合再进行循环均值池化。循环池化保证通道维度上满足Z₃置换等变如果输入通道按顺序轮转输出按相同顺序轮转。模型最终有两个输出头语言模型头对每个核苷酸位置预测被mask掉的token。这个头和普通BERT式预训练完全一样保证模型学到核苷酸层面的语境信息。翻译表征头把编码器输出线性投影到20维氨基酸类别空间IUPAC单字母码用于监督模型在每个位置上能否读出对应氨基酸。这个头会在预训练阶段引入一个辅助loss引导中间表征携带翻译信息并不强制模型显式输出完整蛋白序列——它更像一个多任务学习的辅助监督信号。把这两个头拼在一起的设计让模型同一套参数既能处理序列语法核苷酸之间的依赖又能处理序列语义翻译成蛋白质之后的氨基酸模式而且因为共享主干信息不是简单拼接而是深度融合。3.4 简单类比这就像一个三视角卷积如果觉得上面的结构有点抽象可以做个类比。图像领域里如果任务要求模型对旋转90度等变一种方法是在不同旋转角度下把同一张图输入共享参数的CNN然后对输出做对称化聚合。Equi-mRNA做的事情一模一样三个读框就是三个视角平移一个碱基就是旋转一次视角模型用共享权重的编码器分别处理三个视角最后用结构化的聚合保持旋转即轮转的对称性。这样设计的好处从直觉上就清晰了模型对读框的组织不是靠统计偶然碰出来的而是从架构上就是三个对称的分支。所以哪怕预训练语料里没有刻意标注哪个是编码区、哪个读框正确模型也能自发把三读框信息当成一个完整结构来表征。4. 预训练与下游评测多物种转录本上的效果验证4.1 预训练数据与训练细节预训练语料来自某个公开的mRNA转录本数据库子集包含多种真核物种共约200万条转录本序列。数据预处理时有几个重要步骤把所有序列统一转为RNA字母表T替换为U去掉长度过短150nt的序列并按相似度去重以防止某个高表达转录本在语料中过度重复。预训练任务是经典的掩码语言建模掩码率设为15%其中80%替换为[MASK]、10%替换为随机核苷酸、10%保持原样这个配置和BERT一脉相承。在该基础上辅助的翻译头loss按0.1的权重加进总loss里让模型在预训练阶段就接收氨基酸层面的监督信号。训练使用AdamW学习率2e-4warmup 2000步余弦退火到1e-5batch size设为128条序列每条长度截断至1024个核苷酸。模型是一个12层Transformerhidden size 38412个注意力头参数总量约4600万。共享三分支实现后参数量没有增加但因为一组输入在三读框下要各前向一次激活显存大约是单分支模型的3倍。训练共走了约30万步在4张卡上大约跑了6天。因为共享权重的三路前向可以合并成一个batch实际load balance还算容易训练初期loss下降比普通单读框Transformer明显更快这说明架构带来的归纳偏置在早期训练中有真实收益。4.2 下游任务设置与对比基线评测选了四个方向完全不同的任务确保不是只在某一个任务上刷分翻译效率预测回归给定5UTR和编码区起始段预测核糖体加载密度。指标用决定系数R²。mRNA稳定性预测回归预测转录本半衰期。指标用Spearman相关系数。密码子最优性分类二分类判断一段编码序列的密码子使用是否贴近基因组高表达基因的偏好。指标用AUC。内部核糖体进入位点IRES识别二分类区分序列是否包含IRES元件。指标用AUC。基线选择也很有代表性一个是不含三读框设计、直接对原始核苷酸序列做masked LM的Transformer名字就略过了一个是基于3-mer分词的同规模模型还有一个是先翻译成蛋白序列再用蛋白语言模型编码最后把核苷酸特征和蛋白特征简单拼接的级联方案。结果如表所示模型参数量TE预测R²稳定性Spearman密码子最优性AUCIRES识别AUC核苷酸级Transformer基线45M0.510.420.780.803-mer分词基线48M0.550.450.810.79级联翻译蛋白模型120M0.580.500.820.83Equi-mRNA单读框微调46M0.610.550.860.85Equi-mRNA三读框融合微调46M0.640.530.880.86最值得注意的不是Equi-mRNA全面领先而是它参数量只有级联方案的1/3不到效果反而更好。级联方案的问题在于翻译→蛋白编码这条链路是有损的蛋白语言模型只看到氨基酸序列完全丢失了同义密码子层面的信息而恰恰是密码子偏好对翻译效率和稳定性非常重要。另一个值得注意的点是三读框融合微调在稳定性预测上略逊于单读框微调。这个现象在消融里给了我们一个很有意思的教训后面专门展开。4.3 等变表征带来的直接收益读框识别能力做一个直观的探测实验把编码器输出的三读框通道向量拿出来训练一个线性分类器判断当前读取的是哪个读框准确率几乎是100%。换个等价说法模型内部表示天然编码了读框身份。对比之下普通Transformer基线的读框探测准确率只有62%——比随机33%高一些但远算不上可靠。这说明等变架构确实做到了设计目标模型从第一轮训练就清楚知道自己正在读哪个框并且三个读框之间的位移关系也是结构化的。这为什么重要因为很多实际任务与读框密切相关。比如预测一个序列变异是否会导致移码突变模型需要精确知道当前读框的边界又比如设计一段mRNA疫苗序列时需要保证优化的密码子不会意外引入提前终止密码子。这类任务中模型对读框的敏感性直接决定可行性。5. 消融实验里的意外收获等变约束什么时候失效5.1 去掉等变聚合会发生什么第一个消融是把等变聚合换成普通的自注意力跨通道融合。也就是说三个读框分支之间的信息交互不再用循环池化而是用一层Transformer把三个通道当成三个token做注意力。从参数量角度看这个替代方案参数量更多理论上表达力只强不弱。但实验结果很有意思翻译效率R²从0.64掉到0.57密码子最优性AUC从0.88掉到0.84。原因在于普通自注意力不保证循环对称性。跨通道注意力学到的关系可能依赖具体位置、具体batch是一种模型自己发挥的融合但等变聚合把所有通道间关系锁死在循环置换结构里让模型在通道维度不必浪费能力去学怎么把三者对齐。这就像让一个从未学过乘法表的人做三位数乘法他最终也许能做对但效率和稳定性远不如那个死记硬背过乘法表的人。不能自由发挥在某些时候恰恰是最大的优势。5.2 共享权重的意义参数效率之外还有正则化效果第二组消融对比了三读框共享权重和三读框各用独立编码器两种实现。独立编码器方案的参数量直接乘三从46M涨到138M左右。参数量大了但结果并没有变好——几个任务各有高低但总体没有超出共享权重版本。独立编码器真正的问题是三个分支各自在自己的参数空间里发展出一套表示可能语义对齐程度很低。虽然最终可以通过聚合层强行对齐但学习难度大得多训练更不稳定下游微调收敛也慢。共享权重本质上是一种强正则三个分支被迫用同一套规则理解三个读框这强迫模型学到读框之间的共同结构而不是各行其是。增加的诱导偏差显著压过了参数增加带来的表达力优势。5.3 等变不是万能的两个反例更大价值来自两个反直觉的发现。一是对长度很短的片段小于200个核苷酸比如很多RNA调控元件等变设计的优势几乎消失某些配置下甚至小幅低于普通基线。原因比较合理短序列里读框本来就不完整三条通道之间的循环对应关系被截断破坏等变结构能依赖的平移对称性无法充分展开。这时候模型只能靠有限上下文硬猜和普通Transformer没有本质差别。二是对某些存在复杂剪接的区域强等变反而会引入偏差。剪接会导致同一个转录本前后两段使用不同读框整条序列的读框身份并不是全局统一的。等变编码把全序列的读框关系限制成严格的全局循环对称遇到前后读框不一致的可变剪接区域时模型的表征空间会有一个内在的衔接成本。也就是说全局等变在普通编码区是优点在剪接异构区域却可能变成刚性约束。这也解释了第二组实验里三读框融合微调在稳定性预测上略逊于单读框的现象——稳定性涉及大量UTR区域和部分剪接调控元件全局等变带来的读框刚性在这些区域产生轻微的负面影响。6. 复现与迁移的经验细节训练稳定性与计算成本控制6.1 显存优化三路前向的实际代价与对策前面说过共享权重不增加参数量但三读框同时前向会让激活显存翻三倍。这对小团队来说是不小的压力。实际操作时我建议开gradient checkpointing把激活显存降回去大部分代价是训练时间增加约20%但相比无脑占满显存还是划算得多。更激进的方法是交错执行三路前向而不是同时执行。也就是说同一batch先用读框0过一遍、释放中间激活再用读框1过一遍最后用读框2过一遍只在等变聚合时把三个输出拼起来。这样一个batch的内存占用基本等于普通Transformer只是慢三倍。如果你卡得厉害可以采用这个策略把模型跑起来再判断是否值得升级硬件。6.2 预训练热启动从已有核苷酸模型初始化如果从零开始预训练收敛速度其实不快因为三读框的辅助翻译loss要同时带动三个分支。我试过一个很有效的技巧用已有的单读框普通RNA语言模型权重初始化共享编码器主干然后只随机初始化等变聚合部分。这样可以保留模型在核苷酸层面已经学到的通用语境知识再通过微调阶段把三读框结构叠上去。实测下来预训练步数可以减少约40%最终下游指标基本持平。类似技巧在迁移到特定物种或组织类型时也很有用——先在人类转录本上完成通用预训练再用目标物种数据做几轮领域自适应比从一开始就在目标数据上训练更稳。6.3 微调阶段的选择单通道还是三通道下游微调时一个值得纠结的问题是用融合后的三通道表征还是直接用0号读框分支的表征我们的实验发现对整体编码区内的任务翻译效率、密码子最优性三通道融合略好对UTR为主或涉及剪接调控的任务稳定性预测单通道分支反而略好。一个稳妥的做法是微调时同时保留两个出口主出口用三通道融合输出另加一个只读0号读框分支的辅助分类出口。推理时按任务类型选择哪个出口代价极小效果上两全其美。6.4 推理时的轻量化技巧如果要把模型部署到大批量RNA序列筛选场景三路前向耗时是绕不开的问题。但有一个取舍技巧推理时如果下游任务只关心主读框比如判断一段序列的编码潜力可以只跑0号分支跳过另外两路前向和等变聚合特征质量下降非常小。这是因为预训练时三读框共享权重0号分支本身已经是一个完整的序列编码器等变聚合对它的影响主要体现在通道间的上下文而这个上下文对单序列分类任务的边际增益有限。用这个方法可以把推理吞吐提升约2.5倍对需要扫描全转录组的应用场景帮助很大。最后分享一个小技巧等变编码这类方案真正落地时的价值往往不在刷榜而在下游任务的稳定性和可解释性。我这边的经验是在做等变设计之前先花时间想清楚任务里到底存在什么对称性。不是所有生物学过程都有三读框这样干净的结构强行套等变反而可能给模型加一个不必要的紧箍咒。Equi-mRNA值得借鉴的地方在于它找对了mRNA建模中一个无处不在、结构清晰的对称关系并把它贯彻到了架构层面——而不是停在给loss加个约束的程度。如果你打算把这个思路迁移到其他生物序列任务上我的建议是先做小规模探测实验画出平移一个单位之后表征变化是否符合预期置换的对照图再决定要不要投入成本把等变结构做进主干。对称性选对了事半功倍选错了后面每一步都在还债。
