当RNA也需要“长文本理解能力“，一个16亿参数模型是怎么做到的

你可能没意识到一个很尴尬的事实现在市面上那些号称能读懂RNA的AI模型很多其实连一条完整的信使RNA都读不完。这不是夸张。一条典型的成熟mRNA动辄几千个核苷酸长。而目前主流的RNA基础模型训练时用的上下文长度大多只有1024个token左右。这意味着什么意味着一条2000多个核苷酸的mRNA塞进模型里模型只能看到前面一截后面的内容要么被砍掉要么被硬塞进一个从没见过这么长输入的模型里,勉强运行。这就像让一个只读过短篇小说的人去理解长篇小说的完整情节你把书撕掉一半给他看或者硬塞给他一本没删减的书但他脑子里的阅读习惯根本没为这个长度做过准备。他当然能读出点什么但你不会相信他真正理解了故事的因果脉络。RIBOSPAN这篇论文要解决的就是这件事怎么让一个AI模型从头开始就是按照读完整本书的标准训练出来的而不是读完短篇之后被迫拉伸去应付长篇。为什么这事这么难三个条件同时满足几乎没人做到先说清楚mRNA这东西为什么重要。信使RNA简单说就是细胞里负责把DNA里的遗传信息翻译成蛋白质制造指令的分子。一条完整的mRNA分成三段5非翻译区负责调控翻译怎么开始、编码区真正决定蛋白质长什么样的部分、3非翻译区影响RNA稳定性和降解速度。这三段看起来各管各的但实际上互相牵制得很厉害。编码区里的一个同义突变不改变蛋白质序列只是换了个写法可能会改变RNA的折叠结构进而影响到看似无关的3区域的功能。这也是为什么mRNA疫苗和mRNA药物设计特别麻烦你不能只优化一个局部指标,必须考虑整条链的相互作用。这就带来一个核心矛盾要理解mRNA的功能模型必须能同时看到整条转录本的信息,而不是被切成一段一段分别处理。现在的RNA基础模型基本分两条路。第一条路是像RNA-FM、RiNALMo、AIDO.RNA这样的双向编码器,用的是双向自注意力机制。 双向自注意力一种让模型在理解某个位置的信息时能同时参考它前面和后面所有内容的机制而不是像读书一样只能往前看。这类模型的问题在于上下文长度普遍卡在1024个token根本装不下一条完整的长mRNA。第二条路是像EVA这样的解码器模型用因果注意力靠自回归的方式一个一个生成token能处理更长的序列适合生成任务。但因果注意力有个硬伤每个位置只能看到它前面的内容看不到后面。这对生成没问题但对理解RNA结构和功能就很吃亏,因为RNA的结构折叠、蛋白结合往往涉及序列中相距很远的两端同时参与,单向信息流天然处理不好这种双向依赖。第三条路是一些为了塞下长序列而牺牲某些东西的高效架构。比如HydraRNA主要靠状态空间模块而不是纯注意力RNAret用线性复杂度的双向保持机制BiRNA-BERT用字节对编码把好几个核苷酸压缩成一个token。这些设计确实能处理更长的输入但代价是要么不再是每一层都做全局的两两交互要么牺牲了单核苷酸级别的精度。所以论文提出的核心诉求是能不能同时做到三件事单核苷酸级别的精确表示、每一层都是稠密的双向自注意力、模型规模足够大且原生支持完整转录本长度的上下文。这三件事以前没人同时做到过。RIBOSPAN想成为第一个。模型本身16.1亿参数原生一万个token的上下文先看硬件配置。RIBOSPAN是一个32层的Transformer编码器隐藏维度2048总参数量16.1亿原生预训练上下文长度是10240个token。用的是单核苷酸级别的分词方式也就是每一个A、C、G、U在处理时统一转成T各占一个token位置不做任何压缩合并。 单核苷酸分词把序列里的每一个碱基单独当作一个处理单元不像某些方法把连续几个碱基打包成一个token。这样做的好处是位置分辨率精确到每一个碱基代价是序列变长计算量变大。之所以坚持单核苷酸分词是因为RNA里很多功能性的位点变化就是差一个碱基的事。如果分词的时候把几个碱基捏成一个token相当于把显微镜的分辨率主动降低了,你可能压根看不清那个决定性的单点突变。这就好比检测文件里的错别字如果你把检测粒度设成每五个字看一次那单字的错误你根本发现不了。模型架构上用了RoPE位置编码和SwiGLU前馈网络都是当前大语言模型里比较主流的设计选择这里不算创新点但保证了架构本身足够扎实。真正的关键设计是注意力隔离的序列打包。因为RNA序列长度差异极大从几十个碱基的小RNA到上万个碱基的长转录本都有如果每次训练都按最长序列的长度去处理会浪费大量计算资源在填充的空白token上。RIBOSPAN的做法是把多条不同长度的RNA打包进同一个训练序列里但同时用注意力掩码保证每条RNA的注意力计算只在自己内部进行不会看到隔壁那条毫不相关的RNA的信息位置编码也在每条序列的边界重新归零。这就好比把好几个不同长度的包裹塞进同一辆货车里运输但每个包裹之间用隔板严格分开,货车虽然装得更满更省油但绝不会让A包裹的内容渗透进B包裹里。如果不做这个隔离模型很可能学到序列A后面接着序列B这种完全是训练时拼接产生的假规律而这在真实生物学里毫无意义。训练数据6760万条序列857亿个核苷酸token数据来源两个RNAcentral提供覆盖各种RNA类别的广泛数据和Ensembl提供带有完整CDS和UTR标注的高质量蛋白编码转录本。经过过滤、标准化、精确去重之后最终训练集是6760万条RNA序列857亿个核苷酸token。数据里RNA类别的分布很不均衡rRNA核糖体RNA占了三千多万条是最大的一类而piRNA只有二十万条左右。为了不让评估结果被这种不均衡带偏验证集和测试集用了按类别配额抽样的方式每类RNA至少保证一定数量的样本蛋白编码转录本还按物种做了进一层的分层抽样。训练策略先15%掩码再40%掩码的持续预训练训练用的是掩码语言建模也就是把序列里一部分核苷酸随机遮住,让模型根据上下文去猜被遮住的是什么。 掩码语言建模MLM训练时随机盖住输入序列的一部分内容让模型学习用剩下的上下文去还原被盖住的部分。这是BERT系列模型最经典的自监督训练方式。RIBOSPAN先用15%的掩码率训练一轮然后在这个基础上继续训练把掩码率提高到40%做持续预训练。为什么要这么折腾因为15%的掩码率是标准设置模型看到的大部分上下文都是完整的比较容易靠局部信息猜出答案。而40%的掩码率相当于把将近一半的信息都藏起来模型必须调动更远距离的上下文信息才能补全,这种重度破坏的训练场景更接近实际应用里从残缺信息还原设计意图的需求比如序列优化和重新设计任务。这就像学开车,平时天气好的时候练车遇到暴雨大雾的极端天气反而不知道怎么开。如果只用15%掩码训练模型可能面对轻度信息缺失应付得很好但一旦遇到大段序列被遮盖比如设计任务里大幅度改写一段CDS就会束手无策。40%掩码的持续训练就是专门给模型补上这种极端天气驾驶的经验。论文同时也训练了两个1024长度的基线模型作对照一个是原生1024上下文一个是同样做了40%掩码持续训练的1024版本用来验证长上下文到底带来了多少实际收益。四个模型变体列在下表| 模型变体 | 原生上下文长度 | 掩码率 | 训练阶段 ||---|---|---|---|| RIBOSPAN-10K-15 | 10,240 | 15% | 预训练 || RIBOSPAN-10K-40 | 10,240 | 40% | 在10K-15基础上持续预训练 || RIBOSPAN-1K-15 | 1,024 | 15% | 预训练 || RIBOSPAN-1K-40 | 1,024 | 40% | 在1K-15基础上持续预训练 |重建实验短上下文模型硬拉到长文本会有多惨第一个核心实验是看模型能不能把被遮盖的核苷酸猜对。用了两个指标一个是掩码位置上的重建损失数值越低越好另一个是全局重建准确率数值越高越好衡量的是模型在整条序列所有有效位置上做argmax预测能正确恢复原始核苷酸的比例。结果非常清楚。当输入长度是1024个token的时候各个模型表现都还不错差距不大。但一旦把输入长度拉到10240个token短上下文模型无论是外部参照的AIDO.RNA-CDS还是自家训练的RIBOSPAN-1K系列的表现急剧下滑靠的是直接用RoPE位置编码硬性外推也就是位置编码天生只学过1024长度内的规律硬拉到10倍长度去用效果打折扣是必然的。而原生用10240长度训练的RIBOSPAN-10K系列在10240 token的输入下依然保持稳定的重建表现甚至比在1024长度输入下的表现还要略好一点。具体数字看下表40%掩码10240 token| 模型 | 重建损失 | 全局准确率 ||---|---|---|| AIDO.RNA-CDS | 1.19024 | 0.77791 || RIBOSPAN-1K-15 | 1.27385 | 0.77168 || RIBOSPAN-1K-40 | 1.06730 | 0.80264 || RIBOSPAN-10K-15 | 0.91122 | 0.83769 || **RIBOSPAN-10K-40** | **0.80033** | **0.85887** |这组数字说明两件事原生长上下文训练确实有效硬外推的模型明显吃亏40%掩码的持续训练确实提升了重度破坏场景下的还原能力而且这种提升不是靠牺牲长上下文能力换来的RIBOSPAN-10K-40在两方面都占优。长上下文表征基准光能重建还不够还得看模型有没有理解上下文结构重建准确率高只能说明模型会填空不能证明模型真的理解了序列的上下文组织关系。论文为此专门设计了一个全新的评测基准据其表述这是第一个系统评估RNA基础模型长上下文表征能力的基准。设计思路挺巧妙。取一条完整的mRNA在正中间挖出一小段区域宽度约等于总长度的1/32把这段区域内部的核苷酸顺序打乱但保持成分不变也就是A、C、G、T各有多少个不变只是排列方式变了周围的序列保持原样不动。这样就得到一对原始序列和结构打乱序列。然后看模型的表征也就是每个位置经过模型处理后得到的高维向量在这两种情况下发生了什么变化。具体有三个衡量指标。第一个叫额外上下文分离度ΔCS衡量的是同一种核苷酸比如都是A在被打乱的区域和周围背景区域里模型给它们的表征向量差异是不是变大了。理论上如果模型真的在利用上下文信息被打乱的区域和周围正常区域应该在表征空间里被区分开,数值越高说明模型对局部结构变化越敏感。第二个叫跨区域同碱基相似度Ccross衡量打乱区域和背景区域里同种核苷酸的表征相似程度数值越低说明区域分离得越干净。第三个叫远端表征扩散Ddistal衡量的是这个局部的小改动会不会传染到很远的、根本没被动过的位置上,数值越低说明改动的影响范围控制得越好,只在局部产生变化不会莫名其妙地扩散到几千个碱基之外的地方。这三个指标放在一起看理想的模型应该是ΔCS高对局部变化敏感、Ccross低区域分离干净、Ddistal低影响范围可控不乱扩散。结果发现短上下文模型直接外推到10240长度时ΔCS明显下降Ccross明显上升,也就是模型对局部结构变化的敏感度和区分能力都在退化。用YaRN一种推理时的位置编码缩放技术不需要重新训练就能扩展模型可用的上下文范围能大幅恢复ΔCS和Ccross说明位置编码不匹配确实是退化的主因。 YaRN一种在推理阶段对旋转位置编码RoPE做频率插值调整的技术让原本只训练过短序列的模型能在不重新训练的情况下勉强适应更长的输入序列。但YaRN的代价是Ddistal暴涨。也就是说,用YaRN恢复了模型对局部结构的敏感度但同时让这个局部改动的影响范围失控地扩散到了远处本不该受影响的位置。这就像修好了一个漏水的水龙头但顺带把整间浴室的水管都震裂了。看下面这组10240nt下的对比数据| 模型 | 设置 | ΔCS ↑ | Ccross ↓ | Ddistal ||---|---|---|---|---|| HydraRNA | 直接评估 | 0.313846 | 0.521605 | **0.000667** || AIDO.RNA-CDS | 基础版 | 0.208178 | 0.660208 | 0.004666 || AIDO.RNA-CDS | YaRN | 0.446120 | 0.317370 | 0.025390 || RIBOSPAN-1K-15 | 基础版 | 0.221320 | 0.698757 | 0.006626 || RIBOSPAN-1K-15 | YaRN | **0.451068** | 0.334152 | 0.016084 || RIBOSPAN-10K-15 | 原生10K | 0.405962 | 0.302668 | 0.000785 || RIBOSPAN-10K-40 | 原生10K | 0.405777 | **0.299277** | 0.001164 |HydraRNA作为对照组也提供了一个有意思的信息。它主要靠状态空间模块一种计算效率更高但不是纯粹注意力机制的序列建模方式只有极少数层用了多头注意力。它的Ddistal低到几乎为零说明它把改动的传播范围控制得非常死但ΔCS也明显低于RIBOSPAN-10K系列说明它对局部结构变化的敏感度不够,也就是控制传播的代价是牺牲了灵活捕捉上下文差异的能力。而RIBOSPAN-10K系列无论是15%还是40%掩码版本都是ΔCS较高、Ccross较低、Ddistal也很低的组合,也就是既能敏感捕捉局部结构变化又不会让这种敏感度失控地扩散到无关的远处位置。这里能提炼出一个很有意思的道理位置编码的外推技巧能救回一部分看起来像长上下文理解的表面能力但救不回模型对这个改动应该影响多远这件事的判断力。这种判断力只有让模型在训练阶段就真的经历过完整长序列的输入才能学到。就好比一个只在短跑训练场训练过的运动员你给他吃兴奋剂能让他在马拉松赛道上跑得快一点但他判断该在哪里省力、哪里冲刺的经验这是兴奋剂给不了的只能靠真正跑过马拉松才能积累。RNA类型表征评估不靠下游训练直接看模型脑子里的世界地图画得好不好前面的评测都是在特定任务上验证性能但下游任务的效果好坏一部分归功于预训练模型本身的质量一部分归功于下游微调的技巧,这两者混在一起不容易分开评估纯粹的预训练表征质量。论文因此设计了一个冻结表征评测完全不做任何下游微调直接把预训练模型的输出向量拿出来看这些向量本身有没有把不同类型的RNA自然地聚拢在一起。具体做法是把每条RNA序列所有有效位置的最后一层隐藏状态做平均池化得到一个代表整条序列的向量然后用最近邻分类找10个最相近的邻居看它们的标签是不是一致来评估这个向量空间的组织质量。 最近邻分类一种不需要训练分类器的评估方式直接看某个样本在向量空间里最近的几个邻居是什么类别如果邻居里同类居多说明这个向量空间把同类样本自然地聚在了一起。评测在几个不同的标签体系下进行89955条序列涵盖25种RNA生物类型的总体生物类型评测按功能粗分成看家型调控型编码型三大类的功能评测专门针对7种调控类RNA的评测专门针对长度超过1024nt的长RNA评测以及基于Rfam数据库一个收录RNA家族及其结构同源性的权威数据库的家族评测。结果显示RIBOSPAN-10K系列在总体生物类型评测和功能评测里都拿到了最高的准确率和邻域纯度在长RNA评测里优势更明显,这也符合预期毕竟长上下文原生训练本来就该在长序列上展现更强的表征能力。具体数字10-NN准确率| 模型 | 总体生物类型 | 功能分类 | 调控类型 | 长RNA | Rfam家族 ||---|---|---|---|---|---|| RNA-FM | 0.865033 | 0.954608 | 0.953270 | 0.809924 | **0.990485** || RiNALMo | 0.857395 | 0.944081 | 0.952534 | 0.755867 | 0.985000 || AIDO.RNA-CDS | 0.805125 | 0.954296 | 0.870781 | 0.841973 | 0.952242 || HydraRNA | 0.861264 | 0.976237 | 0.925640 | 0.883771 | 0.987606 || RIBOSPAN-10K-15 | **0.898861** | **0.980621** | **0.959391** | 0.889375 | 0.983818 || RIBOSPAN-10K-40 | 0.898616 | 0.980933 | 0.958087 | **0.891010** | 0.983970 |在Rfam家族评测上RNA-FM表现最强RIBOSPAN紧随其后但没能超越。这提示不同模型在不同任务上各有侧重RIBOSPAN的强势区间集中在长序列和整体类型区分上而在依赖保守序列和结构同源信息的家族识别任务上RNA-FM似乎捕捉到了一些更专门化的模式。值得一提的是两个RIBOSPAN-10K变体之间的差异非常小,说明40%掩码的持续训练虽然在重度破坏场景下提升了重建能力但没有以损害整体表征质量为代价这是一个挺重要的没有免费午餐但也没有额外亏本的验证。建在RIBOSPAN骨架之上一个完整的mRNA生成与重新设计框架有了这个能读懂完整长序列的表征骨架论文接下来做了一件更实际的事:把它改造成一个能主动设计mRNA序列的生成模型。用的是条件化的离散扩散框架。 离散扩散模型一种生成方式训练时先把干净的原始序列逐步腐蚀比如遮盖一部分然后训练模型学习怎么把腐蚀过的序列一步步还原回干净状态生成时就反过来从一个全被遮盖的空白序列开始逐步去噪直到生成一条完整的新序列。和自回归生成像写文章一样一个字接一个字往后写不同扩散模型在每一步去噪的时候是同时更新序列里多个位置的这意味着每一步都能利用序列两侧的完整上下文信息来做决策而不是只能看已经写好的前半部分。这个设计选择跟RIBOSPAN本身的双向架构天然契合。如果用自回归的方式生成那前面提到的双向理解优势就白费了,因为生成过程本身又变成单向的了。用扩散,让5区、编码区、3区的设计可以在同一个统一的表征空间里相互协调而不是先固定一头再往后硬凑。具体实现上扩散过程中加入了时间步和多维度设计条件比如希望序列具备的稳定性、翻译效率等目标属性通过AdaLN-Zero一种条件调制技术让外部条件能够持续地影响模型每一层的表示注入到每个条件扩散模块里。 AdaLN-Zero一种让额外条件信息比如我想要一个翻译效率更高的序列能够动态调节神经网络每一层输出的技术手段最早用在图像生成的扩散模型里这里被搬到RNA序列生成上。特别值得一提的是同义密码子扩散这个设计。因为编码区的序列不是随便改的改动了核苷酸就可能改变对应的氨基酸进而改变蛋白质本身。同义密码子扩散限制候选替换只能在编码同一个氨基酸的不同密码子之间进行这样既能优化编码区在RNA结构、稳定性等方面的性质又能保证最终翻译出来的蛋白质序列完全不变。这就好比给一段文字做同义词替换来优化韵律和节奏但规定每次替换只能换成意思完全一样的词,读者读到的意思一个字没变但整段话的语感对应RNA里就是结构稳定性、翻译效率等可以被悄悄调整。如果不加这个限制模型很可能为了追求某个物理性质指标比如结构更稳定而随意改变编码最后设计出的序列翻译出来的蛋白质根本不是原来想要的那个,这在实际药物设计里是绝对不能接受的事故。论文里也提到这套生成框架之外团队还在开发一个结合强化学习后训练的闭环优化系统用属性预测器给生成的序列打分反馈回训练目标进一步引导生成朝着期望的多维功能特性收敛。不过这部分的完整实验细节论文里说会放在后续的期刊论文里公布。从RIBOSPAN往前看这条路还能走多远这篇论文之后长上下文RNA建模这条线其实还在往前推进。EVA这类因果注意力的解码器模型已经把序列打分、从零生成、局部区域重新设计这些能力扩展到了转录本级别走的是另一条和RIBOSPAN不同但目标相近的路。而HydraRNA、RNAret这些高效架构代表的是在计算效率和表征精度之间做取舍的探索方向,如果未来有办法把稠密双向注意力的表征质量和状态空间模型的计算效率结合到一起可能会催生出下一代真正兼顾精度与规模的RNA基础模型。写在后面读到这篇论文时最让我意外的一点是YaRN带来的副作用。表面上看YaRN是个便宜又好用的技巧,不用重新训练几乎白捡地扩展了模型的可用长度而且确实修复了大部分的表征退化问题。但它把远端扩散指标Ddistal推高了将近10到40倍这个代价在论文表格里安静地躺着很容易被忽略。这提示了一件更普遍的事很多事后修补式的技术方案看起来解决了问题实际上只是把问题从一个可见的维度转移到了另一个不太被关注的维度上。就像给一辆超载的卡车换更硬的减震器表面上开起来更稳了但轮胎和刹车系统承受的压力反而更大了只是你平时看不到。另一个值得琢磨的地方是HydraRNA的对照结果。它把远端扩散控制得极好但代价是上下文分离度也随之打了折扣。这不是简单的谁更好谁更差而是揭示了一种真实存在的架构级权衡:注意力越稠密模型越容易学到丰富但也越容易失控的长距离关联架构越紧凑越受限传播范围就越可控但捕捉复杂上下文关系的能力也跟着受限。这个权衡恐怕不会随着算力的增长而自动消失它更像是信息论层面的一个本质矛盾。如果RNA设计未来真的要走向闭环优化生成、打分、反馈、再生成那这种远端传播的可控性恐怕会变得比准确率本身更重要,因为一旦模型对改一个地方会不会影响八千个碱基之外的功能这件事没有可靠的判断力任何自动化优化系统都有可能在你完全没察觉的情况下把序列改坏。QAQ1RIBOSPAN是什么ARIBOSPAN是一个16.1亿参数的双向RNA基础模型采用单核苷酸级别的分词和稠密双向自注意力机制原生支持长达10240个核苷酸的上下文长度能够对完整的长mRNA转录本进行高分辨率建模同时具备理解RNA表征和生成设计mRNA序列两种能力。Q2RIBOSPAN和普通短上下文RNA模型比优势在哪A普通模型上下文长度大多只有1024个token遇到更长的mRNA只能截断或硬性外推导致重建准确率和长距离表征质量明显下滑。RIBOSPAN从预训练阶段就用10240长度的序列训练在长序列上重建准确率更高同时能更好地把局部结构变化限制在合理范围内不会像外推方案那样让改动效果扩散到无关的远处位置。Q3论文里的同义密码子扩散是干什么用的A这是mRNA生成框架里的一个约束机制限定编码区的替换只能在编码同一氨基酸的不同密码子之间进行这样可以在优化RNA结构稳定性、翻译效率等属性的同时保证最终翻译出来的蛋白质序列完全不变避免设计过程中意外改变蛋白质本身。