训练跨 25 种物种的 mRNA 语言模型仅花费 165 美元：从结构预测到密码子优化的全流程构建

来源：Hugging Face Blog 作者：OpenMed（专注于医疗与生命科学的开源智能体 AI 团队） 系列：Part II: Building the Pipeline, From Structure Prediction to Codon Optimization

背景

在蛋白质人工智能（Protein AI）领域，从概念性的治疗性蛋白质设计到最终可合成的 DNA 序列，中间存在巨大的工程鸿沟。在上一部分（Part I）中，OpenMed 团队梳理了蛋白质 AI 的版图，涵盖了从 AlphaFold 到 ESMFold 的结构预测架构，以及 ProteinMPNN 等开源蛋白质设计工具。

然而，仅仅预测结构或设计氨基酸序列并不足以让蛋白质在活细胞中高效表达。真正的挑战在于密码子优化（Codon Optimization）。遗传密码具有简并性：同一种蛋白质可以由天文数字般不同的 DNA 序列编码，但某些密码子排列方式的表达效率比其他方式高出 100 倍。例如，辉瑞-BioNTech 的新冠疫苗就经过了针对人类表达的密码子优化。

OpenMed 团队的目标是构建一个端到端的蛋白质 AI 流水线，将蛋白质从概念转化为表达就绪的 DNA。该流水线包含三个阶段：

1. 结构预测：确定蛋白质的 3D 形状。
2. 序列设计：设计能折叠成该形状的氨基酸序列。
3. 密码子优化：优化底层 DNA 密码子，确保蛋白质在目标生物体中高效表达。

在前两个阶段，团队使用了成熟的开源工具（Meta 的 ESMFold 和 Baker Lab 的 ProteinMPNN）。而在最核心的密码子优化阶段，OpenMed 团队投入了最大精力，从头构建了新的模型、训练基础设施和评估指标。本文详细记录了这一构建过程，包括架构探索、多物种扩展以及最终的端到端工作流。

核心内容

1. 流水线架构与组件

OpenMed 构建的流水线由三个主要组件构成，分别对应蛋白质工程工作流的不同阶段：

• 结构预测：使用 ESMFold（来自 Meta）。
• 序列设计：使用 ProteinMPNN（来自 Baker Lab）。
• 密码子优化：完全由 OpenMed 团队自主研发的新模型和基础设施。

团队将主要精力集中在 mRNA 优化环节，因为这是连接基因序列与实际生物表达的关键瓶颈。

2. 架构探索：寻找最佳 Transformer 架构

大多数生物语言模型是对自然语言处理（NLP）架构的改编，但密码子序列具有独特的统计特性：

• 小词表：由 64 个密码子组成，加上特殊标记，总词表仅为 69 个 token。
• 强位置依赖性：每三个核苷酸编码一个氨基酸，边界是生物学定义的，而非统计学习的。
• 物种特异性偏差：不同物种对密码子的使用偏好差异巨大。

为了找出最适合密码子级语言建模的架构，团队进行了严格的对比实验。

候选模型

• 基线：小型 CodonBERT（600 万参数，遵循 Sanofi 发布的架构）。
• ModernBERT：代表 NLP 社区最新效率创新的 2024 年架构，包含 RoPE 嵌入、Flash Attention 和交替的局部/全局注意力层。
• RoBERTa：Meta 的 ESM 蛋白质语言模型背后经过验证的“主力”架构。团队假设，既然 RoBERTa 变体能学习氨基酸模式，也可能学习密码子模式。

训练设置

• 数据：使用来自 E. coli RefSeq 的 250,000 个编码序列（CDS）。这是一个干净、注释良好的数据集，其密码子使用模式在文献中已有充分表征，可作为验证基准。
• Tokenizer：将每个密码子映射为一个单独的 token。这种极简设计尊重了生物学定义（每三个核苷酸对应一个氨基酸），不同于 NLP 中基于子词边界的 BPE 分词。
• 硬件与优化：在 4 张 A100 GPU (80GB) 上使用 FSDP 分片进行训练，步数为 15,000 至 25,000 步。所有模型均采用掩码语言建模（MLM），掩码率为 15%，与 ESM-2 的目标一致。

3. 实验结果与关键发现

结果：RoBERTa 显著胜出

实验结果毫不含糊：RoBERTa 在困惑度（Perplexity）上比 ModernBERT 高出 6 倍（4.01 vs 26.24）。尽管 ModernBERT 拥有更现代的注意力模式和高效架构，但在密码子序列上，经典的 RoBERTa 设计表现远超前者。

关键洞察 1：预训练的 NLP 权重无法迁移到生物学

团队尝试将 ModernBERT 从其在英语语料库上的预训练检查点初始化，期望利用其学到的注意力模式作为起点，但效果不佳。

• 原因分析：ModernBERT 在英语文本上的预训练引入了归纳偏差（如子词频率分布、位置注意力模式），这些偏差反而干扰了对密码子统计规律的学习。
• 对比：RoBERTa 从零开始随机初始化，纯粹在生物数据上训练，没有这种“包袱”。这与 ESM-2 和 ProtTrans 等主流生物模型的做法一致——它们都在生物数据上从头训练，而非微调自 NLP 检查点。

关键洞察 2：超参数调优解锁了生物一致性

这是最具实践意义的发现。团队对比了 CodonRoBERTa-large v1 和 v2：

• 差异：两者架构、数据、参数量完全相同。唯一区别是 v2 的学习率减半，且预热步数更长（2,000 步 vs 1,000 步）。
• 结果：
  • v2 的困惑度略高（4.10 vs 4.01），意味着它在预测被掩码的确切密码子时稍不准确。
  • 但是，v2 预测的密码子似然度与真实生物密码子偏好（通过密码子适应指数 CAI 衡量）的相关性提高了 16 倍。
• 结论：较慢的训练 schedule 让模型能够收敛到捕捉真实生物信号的表示中，而不是过拟合于表面统计特征。这证明 MLM 损失 alone 不能衡量生物相关性，必须结合领域特定的指标（如 CAI）。

4. 扩展至多物种与最终系统

基于 CodonRoBERTa-large-v2 的成功，团队将其扩展至 25 种物种，并在 55 个 GPU 小时内训练了 4 个生产级模型。最终构建了一个**物种条件化（species-conditioned）**的系统，这是目前其他开源项目中未提供的功能。

关键要点

• 架构选择：对于密码子级语言建模，经典的 RoBERTa 架构显著优于代表最新 NLP 效率创新的 ModernBERT，困惑度低 6 倍。
• 预训练陷阱：来自自然语言处理（NLP）的预训练权重（如英语文本）会引入不利于生物序列学习的归纳偏差，生物语言模型最好从头训练。
• 评估指标的重要性：仅优化掩码语言建模（MLM）损失是不够的。必须使用领域特定指标（如 CAI 相关性）来评估模型的生物实用性。
• 超参数敏感性：降低学习率并延长预热时间（Warmup）可以显著提升模型对真实生物信号的捕捉能力，即使困惑度略有上升。
• 端到端流水线：OpenMed 成功整合了 ESMFold（结构）、ProteinMPNN（设计）和自研 CodonRoBERTa（优化），实现了从蛋白质概念到合成就绪 DNA 的全流程自动化。
• 多物种支持：最终系统支持 25 种物种