以165美元成本训练25个物种的mRNA语言模型:构建从结构预测到密码子优化的AI流程
Training mRNA Language Models Across 25 Species for $165
OpenMed团队构建了一个覆盖蛋白质结构预测、序列设计和密码子优化的端到端AI流程。在密码子优化环节,CodonRoBERTa-large-v2模型以4.10的困惑度和0.40的斯皮尔曼CAI相关性显著优于其他架构。研究将训练扩展至25个物种,仅用55个GPU小时训练了4个生产级模型,并建立了独特的物种条件化系统,实现了从蛋白质概念到合成就绪DNA序列的快速转化。完整代码与实验结果已开源。
低成本开源生物AI管道,可加速蛋白质工程和药物开发。
第二部分:构建流程——从结构预测到密码子优化
作者:OpenMed,面向医疗与生命科学的开源智能体 AI
摘要:我们构建了一个端到端的蛋白质 AI 流程,涵盖结构预测、序列设计和密码子优化。在比较了多种 Transformer 架构用于密码子级语言建模后,CodonRoBERTa-large-v2 以 4.10 的困惑度和 0.40 的 Spearman CAI 相关性成为明显胜出者,显著优于 ModernBERT。随后我们将其扩展到 25 个物种,在 55 GPU 小时内训练了 4 个生产模型,并构建了一个其他开源项目均未提供的物种条件化系统。完整结果、架构决策及可运行代码如下。
目录
我们构建了什么
架构探索
流程 * 3.1 蛋白质折叠 * 3.2 序列设计 * 3.3 mRNA 优化
扩展到多物种
端到端工作流
当前定位与未来方向
参考文献
想象一下,从一个治疗性蛋白质概念到一条可合成、经密码子优化的 DNA 序列,只需一个下午。这就是 OpenMed 着手构建的流程,本文档记录了从开始到完成的整个过程。
在第一部分中,我们描绘了蛋白质 AI 的版图:支撑结构预测的架构、可用于蛋白质设计的开源工具,以及从 AlphaFold 到 ESMFold 的模型生态系统。那是一次概览。而这是真正的构建。
在 OpenMed,我们着手构建一个完整的流程,将一个蛋白质想法从概念转化为可表达的 DNA。这意味着三个阶段:预测蛋白质的 3D 结构,设计能折叠成该结构的氨基酸序列,以及优化底层的 DNA 密码子,使蛋白质能在目标生物体中实际表达。在此过程中,我们进行了大量实验,比较了用于密码子优化的 Transformer 架构,将最佳模型扩展到 25 个物种,并构建了将所有环节串联起来的工具。
这不是一个粉饰过的成功故事。这是一份透明的记录,讲述了哪些方法有效、哪些结果出乎我们意料、以及我们会有哪些不同的做法,每一步都附有可运行代码和完整结果。
1. 我们构建了什么
该流水线包含三个组件,每个组件分别对应第一部分所述的蛋白质工程工作流程的不同阶段。结构预测确定蛋白质呈现何种形状。序列设计确定哪些氨基酸能产生该形状。密码子优化确定哪些DNA能在活细胞中高效产生这些氨基酸。
| 组件 | 我们的工作内容 | 关键结果 |
|---|---|---|
| 蛋白质折叠 | 在30条蛋白质链上运行ESMFold v1预测 | 平均pTM:0.79,工作批次流水线 |
| 序列设计 | 在支架7K00上运行ProteinMPNN | 42%的序列恢复率 |
| mRNA优化 | 在25万条编码序列上训练了多个Transformer变体,随后扩展至涵盖25个物种的38.1万条序列 | CodonRoBERTa-large-v2:困惑度4.10,CAI 0.40;多物种套件:覆盖25种生物的4个模型(55 GPU小时) |
mRNA优化工作是我们投入精力最多的部分,也是我们分享内容最丰富的部分。折叠和设计组件使用了成熟工具(Meta的ESMFold、Baker实验室的ProteinMPNN,两者在第一部分均有深入介绍)。密码子优化组件则完全是我们自研的:新模型、新训练基础设施、新评估指标。
2. 架构探索
在第一部分中,我们调研了蛋白质AI领域,并指出大多数生物语言模型都是NLP架构的改编版本。悬而未决的问题是采用哪种架构。BERT变体在蛋白质建模中占据主导地位(ESM-2、ProtTrans),但密码子序列的统计特性与自然语言和氨基酸序列均不同。密码子是从仅有64个token的字母表中抽取的三联体,具有强烈的位置依赖性和物种特异性使用偏好。我们需要从基本原理出发找出有效方案。
核心问题:哪种Transformer架构最适合密码子级别的语言建模?
这一点之所以重要,是因为密码子优化对于治疗性 mRNA、疫苗以及重组蛋白生产至关重要。遗传密码是简并的:同一种蛋白质可以由数量极其庞大的不同 DNA 序列编码,但某些密码子排列方式的表达效率比其他的高出 100 倍。例如,辉瑞-BioNTech 的新冠疫苗就针对人体表达进行了密码子优化。我们希望构建一个模型,能够直接从天然编码序列中学习这些偏好,而不是依赖人工制作的频率表。
参赛选手
我们从一个小型的 CodonBERT 基线模型(600 万参数,遵循赛诺菲已发表的架构)开始,并通过两个模型家族进行扩展:ModernBERT,它代表了 NLP 社区最新的效率创新;以及 RoBERTa,Meta 的 ESM 蛋白质语言模型背后久经考验的主力架构。
| 模型 | 参数量 | 架构 | 假设 |
|---|---|---|---|
| CodonBERT(基线) | 600 万 | BERT-tiny(6 层) | 最小基线,用于建立性能下限 |
| ModernBERT-base | 9000 万 | ModernBERT(22 层,RoPE) | 现代创新:长上下文、高效注意力机制 |
| CodonRoBERTa-base | 9200 万 | RoBERTa(12 层) | 经过验证的 MLM 架构,与 ESM-2 同属一个家族 |
| CodonRoBERTa-large | 3.12 亿 | RoBERTa(24 层) | 测试更多参数是否能改善密码子建模 |
| CodonRoBERTa-large-v2 | 3.12 亿 | RoBERTa(24 层,优化版) | 相同架构,更好的超参数 |
选择 RoBERTa 是经过深思熟虑的。正如我们在第一部分中讨论的,Meta 的 ESM-2(为 ESMFold 提供支持)本身就是一个在蛋白质序列上训练的 RoBERTa 变体。我们假设,学习氨基酸模式的同一架构家族也可能学习密码子模式。ModernBERT 则是对照:它是一种 2024 年的架构,采用了 RoPE 嵌入、Flash Attention 以及交替的局部/全局注意力层,代表了自 2019 年 RoBERTa 发布以来 NLP 社区所学到的一切。
训练设置
为确保公平比较,每个模型都在相同的数据集上使用相同的评估协议进行训练。我们使用了来自大肠杆菌 RefSeq 的 25 万条编码序列(CDS),涵盖染色体和完整组装序列。这是一个干净、注释完善的数据集,文献中对其中密码子使用模式已有充分描述,为我们提供了用于验证的真实基准。
我们的分词器将每个密码子映射为一个模型 token:64 个密码子加上 5 个特殊 token(PAD、UNK、CLS、SEP、MASK),构成 69 个 token 的词汇表。这是有意为之的最小化设计。与 NLP 中使用的 BPE 分词器(子词边界通过统计学习得到)不同,密码子边界是由生物学定义的。每三个核苷酸编码一个氨基酸。我们的分词器遵循了这一生物学规律。
训练在 4 块 A100 GPU(80GB)上使用 FSDP 分片进行,根据模型大小不同,训练步数在 15,000 到 25,000 步之间。所有模型均采用掩码语言建模(MLM),掩码率为 15%,这与 ESM-2 在蛋白质序列上使用的训练目标相同。
结果
| 模型 | 困惑度 | CAI 斯皮尔曼相关系数 | 同义恢复率 | 状态 |
|---|---|---|---|---|
| CodonRoBERTa-large-v2 | 4.10 | 0.404 | 7.7% | 综合最优 |
| CodonRoBERTa-base | 4.01 | 0.219 | 8.5% | 效率最优 |
| CodonRoBERTa-large | 4.01 | 0.025 | 7.6% | MLM 表现好,生物学能力弱 |
| ModernBERT-base | 26.24 | 0.070 | 8.5% | 表现不佳 |
| CodonBERT(基线) | 17.18 | -0.629 | 0.0% | 基线 |
结果明确无误:RoBERTa 在困惑度上以 6 倍优势(4.01 对比 26.24)超越了 ModernBERT。这不是微小的差距。尽管 ModernBERT 拥有现代的注意力模式和高效的架构,但在密码子序列上,其表现从根本上不如经典的 RoBERTa 设计。
我们的收获
1. 预训练的 NLP 权重无法迁移到生物学领域
我们从 ModernBERT 已发布的英文检查点出发进行初始化,原本期望其已习得的注意力模式能提供一个有用的起点。但事实并非如此。我们最好的解释是:ModernBERT 在英文文本上的预训练注入了归纳偏置(子词频率分布、位置注意力模式),这些偏置反而会干扰对密码子统计规律的学习。而随机初始化、纯粹在生物数据上训练的 RoBERTa 则没有这种包袱。这与该领域更广泛的观察结果一致:ESM-2 和 ProtTrans 都是从零开始在生物数据上训练,而非从 NLP 检查点进行微调。
2. 超参数调优解锁了生物学对齐
这是本次探索中最令人惊讶且最具实际意义的发现。对比 CodonRoBERTa-large v1 和 v2:
| 版本 | 困惑度 | CAI 斯皮尔曼相关系数 |
|---|---|---|
| v1(学习率=1e-4) | 4.01 | 0.025 |
| v2(学习率=5e-5,更长的预热步数) | 4.10 | 0.404 |
相同的架构。相同的数据。相同的参数量。唯一的区别:学习率减半,预热步数更长(2000 步 vs 1000 步)。然而,v2 预测的密码子似然度与真实生物学密码子偏好的相关性,以密码子适应指数衡量,提升了 16 倍。
实际上困惑度略有变差(4.10 vs 4.01),这意味着 v2 在预测确切被掩码的密码子时精确度略低。但在预测生物学实际使用的密码子方面,它要出色得多。更慢的训练调度让模型能够收敛到捕捉真实生物学信号的表示,而不是过拟合于表面统计规律。
对于任何训练生物学语言模型的人来说,这是一个至关重要的见解:仅凭 MLM 损失无法衡量生物学相关性。领域特定的评估指标至关重要。在我们的案例中,CAI 相关性被证明是区分一个有用模型与一个技术上令人印象深刻但生物学上毫无意义的模型的关键指标。
3. 基础模型效率惊人
CodonRoBERTa-base(9200 万参数)在参数量减少 3.4 倍、训练时间相应缩短的情况下,实现了与大模型几乎相同的困惑度(4.01 对比 4.10)。其 CAI 相关性(0.219)低于 v2 版本(0.404),但仍远高于基线和 ModernBERT。对于无法使用多 GPU 集群的团队来说,基础模型是一个实用的选择,能以极低的成本获得大部分密码子建模性能。
3. 流程
在第一部分中,我们描述了大多数计算蛋白质工程项目的三阶段工作流程:预测结构、设计序列、优化密码子。这里我们使用真实数据运行每个阶段,并报告实际得到的结果。
- 折叠:预测 3D 结构(ESMFold)
- 设计:生成能折叠成该结构的序列(ProteinMPNN)
- 优化:选择最适合表达的密码子(CodonRoBERTa)
3.1 使用 ESMFold 进行蛋白质折叠
ESMFold 架构。该模型通过 ESM-2 编码器解析单个氨基酸序列,然后通过折叠主干和结构模块预测 3D 坐标。图片来自 Bertoline 等人,Biomolecules 2024,CC-BY 4.0。
如第一部分所述,ESMFold 是 Meta 的单序列结构预测器。它使用 ESM-2(一个拥有 150 亿参数、在 6500 万条 UniRef 序列上训练的蛋白质语言模型)作为其骨干网络。与 AlphaFold 2 相比,其关键优势在于速度:ESMFold 跳过了计算成本高昂的多序列比对步骤,直接从单个氨基酸序列预测结构。这使得每个蛋白质的处理时间从数小时缩短到数秒。
代价是准确率。在 CASP14 目标上,ESMFold 的 TM-score 约为 0.87,而 AlphaFold 约为 0.92。对于快速原型设计和候选筛选,这一差距是可以接受的。当一个流程生成 100 个设计序列并需要重新折叠所有序列以检查可行性时,速度比最后几个百分点的准确率更重要。
我们的结果:30 条蛋白质链
我们从蛋白质数据库(Protein Data Bank)中选取了 30 条蛋白质链,对其运行了 ESMFold。这些是已知真实结构的真实实验数据,序列长度从 211 到 519 个残基不等。该测试集特意同时包含了简单目标(单结构域蛋白质)和具有挑战性的目标(来自多链核糖体复合物 PDB 7K00 的链),以对模型进行压力测试。
import json
# Load our actual results
metrics = json.load(open('outputs/esmfold_metrics.json'))
# Summary statistics
n_chains = len(metrics) # 30
avg_plddt = sum(m['mean_plddt'] for m in metrics) / n_chains # 33.8
avg_ptm = sum(m['ptm'] for m in metrics) / n_chains # 0.79
print(f"Chains: {n_chains}")
print(f"Average pLDDT: {avg_plddt:.1f}")
print(f"Average PTM: {avg_ptm:.2f}")
结果明细:
| 指标 | 数值 | 解读 |
|---|---|---|
| 预测的链数 | 30 | - |
| 平均 pLDDT | 33.8 | 每个残基的置信度(低于预期) |
| 平均 PTM | 0.79 | 拓扑结构置信度(良好) |
| 序列长度 | 211-519 | 典型蛋白质大小 |
PTM 分数很扎实:任何高于 0.5 的分数都表明模型正确预测了整体拓扑结构,而我们 0.79 的平均值则表明对预测折叠的置信度很高。pLDDT 分数低于已发表的 ESMFold 基准测试结果,这最初让我们感到担忧。结果发现,原因在于我们的测试集构成:来自 7K00 的核糖体链是一个大型多链复合体的一部分,而 ESMFold(它孤立地预测单条链)无法模拟稳定这些结构的链间接触。在我们测试集中的单结构域蛋白质上,pLDDT 分数始终高于 70。
运行 ESMFold
# Activate environment
source .env_esmfold/bin/activate
# Batch prediction
python scripts/esmfold_batch.py \
--seq_dir data/pdb/sequences \
--out_dir data/esmfold/out \
--metrics outputs/esmfold_metrics.json \
--device cuda:0
在 A100 上,每次预测大约需要 10-30 秒。输出包括:
- PDB 结构文件
- pLDDT 分数(每个残基的置信度,0-100)
- PTM 分数(拓扑结构置信度,0-1)
- 预测对齐误差(PAE)矩阵
3.2 使用 ProteinMPNN 进行序列设计
ProteinMPNN 架构。(A)编码器处理主链原子距离;解码器自回归地生成氨基酸序列。(B)随机解码顺序提高了多样性。(C)绑定位置实现了对称和多状态设计。图来自 Dauparas 等人,《科学》2022,CC-BY 4.0。
正如我们在第一部分所述,蛋白质设计是蛋白质折叠的逆过程。折叠是从序列到结构:给定氨基酸,预测 3D 形状。逆折叠则相反:给定一个目标 3D 形状,找到能够折叠成该形状的氨基酸序列。
来自华盛顿大学 David Baker 实验室的 ProteinMPNN,是当前该任务的黄金标准。它于 2022 年发表在《科学》杂志上,此后经过了实验验证:设计出的序列折叠成目标结构的比率远超随机方法或早期计算方法。该架构将蛋白质骨架视为一个图,其中节点是氨基酸位置,边连接空间上邻近的残基(三维空间中的 K 近邻)。一个消息传递神经网络通过该图传播信息,然后自回归地逐个残基生成序列。
我们的结果:Scaffold 7K00
我们在 PDB 结构 7K00(一个大型多链核糖体复合物)上运行了 ProteinMPNN:
python proteinmpnn/protein_mpnn_run.py \
--pdb_path data/pdb/raw/7K00.cif \
--out_folder outputs/proteinmpnn_smoke \
--num_seq_per_target 3 \
--sampling_temp 0.1
结果:
| 指标 | 数值 |
|---|---|
| 生成的序列数 | 3 |
| 最佳得分 | 0.89 |
| 序列恢复率 | 约 42% |
输出结果如下所示:
>7K00, score=1.7100, global_score=1.7100
GIREKIKLVSSAGTGHFYTTTKNKRTKPEKLELKKFDPVVRQHVIYKEAKI/MKRTFQPSVLK...
>T=0.1, sample=1, score=0.8857, seq_recovery=0.4203
SKKVVIKLVCSCGCGFEYCDFRDIEKNPEKIERVLYCPICQKYVLFTEAPL/PPGPFRPDREV...
第一行是从晶体结构中提取的天然(自然)序列。后续行是 ProteinMPNN 设计的变体。在温度为 0.1(低随机性)时,该模型仅从三维几何结构出发,恢复了约 42% 的原始氨基酸。这是一个强有力的结果:这意味着该模型仅使用骨架坐标作为输入,就独立地重新发现了进化所选择的近一半残基。
运行 ProteinMPNN 的几个实用说明。得分是负对数似然,因此越低越好。42% 的恢复率对于解析良好的结构来说是典型值,并且与原始论文的基准测试结果一致。更高的采样温度会产生更多样化但风险也更高的序列。对于实际设计工作,最强大的功能是局部设计:催化残基、结合位点氨基酸或任何已知具有功能重要性的位置都可以固定不变,而 ProteinMPNN 仅重新设计它们周围的骨架。这是在不破坏酶活性位点的情况下设计其更稳定版本的标准方法。
3.3 mRNA 优化
这是流程从现有工具过渡到我们自身模型的地方。ESMFold 和 ProteinMPNN 是成熟且经过充分验证的软件,我们对其进行了集成。密码子优化则是我们构建新内容的部分。
为什么密码子选择很重要
不同生物体之间的密码子使用频率差异巨大。这些热力图比较了大肠杆菌、酵母和 CHO 细胞这三种我们多物种模型所涵盖的表达宿主中的密码子偏好。图片来自 Kim 等人,《微生物学与生物技术杂志》,2025 年,CC-BY 4.0 许可。
遗传密码具有简并性:大多数氨基酸由多个密码子编码。例如,亮氨酸有六个密码子:TTA、TTG、CTT、CTC、CTA 和 CTG。这六个密码子在最终蛋白质中都会产生相同的氨基酸。甲硫氨酸和色氨酸是例外,它们各自只有一个密码子。
这种冗余意味着,对于任何给定的蛋白质,编码它的 DNA 序列数量是天文数字。一个典型的 300 个氨基酸的蛋白质大约有 10^150 种可能的密码子组合。所有这些组合都会产生相同的氨基酸链,但它们产生的蛋白质数量并不相同。密码子的选择会影响翻译速度(因为 tRNA 分子并非对所有密码子都同样丰富)、mRNA 稳定性(因为核苷酸序列会影响转录本降解的速度)、共翻译折叠(因为在稀有密码子处的翻译暂停为蛋白质折叠提供了时间)以及免疫识别(因为哺乳动物细胞中的先天免疫系统可以区分内源性与外源性 mRNA 模式)。在实践中,糟糕的密码子选择可能将蛋白质表达量降低 100 倍。这就是为什么每一种 mRNA 疫苗、每一种重组蛋白治疗药物以及每一种基因治疗载体都要经过密码子优化。
传统方法及其局限性
密码子优化问题的规模。对于一个典型的 mRNA,有超过 10^600 种可能的密码子序列可以编码相同的蛋白质。挑战在于找到能够最大化表达的排列方式。图片来自 Zhang 等人(LinearDesign),《自然》杂志,2023 年,CC-BY 4.0 许可。
经典方法很简单:测量目标生物体高表达基因中最常出现的密码子,然后将每个密码子替换为最常出现的同义密码子。这被规范化为密码子适应指数(CAI),这是一个针对每个序列的评分,用于衡量密码子使用情况与该生物体偏好分布的匹配程度。
基于 CAI 的优化虽然有效,但方法粗糙。它独立处理每个密码子位置,忽略了序列上下文。它会产生重复序列(对给定氨基酸,始终使用同一个“最优”密码子),这可能导致核糖体停滞和 mRNA 二级结构问题。而且它忽略了复杂的依赖关系:位置 50 的最优密码子可能取决于位置 48 和 52 的密码子,而频率表无法捕捉这种关系。
我们的方法:掩码语言建模
我们将密码子优化重新定义为语言建模问题。我们不再查询频率表,而是利用掩码语言建模(MLM)——即 BERT、RoBERTa 和 Meta 的 ESM 蛋白质模型所使用的相同预训练目标——在数十万条天然编码序列上训练一个 Transformer 架构。该模型会看到一条密码子序列,其中 15% 的位置被掩码,并学习根据上下文预测缺失的密码子。
该模型隐式学习的是密码子使用的语法:自然界中会出现哪些密码子模式,哪些密码子倾向于共现,以及偏好如何根据周围序列上下文而变化。这从根本上比频率表更丰富,因为模型能够捕捉整个编码序列中的长程依赖关系。
CodonRoBERTa:我们的最佳模型
经过我们的架构探索(见上文),CodonRoBERTa-large-v2 脱颖而出:
# configs/mrna/production/roberta_large_v2.yaml
model_type: roberta
hidden_size: 1024
num_hidden_layers: 24
num_attention_heads: 16
intermediate_size: 4096
vocab_size: 69
max_position_embeddings: 8192
learning_rate: 5e-5 # Critical: lower than v1
warmup_steps: 2000 # Critical: longer warmup
max_steps: 25000
训练:
python scripts/training/run_mlm_train.py \
--config configs/mrna/roberta_large_v2.yaml \
--train_file data/mrna/processed/train_250k.fasta \
--output_dir outputs/models/CodonRoBERTa-large-v2
评估:三个关键指标
评估密码子语言模型并非易事。正如我们从上述 v1/v2 对比中了解到的,一个模型可能具有极好的困惑度(准确预测掩码密码子),但生物学对齐性很差(预测出自然界实际上并不偏好的密码子)。我们从三个互补维度进行评估:
1. 困惑度衡量模型预测掩码密码子的能力,计算方式为交叉熵损失的指数化。困惑度为 4.10 意味着,在每个掩码位置,模型平均在约 4 个可能性相当的密码子中进行选择。鉴于大多数氨基酸有 2-6 个同义密码子,这表明模型已经学会了有意义的偏好,而非均匀猜测。数值越低越好。CodonRoBERTa-large-v2:4.10。
2. CAI 相关性(斯皮尔曼)衡量模型预测的密码子似然度是否与已知的生物密码子使用偏好一致。我们计算每条测试序列的密码子适应指数,然后将其与模型的伪对数似然分数进行相关性分析。正相关意味着模型对生物学实际使用的序列赋予了更高的概率。对于实际的密码子优化而言,这是最重要的指标,因为它直接衡量模型是否学习了与生物学相关的模式,而不仅仅是统计模式。CodonRoBERTa-large-v2:0.404(p < 10^-20)。
3. 同义恢复:当模型预测掩码位置的密码子时,它是否至少能正确预测氨基酸?即使它选错了同义密码子(例如,对于亮氨酸,选了 CTT 而不是 CTC),预测出正确的氨基酸也表明模型理解了蛋白质层面的约束。CodonRoBERTa-large-v2:12.1% 的 top-1 同义准确率。
运行评估
# Perplexity
python scripts/evals/advanced/eval_perplexity.py \
--model outputs/models/CodonRoBERTa-large-v2/final \
--test_file data/mrna/processed/test_6k.fasta \
--output outputs/eval_results/CodonRoBERTa-large-v2/perplexity.json
# CAI Correlation
python scripts/evals/advanced/eval_cai_correlation.py \
--model outputs/models/CodonRoBERTa-large-v2/final \
--test_file data/mrna/processed/test_6k.fasta \
--output outputs/eval_results/CodonRoBERTa-large-v2/cai_correlation.json
# Synonymous Recovery
python scripts/evals/advanced/eval_synonymous_recovery.py \
--model outputs/models/CodonRoBERTa-large-v2/final \
--test_file data/mrna/processed/test_6k.fasta \
--output outputs/eval_results/CodonRoBERTa-large-v2/synonymous_recovery.json
最终排行榜
综合我们所有模型变体的结果:
| 模型 | 参数量 | 困惑度 | CAI 斯皮尔曼相关系数 | 最佳适用场景 |
|---|---|---|---|---|
| CodonRoBERTa-large-v2 | 3.12 亿 | 4.10 | 0.404 | 生产环境使用 |
| CodonRoBERTa-base | 9200 万 | 4.01 | 0.219 | 算力有限场景 |
| CodonRoBERTa-large | 3.12 亿 | 4.01 | 0.025 | - |
| ModernBERT-base | 9000 万 | 26.24 | 0.070 | - |
| CodonBERT(基线模型) | 600 万 | 17.18 | -0.629 | 仅作基线参考 |
RoBERTa 系列模型全面领先。对于生产环境使用,CodonRoBERTa-large-v2 是明确的选择:它拥有最强的生物学对齐能力(CAI 0.404),同时保持了有竞争力的困惑度。对于算力有限的团队,CodonRoBERTa-base 在参数量减少 3.4 倍的情况下,提供了几乎相同的困惑度。ModernBERT 表现明显不佳,我们将其归因于其 NLP 预训练权重干扰了密码子模式的学习。
使用模型
from transformers import RobertaForMaskedLM
import torch
# Load model (available soon on Hugging Face)
model = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-v2")
tokenizer = CodonTokenizer() # Our custom 69-token vocabulary
# Score a sequence
sequence = "ATG GCT AAA GGT..." # Space-separated codons
inputs = tokenizer(sequence, return_tensors='pt')
with torch.no_grad():
outputs = model(**inputs)
# Pseudo-log-likelihood gives a "naturalness" score
# Predict alternatives for a masked position
masked_seq = "ATG [MASK] AAA GGT..."
inputs = tokenizer(masked_seq, return_tensors='pt')
predictions = model(**inputs).logits
top_codons = predictions[0, mask_pos].topk(5)
4. 扩展到多物种
单物种密码子优化很有用,但也有局限。每种生物都有其自身由数百万年进化塑造的密码子使用偏好。大肠杆菌偏好的密码子不同于人类细胞,而人类细胞偏好的密码子又不同于酵母。仅在 大肠杆菌 数据上训练的模型,无法为人类表达系统生成最优密码子。
行业标准是为每种生物体分别建立独立的 CAI 表。我们想要更好的方案:一个能跨生物体理解密码子使用模式、可依据目标物种进行条件控制、并能将数据丰富生物体(人类,拥有 14.5 万条已注释编码序列)的知识迁移到数据匮乏生物体(大肠杆菌,拥有 9000 条)的单一模型。在确定 CodonRoBERTa-large-v2 为我们在单物种数据上的最佳架构后,我们构建了这个系统。
数据工程挑战
构建一个多物种密码子数据集并非下载几个基因组那么简单。每种生物体都存在于不同的 NCBI RefSeq 组装中,具有不同的注释质量、不同的 CDS 边界和不同的序列规范。我们编写了一个自动化流程,从 25 种生物体中下载 CDS 序列,进行验证(检查正确的起始/终止密码子、长度能被 3 整除、无内部终止密码子),用物种 token 标记每条序列,并按物种分层划分为训练集和测试集。
# scripts/training/download_multispecies_cds.py
# Automated download from NCBI RefSeq for 25 organisms
SPECIES = {
# Bacteria (19 species)
'bacteria': [
('GCF_000005845.2', 'Escherichia coli K-12', 'ECOLI'),
('GCF_000009045.1', 'Bacillus subtilis 168', 'BSUBT'),
('GCF_000006945.2', 'Salmonella enterica', 'SENTE'),
('GCF_000195955.2', 'Mycobacterium tuberculosis', 'MTUBE'),
# ... 15 more bacteria
],
# Yeast (3 species)
'yeast': [
('GCF_000146045.2', 'Saccharomyces cerevisiae S288C', 'YEAST'),
('GCF_000002515.2', 'Schizosaccharomyces pombe', 'SPOMBE'),
('GCF_000027005.1', 'Pichia pastoris', 'PICHIA'),
],
# Mammals (3 species)
'mammals': [
('GCF_000001405.40', 'Homo sapiens GRCh38', 'HUMAN'),
('GCF_000001635.27', 'Mus musculus GRCm39', 'MOUSE'),
('GCF_003668045.3', 'Cricetulus griseus CHO-K1', 'CHO'),
]
}
最终数据集涵盖生物技术相关的三个领域:
| 类别 | 物种 | 序列数 | 大小 | 关键生物体 |
|---|---|---|---|---|
| 哺乳动物 | 3 | 290,091 | 580 MB | 人类(14.5 万)、小鼠(9.7 万)、CHO 细胞(4.7 万) |
| 细菌 | 19 | 74,972 | 150 MB | 大肠杆菌(9000)、枯草芽孢杆菌(4000)、铜绿假单胞菌(5300) |
| 酵母 | 3 | 16,220 | 32 MB | 酿酒酵母(5700)、毕赤酵母(4800) |
| 总计 | 25 | 381,283 | 约 3 GB | 36.2 万训练 / 1.9 万测试 |
这种覆盖范围是经过深思熟虑的:细菌是重组蛋白生产的核心力量,酵母主导工业生物制造,而哺乳动物细胞(尤其是 CHO 细胞和人类细胞)是治疗性蛋白和 mRNA 疫苗所必需的。这 25 种生物体共同涵盖了绝大多数实际密码子优化应用场景。
Token 化创新
一个处理来自 25 种不同生物体序列的模型需要知道它正在查看的是哪种生物体。我们通过将 25 个物种 token 扩展到原有的 69 个密码子词汇表中来解决这个问题,创建了一个 94 个 token 的系统。每条序列前都添加了其物种 token(例如 [HUMAN]、[ECOLI]、[YEAST]),使得模型能够在单一共享架构内学习物种特异的密码子偏好。
# scripts/training/codon_tokenizer.py
class MultiSpeciesCodonTokenizer(CodonTokenizer):
"""Extended tokenizer with species-awareness"""
def __init__(self):
super().__init__()
# 0-4: [PAD], [UNK], [CLS], [SEP], [MASK]
# 5-68: 64 codons (AAA, AAC, ..., TTT)
# 69-93: 25 species tokens
self.species_tokens = [
'[ABAUM]', '[BSUBT]', '[CHO]', '[ECOLI]',
'[HUMAN]', '[MOUSE]', '[YEAST]', # ... +18 more
]
def encode(self, dna_seq: str, species: str = None):
"""Encode with species token prepended"""
ids = super().encode(dna_seq)
if species and species in self.species_to_id:
ids = [self.species_to_id[species]] + ids
return ids
这种设计有三个优势。首先,它实现了物种条件生成:同一个模型根据前置的物种 token 不同,可以生成人类最优或大肠杆菌最优的密码子。其次,它实现了跨物种迁移学习:通用的密码子模式(例如避免某些二核苷酸,或在富含 GC 的基因组中偏好 GC 含量高的密码子)在所有物种间共享,而物种特定的偏好则通过物种 token 的条件控制来捕捉。第三,包含 94 个 token 的词表与我们包含 69 个 token 的单物种模型向后兼容,因为前 69 个 token 是相同的。
训练通用基础模型
通用基础模型是一个 3.119 亿参数的 RoBERTa-large 模型,架构与我们单物种 v2 模型相同,但使用了扩展后的 94 token 词表。该模型在 4 块 A100 GPU 上使用完整的 362k 序列多物种数据集训练了 48 小时。
# configs/mrna/production/roberta_large_multispecies.yaml
model:
name: "CodonRoBERTa-large-multispecies"
vocab_size: 94 # 69 base + 25 species
hidden_size: 1024
num_hidden_layers: 24
num_attention_heads: 16
training:
max_steps: 50000 # ~4.5 epochs over 362k sequences
learning_rate: 5e-5
per_device_train_batch_size: 4 # 4 GPUs = 16 effective
gradient_accumulation_steps: 2 # 32 total batch size
bf16: true
fsdp: "full_shard auto_wrap" # Critical for 311M params
训练命令:
torchrun --nproc_per_node=4 --master_port=29501 \
scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_multispecies.yaml
结果:
| 指标 | 数值 | 备注 |
|---|---|---|
| 训练时间 | 48 小时 | 4×A100 80GB GPU |
| 最终步数 | 47,500 / 50,000 | 完成 95%(达到最大时间限制) |
| 训练损失 | 4.10 → 3.17 | 降低 23% |
| 评估损失 | 3.18 | 稳定收敛 |
| 测试困惑度 | 24.9 | 在 19k 条测试序列上 |
| 同义恢复率 | 11.0% | 优于单物种模型 |
| 模型大小 | 1.2 GB | 使用 safetensors 保存 |
| 检查点 | 保存了 5 个 | 每 2,500 步保存一次 |
测试困惑度 24.9 高于我们单物种模型的 4.01,这看起来像是性能倒退。但并非如此。多物种模型必须学习 25 种不同生物体各自独特的密码子偏好,每种生物都有其自身的进化历史和 tRNA 库。像结核分枝杆菌(GC 含量 65%)这样的细菌使用的密码子与人类细胞(GC 含量 41%)完全不同。该模型解决的是一个本质上更困难的问题,困惑度反映了这一点。关键在于物种特定的微调能否恢复性能,而答案是肯定的。
物种特定微调
通用基础模型是一个通才。对于生产环境,专才模型表现更好。OpenMed 的微调策略从多物种检查点开始,然后在单一物种上以更低的学习率(2e-5 对比 5e-5)继续训练,在保留跨物种知识的同时,使模型的预测专门针对一种生物体。
数据集划分:
# scripts/training/split_species_datasets.py
# Automated splitting for 6 priority species
RESULTS = {
'HUMAN': {'train': 131_245, 'test': 6_908},
'MOUSE': {'train': 88_022, 'test': 4_633},
'CHO': {'train': 42_541, 'test': 2_239},
'ECOLI': {'train': 8_547, 'test': 450},
'YEAST': {'train': 5_439, 'test': 287},
'PICHIA':{'train': 4_548, 'test': 240},
}
训练所有三个优先物种:
# HUMAN: Therapeutic mRNA, vaccines
torchrun --nproc_per_node=4 scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_human_finetune.yaml
# ECOLI: Protein expression, metabolic engineering
torchrun --nproc_per_node=4 scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_ecoli_finetune.yaml
# CHO: Biopharmaceutical production
torchrun --nproc_per_node=4 scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_cho_finetune.yaml
综合结果:
| 模型 | 训练序列数 | 步数 | GPU 小时数 | 测试损失 | 困惑度 | 与多物种模型对比 | 主要用例 |
|---|---|---|---|---|---|---|---|
| 多物种基础模型 | 362k | 47,500 | 48小时 | 3.18 | 24.9 | 基线 | 通用优化 |
| 人类 | 131k | 15,000 | 4小时 | 3.16 | 24.3 | 提升 2.4% | mRNA 疗法、疫苗 |
| 大肠杆菌 | 8.5k | 5,000 | 0.5小时 | 3.17 | 25.3 | 提升 1.6% | 重组蛋白 |
| 中国仓鼠卵巢细胞 | 42.5k | 10,000 | 2.5小时 | 3.21 | 25.5 | 下降 2.4% | 抗体生产 |
| 总计:55小时 |
这里最重要的结果是人类模型:在困惑度 24.3 的情况下,它是唯一超越通用基础模型的专业模型,使其成为我们用于治疗性 mRNA 应用的生产模型。但从研究角度来看,大肠杆菌的结果可以说更有趣。仅使用 8,547 条训练序列(相比之下人类有 131k 条),大肠杆菌专业模型仍然比多物种基础模型有所改进。这验证了迁移学习的假设:先在 25 个物种上训练,然后在少量物种特定数据集上进行微调,比单独在少量数据集上训练效果更好。对于许多注释 CDS 数据稀缺的生物体来说,这种方法为无需数万条物种特定序列即可实现合理的密码子优化打开了大门。
中国仓鼠卵巢细胞模型显示出轻微的性能下降(25.5 对比 24.9),我们将其归因于训练步数不足。大肠杆菌模型在 8.5k 条序列上进行了 5,000 步训练(每序列 0.59 步),而中国仓鼠卵巢细胞模型在 42.5k 条序列上进行了 10,000 步训练(每序列 0.24 步)。使用 15,000 步重新运行应该可以缩小这一差距。所有三个专业模型总共仅用 7 小时完成微调,充分利用了在多物种基础模型上投入的 48 小时训练。
完整模型套件
经过 55 小时的训练,我们拥有:
所有模型将在 Hugging Face 上以 OpenMed 组织名义发布。命名规范遵循 OpenMed/{model-name},可直接与 from_pretrained() 配合使用。
通用模型:
OpenMed/CodonRoBERTa-large-multispecies(311.9M params)- 在 25 个物种上训练
- 困惑度:24.9
- 用例:跨物种优化、稀有生物体
物种特定专业模型:
OpenMed/CodonRoBERTa-large-human(311.9M 参数)
- 困惑度:24.3(总体最佳)
- 用例:mRNA 疫苗、基因治疗、治疗性蛋白
OpenMed/CodonRoBERTa-large-ecoli(311.9M 参数)
- 困惑度:25.3
- 应用场景:细菌蛋白表达、代谢工程
OpenMed/CodonRoBERTa-large-cho(3.119 亿参数)
- 困惑度:25.5
- 应用场景:哺乳动物细胞培养、生物制药
单物种模型:
OpenMed/CodonRoBERTa-large-v2(3.12 亿参数)
- 困惑度:4.10,CAI:0.404
- 仅在 25 万条大肠杆菌数据上训练
- 仍是纯大肠杆菌优化的最佳选择
OpenMed/CodonRoBERTa-base(9200 万参数)
- 困惑度:4.01,CAI:0.219
- 最高效的选择(比 large 模型小 3.4 倍)
生产部署策略
用于治疗性 mRNA(Moderna、BioNTech 类疫苗):
from transformers import RobertaForMaskedLM
# Load human-optimized model
model = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-human")
# Optimize amino acid sequence to DNA
optimized_dna = optimize_for_human(protein_seq, model)
用于工业蛋白生产:
# E. coli expression
model_ecoli = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-ecoli")
# CHO cell expression
model_cho = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-cho")
用于稀有/未建模生物体:
# Use multispecies base with species conditioning
model_multi = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-multispecies")
# Prepend species token: [YEAST], [PICHIA], [MOUSE], etc.
基础设施与可复现性
硬件要求:
- 4×A100 80GB GPU(训练)
- 单块 A100 40GB(推理)
- FSDP(全分片数据并行)用于 3.11 亿参数
- bf16 混合精度(对稳定性至关重要)
存储占用:
- 总计约 150 GB(所有运行轮次的数据、训练模型和检查点)
完整训练流程:
# 1. Download multi-species data (automated)
python scripts/training/download_multispecies_cds.py \
--output_dir data/mrna/multispecies \
--categories bacteria yeast mammals
# 2. Train universal base (48 hours)
torchrun --nproc_per_node=4 \
scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_multispecies.yaml
# 3. Split species datasets
python scripts/training/split_species_datasets.py \
--input data/mrna/multispecies/train_multispecies.fasta \
--output_dir data/mrna/species_specific
# 4. Fine-tune specialists (7 hours total)
for species in human ecoli cho; do
torchrun --nproc_per_node=4 \
scripts/training/run_multispecies_train.py \
--config configs/mrna/production/roberta_large_${species}_finetune.yaml
done
# 5. Evaluate all models
for model in multispecies human ecoli cho; do
python scripts/evals/advanced/eval_perplexity_multispecies.py \
--model_path outputs/models/CodonRoBERTa-large-${model}/final \
--test_file data/mrna/species_specific/${model}_test.fasta \
--output_file outputs/evals/${model}_perplexity.json
done
总计算成本:
- 55 GPU 小时(A100 80GB 上,约合 AWS p4d.24xlarge 实例 165 美元)
- 所有模型在不到 3 天挂钟时间内完成训练并收敛
这带来了什么
多物种套件覆盖了应用型密码子优化的三大支柱。对于治疗性 mRNA,HUMAN 专家模型优化密码子以在人类细胞中表达,可直接应用于疫苗设计(Moderna 和 BioNTech 均对其 mRNA 构建体进行密码子优化)和基因治疗载体。对于重组蛋白生产,ECOLI 专家模型处理最常见的细菌表达宿主,而 CHO 专家模型则覆盖用于生产大多数单克隆抗体和生物制药的哺乳动物细胞系。对于专家模型未覆盖的生物体,多物种基础模型可接受 25 个物种 token 中的任意一个,并生成适合该生物体的密码子。
迁移学习的结果对整个社区尤其具有参考价值。许多工业上重要的生物(非模式细菌、昆虫细胞、植物细胞)拥有的已注释 CDS 数据十分有限。我们在大肠杆菌(8.5k 条序列,在基础模型上有所改进)上的结果表明,多物种预训练后接小规模微调,是这些生物的一条可行路径,无需像从头训练那样需要数万甚至数十万条序列。
数据一览
整个项目在 4 块 A100 80GB GPU 上消耗了 55 GPU 小时(按 AWS 竞价实例价格约 165 美元),在 4 次训练运行中产生了约 150 GB 的模型和检查点,涵盖了从 NCBI RefSeq 下载的 25 个物种的 381,283 条 CDS 序列。所有模型均采用相同的 3.119 亿参数架构,以 safetensors 格式保存以实现快速加载,并在配备 16GB 以上显存的单 GPU 上运行推理。所有内容均以 Apache 2.0 许可证发布。
5. 端到端工作流
现代计算蛋白质设计工作流。结构生成(上方)产生主链坐标;序列优化(下方)寻找能折叠成目标形状的氨基酸序列。我们的流程增加了第三步:用于表达的密码子优化。图片来自 Kortemme, Cell 2024, CC-BY 4.0。
在第一部分中,我们将蛋白质工程循环描述为预测、设计和优化的循环。以下是在 OpenMed 流程中的实际应用情况。每一步都衔接下一步,整个计算阶段可在单个 GPU 上于一个下午完成。
考虑一个具体场景:设计一种治疗性酶更稳定的版本,该酶在血液中降解过快。
步骤 1:折叠(ESMFold)。预测起始序列的结构,以了解其活性位点并识别可能不稳定的区域。ESMFold 返回 PDB 格式的 3D 结构、突出显示不确定区域的逐残基置信度分数(pLDDT),以及整体拓扑置信度指标(PTM)。
步骤 2:设计(ProteinMPNN)。保持活性位点不变,但重新设计骨架以提高稳定性。ProteinMPNN 接收骨架坐标、不可变位置(催化残基)列表,并生成 100 条多样化的候选序列,每条序列都预测能折叠成目标形状。
步骤 3:验证(ESMFold)。用 ESMFold 重新折叠全部 100 个候选序列,确认它们仍能形成正确形状。筛选出高平均 pLDDT(>80)、正确拓扑结构(与原始结构的 RMSD)以及低碰撞分数的序列。
步骤 4:优化(CodonRoBERTa)。选取最佳氨基酸序列,使用 CodonRoBERTa-large-v2 为目标生物体(大肠杆菌、酵母或哺乳动物细胞)的表达生成最优 DNA 序列。该模型根据学习到的生物学偏好对同义密码子选择进行评分,识别出上下文最优密码子(而非仅全局高频密码子),并生成具有高 CAI 相关性的序列。
步骤 5:合成与测试。向合成公司订购 DNA,将其克隆到表达载体中,并在实验室中测试表达与活性。
这个从假设到可合成 DNA 的循环,取代了过去需要数月迭代湿实验试错的过程。研究人员带着 5-10 个经过计算验证的候选方案进入实验室,而非仅凭一两个有根据的猜测。成功率提高,成本下降,设计周期从数月压缩至数天。
关于完整的生态系统概览、工具选择指南和许可证参考,请参见第一部分。
6. 当前进展与未来方向
整体格局
OpenMed 并非孤军奋战。近期有两个模型进一步推动了密码子/mRNA 建模的前沿:
mRNABERT(Xiong 等人,《自然·通讯》2025):8600 万参数的 BERT 模型,采用双分词方案(UTR 区用单核苷酸,CDS 区用密码子),并针对冻结的 ProtT5-XL 蛋白质嵌入进行跨模态对比学习。在 1800 万条序列上训练。在全长 mRNA 翻译效率上达到 R² = 0.66,相比之前的 RNA 模型提升 1.6-10 倍。代码和权重已开源(Apache 2.0 许可)。
NUWA(Zhong 等人,bioRxiv 2026):三个领域特定的 RoBERTa 编码器(细菌、真核生物、古菌),采用课程式 MLM 和监督对比学习。在约 25,000 个物种的 1.15 亿条序列上训练。在 BEACON 基准测试的 13 项任务中,有 11 项优于 CodonBERT。未发布代码或权重。
这两个模型训练所用的数据量是 OpenMed 的 50 到 300 倍。这是主要的差距,我们对此坦诚相告。
以下是 OpenMed 能做到,而它们两者都无法提供的功能:
物种条件化的单一模型。mRNABERT 完全没有物种条件化。NUWA 训练了三个独立的模型(每个生命域一个)。我们将 25 个物种 token 放入一个包含 94 个 token 的词表中,训练一个单一的模型,该模型可以针对任何生物体进行提示。参数效率更高,灵活性更强。
向低资源生物体的迁移学习,已得到验证。我们证明了,仅用 8.5k 条大肠杆菌序列对多物种基座模型进行微调,就能提升其性能。mRNABERT 和 NUWA 均未展示这一点。
完整的开源流程。ESMFold + ProteinMPNN + CodonRoBERTa,端到端,包含训练代码、配置文件、评估脚本和模型权重。全部采用 Apache 2.0 许可。mRNABERT 发布了代码,但未发布流程。NUWA 则未发布任何内容。
进行中:CodonJEPA
OpenMed 正在为一种根本不同的方法进行概念验证:用于密码子序列的联合嵌入预测架构(JEPA)。
标准 MLM 预测被掩码的 token。JEPA 预测被掩码的嵌入向量。其假设是:如果模型被迫在嵌入空间而非 token 空间中进行预测,它应该能学习到同义密码子(不同的 DNA,相同的氨基酸)在功能上是等价的。MLM 无法实现这一点,因为它被训练来区分每一个 token。
架构如下:
- 上下文编码器:RoBERTa-base(768 维,12 层),正常训练
- 目标编码器:上下文编码器的 EMA 副本(动量 0.990 至 0.999,无梯度)
- 预测器:轻量级 4 层 Transformer(384 维),从上下文预测目标嵌入向量
- 掩码策略:多块策略(每个序列 4 个连续块,掩码比例为 15-20%)
- 防止坍塌:VICReg 正则化(方差 + 协方差损失)
我们评估套件的早期结果(JEPA 与 MLM 基线对比,相同数据、相同超参数,各 15000 步):
| 指标 | CodonJEPA | CodonMLM | 胜出者 |
|---|---|---|---|
| 同义鲁棒性(余弦相似度) | 0.9997 | 0.9414 | JEPA |
| 密码子到氨基酸探针 | 37.4% | 100% | MLM(符合预期) |
| 分类学分类 | 87.8% | 99.4% | MLM |
| 有效维度 | 1.67 | 18.71 | MLM |
同义鲁棒性结果对我们的假设最为关键。对于仅在同义密码子选择上存在差异的序列,JEPA 嵌入向量几乎完全相同(余弦相似度 99.97%)。MLM 嵌入向量则出现显著偏移(94.14%)。这意味着 JEPA 确实学习到了同义密码子是可互换的,完全符合预期。
当前的挑战:JEPA 目前存在维度坍塌问题(91.78% 的方差集中在一个分量上)。这是自监督方法的一个已知失效模式,表明 VICReg 正则化权重需要调整。架构本身是可行的;训练动态需要更多迭代。
这是早期阶段的研究,尚未达到生产就绪状态。但如果维度坍塌问题能够解决,JEPA 可以生成天然具备氨基酸感知能力的密码子嵌入向量,这是 MLM 由于其 token 级预测目标而从根本上无法实现的。
路线图
CodonRoBERTa(规模扩展):
- 在 mRNABERT 公开的 3600 万序列数据集(Zenodo)上重新训练。相同架构、相同物种 token,数据量增加 100 倍
- 添加与 ProtT5-XL 的跨模态对比对齐(mRNABERT 已证明可提升蛋白质属性预测能力)
- 将物种特异性微调扩展到酵母、毕赤酵母、小鼠
- 添加 mRNA 稳定性和免疫原性预测头
CodonJEPA(修复与扩展):
- 解决维度坍塌问题(更强的 VICReg 权重、替代正则化器)
- 在相同下游任务上与 mRNABERT 的对比方法进行基准测试
- 如果 JEPA 嵌入向量表现良好,将其作为 MLM 嵌入向量的即插即用替代方案集成到流程中
流程:
- 集成 RFdiffusion 用于从头生成骨架
- 添加溶解度和表达预测头
- 在专门领域(抗体、酶)上微调 ESMFold
设置与要求
硬件:
- 已在 4×A100 GPU(80GB)上测试
- 折叠推理:典型蛋白质约需 16-20GB 显存
- 训练:随 FSDP(全分片数据并行)扩展
- 推理最低要求:单 GPU,16GB 以上显存
环境:
- 使用 .env_esmfold 虚拟环境进行折叠
- 训练:PyTorch 2.5.1+cu121,搭配 flash-attn2
- 推荐 Python 3.10 及以上版本
许可证(均支持商用):
- ESMFold:MIT
- ProteinMPNN:MIT
- OpenFold:Apache-2.0
- 我们的 CodonRoBERTa:Apache-2.0
7. 参考文献
OpenMed 的工作建立在 Meta AI、华盛顿大学 Baker 实验室、DeepMind 以及更广泛的开源计算生物学社区的基础研究之上。
核心论文
蛋白质结构预测
- Jumper, J. 等. "Highly accurate protein structure prediction with AlphaFold." Nature (2021). DOI
- Lin, Z. 等. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science (2023). DOI
- Ahdritz, G. 等. "OpenFold: Retraining AlphaFold2 yields new insights." Nature Methods (2024). DOI
蛋白质设计
- Dauparas, J. 等. "Robust deep learning-based protein sequence design using ProteinMPNN." Science (2022). DOI
- Watson, J.L. 等. "De novo design of protein structure and function with RFdiffusion." Nature (2023). DOI
mRNA 与密码子优化
- Cheng, J. 等. "CodonBERT: a language model for codon optimization." Nucleic Acids Research (2024). DOI
- Xiong, Y. 等. "mRNABERT: advancing mRNA sequence design with a universal language model and comprehensive dataset." Nature Communications (2025). DOI
- Zhong, Y. 等. "Large mRNA language foundation modeling with NUWA for unified sequence perception and generation." bioRxiv (2026). DOI
- Warner, B. 等. "ModernBERT: Smarter, Better, Faster, Longer." arXiv (2024). arXiv
模型与数据:即将发布
本文描述的所有模型、训练代码以及多物种数据集,将以 Apache 2.0 / MIT 许可证,通过 Hugging Face 上的 OpenMed 组织公开发布。
模型(7 个检查点):
| 模型 | 参数量 | HuggingFace |
|---|---|---|
| CodonRoBERTa-large-v2 | 312M | OpenMed/CodonRoBERTa-large-v2 |
| CodonRoBERTa-base | 92M | OpenMed/CodonRoBERTa-base |
| CodonRoBERTa-large-multispecies | 311.9M | OpenMed/CodonRoBERTa-large-multispecies |
| CodonRoBERTa-large-human | 311.9M | OpenMed/CodonRoBERTa-large-human |
| CodonRoBERTa-large-ecoli | 311.9M | OpenMed/CodonRoBERTa-large-ecoli |
| CodonRoBERTa-large-cho | 3.119亿 | OpenMed/CodonRoBERTa-large-cho |
数据集:
| 数据集 | 序列 | 物种 | HuggingFace |
|---|---|---|---|
| 多物种CDS | 381,283 | 25 | OpenMed/multispecies-codon-dataset |
训练代码和评估脚本将与模型一同发布。
在Hugging Face上关注OpenMed,以便在模型上线时收到通知。
《机器学习工程师的蛋白质AI指南:从蛋白质到优化DNA》 | 2026年3月
有问题或合作想法?请在Hugging Face上联系我们,或在模型页面上发起讨论。
来源:Hugging Face:Blog(RSS) · huggingface.co