OpenMed CodonRoBERTa 多物种 mRNA 语言模型发布
TL;DR – OpenMed 发布了一个端到端的蛋白质工程流水线,该流水线增加了一个由新型转换器语言模型家族驱动的密码子优化阶段,其中 CodonRoBERTa‑large‑v2 达到了 4.10 的困惑度和 0.404 的 CAI 斯皮尔曼相关系数,并且一个覆盖 25 种生物的多物种套件仅用了 55 GPU 小时(约 165 美元)进行了训练。该套件包括一个通用的 311 M 参数模型和三个专门模型(人类、E. coli、CHO),这些模型优于之前的方法,并能够快速生成可表达的 DNA。
1. 发布内容
| 组件 | 描述 | 关键结果 |
|---|---|---|
| 蛋白质折叠 | ESMFold v1 对 30 条蛋白质链的预测 | 平均 PTM = 0.79(高拓扑置信度) |
| 序列设计 | ProteinMPNN 在支架 7K00 上 | 42 % 氨基酸恢复 |
| mRNA 优化 | 新型转换器模型在 250k E. coli CDS 上训练,然后扩展到 25 个物种的 381k CDS | CodonRoBERTa‑large‑v2:困惑度 4.10,CAI 斯皮尔曼相关系数 0.404;覆盖 25 种生物的 4 个生产模型在 55 GPU 小时内训练完成 |
折叠和设计阶段重复使用了已建立的开源工具(ESMFold、ProteinMPNN)。密码子优化阶段是全新的:一套 RoBERTa 风格的语言模型,在密码子级序列上训练并带有物种条件。
2. 架构探索 – 哪种转换器最适合密码子?
| 模型 | 参数 | 架构 | 训练数据(250k E. coli CDS) |
|---|---|---|---|
| CodonBERT (基线) | 6 M | BERT‑tiny(6 层) | – |
| ModernBERT‑base | 90 M | ModernBERT(22 层,RoPE) | – |
| CodonRoBERTa‑base | 92 M | RoBERTa(12 层) | – |
| CodonRoBERTa‑large | 312 M | RoBERTa(24 层) | – |
| CodonRoBERTa‑large‑v2 | 312 M | RoBERTa(24 层,精细化学习率 & 预热) | – |
在 E. coli 测试集上的结果
| 模型 | 困惑度 | 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 % |
| ModernBERT‑base | 26.24 | 0.070 | 8.5 % |
| CodonBERT (基线) | 17.18 | –0.629 | 0 % |
要点
- RoBERTa 在困惑度上比 ModernBERT 好六倍,在 CAI 相关系数上也好六倍,表明经典的 RoBERTa MLM 架构更适合密码子序列的统计结构。
- 用英文文本进行预训练会损害性能 – 从英文检查点初始化 ModernBERT 降低了其学习密码子模式的能力。
- 超参数调优比模型大小更重要 – v2 检查点(学习率 5e‑5,预热 2k 步)尽管困惑度略高,但 CAI 相关系数提升了 16 倍。
- 92 M 参数的基础模型是一个实用的低成本替代方案 – 它与大模型的困惑度相当,但参数仅为其一三。
3. 端到端流水线
3.1 使用 ESMFold 进行蛋白质折叠
- 工具:Meta 的 ESMFold v1(单序列,无 MSA)。
- 在 30 条 PDB 链上的性能:平均 PTM = 0.79,每残基 pLDDT ≈ 34(在多链目标上较低)。折叠在 A100 上每蛋白质耗时 10–30 秒。
3.2 使用 ProteinMPNN 进行序列设计
- 工具:ProteinMPNN(基于图的逆折叠)。
- 在支架 7K00 上的结果:生成了三个序列;最佳序列在温度 0.1 时恢复了 42 % 的原生氨基酸。
3.3 使用 CodonRoBERTa 进行 mRNA 优化
- 问题 – 传统的 CAI 表格独立替换每个密码子,忽略上下文,导致重复且有时次优的序列。
- 解决方案 – 在密码子标记上的掩码语言建模(MLM)学习上下文密码子偏好。模型从周围上下文预测被掩码的密码子,捕获长距离依赖关系。
- 评估指标
- 困惑度 – CodonRoBERTa‑large‑v2 为 4.10(≈ 每个掩码有 4 个等可能的密码子)。
- CAI 斯皮尔曼相关系数 – 0.404(p < 10⁻²⁰),表明与生物学上偏好的密码子用法有很强的一致性。
- 同义恢复 – 12.1 % 的 top‑1 同义预测,表明模型尊重氨基酸约束。
- 使用示例(Python,Hugging Face Transformers):
from transformers import RobertaForMaskedLM model = RobertaForMaskedLM.from_pretrained("OpenMed/CodonRoBERTa-large-v2 tokenizer = CodonTokenizer() # 69‑token 密码子词汇表 seq = "ATG GCT AAA GGT ..." inputs = tokenizer(seq, return_tensors="pt with torch.no_grad(): scores = model(**inputs).logits
4. 扩展到 25 种物种 – 一个了解每种生物的单一模型
4.1 数据工程
- 下载了 25 种生物(19 种细菌,3 种酵母,3 种哺乳动物)的 RefSeq CDS。*
- 验证步骤:起始/终止密码子,长度 % 3 = 0,无内部终止密码子。*
- 最终数据集:381,283 条序列(≈ 3 GB),按每种物种 95 %/5 % 划分训练/测试集。
4.2 物种条件化分词器
- 将 69‑token 密码子词汇表扩展了 25 个特殊物种标记,得到 94‑token 词汇表。*
- 每个序列都带有其物种标记前缀(例如,
[HUMAN]),使模型能够在跨物种共享参数的同时学习物种特异性的密码子偏好。
4.3 通用基础模型
- 架构:RoBERTa‑large(311.9 M 参数),使用 94‑token 词汇表。*
- 训练:在 4 × A100 GPU 上 48 小时,50k 步(≈ 4.5 个 epoch),bf16 混合精度,完全分片数据并行(FSDP)。*
- 测试困惑度:24.9 – 高于单物种模型,因为模型必须捕捉 25 种不同的密码子用法模式。
4.4 物种特定微调
| 物种 | 训练序列 | 微调步骤 | GPU 小时 | 测试困惑度 |
|---|---|---|---|---|
| 人类(治疗用 mRNA) | 131k | 15k | 4h | 24.3 |
| E. coli(细菌表达) | 8.5k | 5k | 0.5h | 25.3 |
| CHO(生物药 mammalian) | 42.5k | 10k | 2.5h | 25.5 |
人类专家模型在使用远少参数的情况下优于通用基础模型(困惑度降低 2.4%),证实了从多物种基础模型进行迁移学习的好处。
4.5 模型套件(已在 Hugging Face 上发布)*
OpenMed/CodonRoBERTa-large-multispecies– 通用 311 M 参数模型。*OpenMed/CodonRoBERTa-large-human– 人类细胞专用模型(总体表现最佳)。*OpenMed/CodonRoBERTa-large-ecoli– E. coli 专用模型。*OpenMed/CodonRoBERTa-large-cho– CHO 细胞专用模型。*OpenMed/CodonRoBERTa-large-v2– 单物种 E. coli 模型,具有最低困惑度(4.10)和最高 CAI 相关系数(0.404)。*OpenMed/CodonRoBERTa-base– 92 M 参数的高效替代模型。*
所有模型均采用 Apache-2.0 许可证,以 safetensors 格式存储,可通过 from_pretrained() 加载。*
5. 端到端工作流示例
- 折叠 – 使用 ESMFold 获得目标蛋白质的 3D 结构。
- 设计 – 运行 ProteinMPNN,固定功能残基,以生成候选氨基酸序列。
- 验证 – 用 ESMFold 重新折叠候选序列;保留具有高 pLDDT/PTM 的序列。
- 优化 – 将选定的氨基酸序列传递给适当的 CodonRoBERTa 专家模型(例如,
CodonRoBERTa-large-human)以获得密码子优化的 DNA 字符串。 - 合成 – 订购 DNA,克隆,并在选定的宿主中测试表达。
整个计算循环在单个 GPU 上运行不到一小时,对于典型蛋白质,将历史上需要数周湿实验迭代的过程压缩了。
6. 在格局中的定位及未来方向
6.1 与当代模型的比较
| 模型 | 参数 | 数据规模 | 物种条件 | 报告的 CAI / 翻译效率指标 |
|---|---|---|---|---|
| OpenMed CodonRoBERTa | 312 M | 381k CDS(25 种物种) | 单一模型带有 25 个标记 | CAI 斯皮尔曼 0.404(人类专家) |
| mRNABERT(Xiong 等,2025) | 86 M | 18M 序列 | 无物种标记 | 在翻译效率上的 R² 0.66 |
| NUWA(Zhong 等,2026) | – | 115M 序列,~25k 物种 | 三个特定领域模型 | 在 11/13 基准上优于 CodonBERT |
OpenMed 的独特贡献
- 一个 单一的、物种条件化的模型,可以针对所覆盖的 25 种生物中的任何一种进行提示,这与 NUWA 的三个独立领域模型不同。*
- 展示了 迁移学习:在通用基础模型上仅用 8.5k E. coli CDS 进行微调即可获得一个专家模型,其性能优于基础模型。*
- 提供了一个 完整、可运行的流水线(折叠 → 设计 → 密码子优化),所有代码和权重均在宽松许可证下发布。*
6.2 进行中的研究 – CodonJEPA
- 目标 – 用 Joint Embedding Predictive Architecture(JEPA)替换 token 级别的 MLM,使同义密码子在嵌入空间中不可区分。*
- 早期结果(相同数据上的 15k 步):
- 同义鲁棒性(余弦相似度):JEPA 为 0.9997,MLM 为 0.9414。*
- 观察到维度坍塌(91% 的方差在一个分量中);正在进行更强的 VICReg 正则化工作。*
- 如果坍塌得到解决,JEPA 可能提供固有氨基酸感知的密码子嵌入,这是 MLM 无法实现的。*
6.3 路线图(未来 12 个月)
- 扩展 CodonRoBERTa 在公开的 36M 序列 mRNABERT 数据集上,并添加与 ProtT5‑XL 的对比对齐。*
- 扩展微调 到其他生物(酵母、小鼠、Pichia)并发布更新的专家检查点。*
- 完成 CodonJEPA 通过修复坍塌,与 mRNABERT 进行基准测试,并将其作为可选嵌入层集成到流水线中。*
- 丰富流水线 加入 RFdiffusion 用于 de novo 骨架生成,并添加表达稳定性预测头。*
7. 实际部署
- 推理硬件 – 单个 A100(≥ 16 GB VRAM)或任何具有 ≥ 12 GB 显存的 GPU,用于 92 M 参数基础模型。*
- 软件栈 – PyTorch 2.5+,
transformers≥ 4.40,flash‑attn2用于高效注意力,FSDP 用于训练。* - 许可证 – 所有组件(ESMFold、ProteinMPNN、OpenMed 模型)均为 MIT 或 Apache-2.0,允许商业使用。*
- 成本 – 整个训练工作量在 AWS spot 实例上的成本约为 165 美元(4 × A100,55 GPU 小时)。*
8. 参考文献
- Jumper, J. et al. “Highly accurate protein structure prediction with AlphaFold.” Nature (2021).*
- Lin, Z. et al. “Evolutionary‑scale prediction of atomic‑level protein structure with a language model.” Science (2023).*
- Dauparas, J. et al. “Robust deep learning‑based protein sequence design using ProteinMPNN.” Science (2022).*
- Cheng, J. et al. “CodonBERT: a language model for codon optimization.” Nucleic Acids Research (2024).*
- Xiong, Y. et al. “mRNABERT: advancing mRNA sequence design with a universal language model and comprehensive dataset.” Nature Communications (2025).*
- Zhong, Y. et al. “Large mRNA language foundation modeling with NUWA for unified sequence perception and generation.” bioRxiv (2026).*
- Warner, B. et al. “ModernBERT: Smarter, Better, Faster, Longer.” arXiv (2024).*
完整的代码、训练好的检查点以及 25 物种 CDS 数据集将在 Hugging Face 上的 OpenMed 组织下发布。
摘要: OpenMed 发布了一个端到端的蛋白质工程流水线,包含 CodonRoBERTa-large-v2(困惑度 4.10,CAI 0.404)以及一个在 55 GPU 小时内训练完成、成本为 165 美元的 25 物种密码子优化模型套件。
标题: OpenMed CodonRoBERTa 多物种 mRNA 语言模型发布