开放式研究对推动创新至关重要,AI 与科学领域的众多突破正是通过开放式协作实现的。在数字生物学研究领域,NVIDIA Clara 为这种开放合作提供了支持。
Clara 是一个开源系列,涵盖生物学、化学和人类健康领域的模型、工具与方法。该系列包含适用于小分子生成式设计、合成路径预测、ADMET 性质预测以及蛋白质结构与序列协同设计等应用场景的模型。如需详细了解面向生物学和化学领域的 NVIDIA Clara 模型与工具,请访问 NVIDIA-Digital-Bio GitHub 资源库。
本文将介绍 Clara 开放模型系列的新成员 CodonFM。CodonFM 是一种专注于 RNA 的生物学语言模型,本文将阐述其模型设计原理,并探讨其在变异效应预测、mRNA 设计等多种任务中的应用方法。
CodonFM:面向 RNA 的开放基础模型
今天,NVIDIA 宣布推出 CodonFM,这是一个全新的先进 RNA 基础模型,现已加入 Clara 开放模型系列。CodonFM 通过读取由三个核酸组成的密码子来处理 RNA 序列,将每个密码子视作句子中的单词,而非孤立的核酸字母。这种方法借鉴自然语言中的语法分析方式,使模型能够学习遗传密码中复杂的“语法规则”。由此构建的模型,能够深入理解不同生物体之间密码子使用偏好所呈现出的复杂且依赖于上下文的模式。
一些常见的生物学语言模型是蛋白质语言模型,它们能够独立模拟蛋白质序列中的每个氨基酸残基。然而,这类模型忽略了同一个氨基酸可以通过不同的密码子(即同义变体)进行编码的事实,并且在细胞的蛋白质合成过程中,不同的同义密码子变体可能导致蛋白质产量的差异。
通过解析同义变体,CodonFM 能够理解编码相同氨基酸的不同 RNA 序列如何影响生物学功能,进而预测 mRNA 稳定性、翻译效率和蛋白质产量等特性。此外,该模型在预测与基因突变相关疾病风险方面的表现也得到提升。
CodonFM 基于 BERT 式双向编码器架构构建,使模型能够充分理解完整的输入 RNA 序列。凭借包含 2046 个密码子 tokens(6138 个核糖核酸)的大型上下文窗口,该模型可识别在数十亿年进化过程中不断优化的复杂长序列模式。
为了学习这种生物语言,CodonFM 利用来自 22000 个物种的 1.31 亿个蛋白质编码序列进行训练,这些序列包含从美国国立卫生研究院 – 美国国家生物技术信息中心 RefSeq 数据库中提取的数千亿个密码子 tokens。
CodonFM 提供多种不同的模型规模(80M、600M 和 1B 参数)以及两种预训练方法。随着模型规模的增大,其对编码相同氨基酸的同义密码子的区分能力也相应提升。密码子混淆现象(即模型错误预测一个同义密码子出现频率的情况)的减少,表明模型对密码子使用模式及与翻译相关的序列上下文有了更深入的理解(图 1)。此外,密码子混淆(即模型错误预测一个同义密码子出现频率的情况)的减少,反映了模型对密码子使用模式及翻译相关序列上下文的理解更加深入(图 1)。
此外,不同的预训练方法为生成的模型带来了各自独特的优势:
- 随机密码子遮罩:对序列中的密码子进行随机遮罩,而不考虑其在预训练语料库中的出现频率。该方法可训练 CodonFM 根据上下文环境预测缺失的密码子,有助于模型学习各类编码区域中遗传密码的基本语法结构。
- 密码子加权遮罩:在随机遮罩的基础上构建,通过结合密码子使用偏好,有选择性地遮罩在特定局部上下文中较为罕见的密码子。这种方法使模型能够更有效地捕捉与物种特异性功能相关的密码子选择模式,而非对所有密码子一视同仁。
正如本文基准测试所示,CodonFM 展现出明确的缩放规律:随着模型规模和数据集规模的增大,模型在多种应用场景中的准确性均有所提升,例如同义突变与错义突变分类、mRNA 翻译效率预测以及蛋白质丰度预测。
在生物任务中应用 CodonFM
CodonFM 在零样本和微调设置下均展现出广泛的适用性,能够支持多种分子与临床应用场景。本节将重点介绍 CodonFM 在不同生命科学应用中的表现,并为各项任务提供相应的模型实现代码片段。
预测突变效应的大小
CodonFM 通过对编码序列本身进行建模,捕捉密码子的上下文、冗余性及调控模式,而无需显式依赖蛋白质结构。这使得经过微调的 1B 参数 Encodon 模型(由 NVIDIA 与 Arc Institute 合作开发)在检测致病性错义突变方面表现出优异的性能,在区分疾病相关氨基酸替换与良性变异方面具有很高的准确性。
更重要的是,CodonFM 将这一功能进一步拓展至解释同义词变体这一更具挑战性的问题。同义突变不改变蛋白质序列,传统上难以被预测模型识别。Encodon 能够检测密码子使用偏好以及翻译水平上的细微变化,在 ClinVar 数据集中实现了对致病性与良性同义突变的高效区分,展现出其在解析沉默突变方面的独特能力。
以下代码段展示了如何利用预训练的 Encodon 模型进行突变评分任务:
# Task: score effect of single synonymous/missense mutation at a given codon
# Output: log-likelihoods for ref/alt codons + LLR per variant
# More details can be found here in the source code of CodonFM
# src/data/preprocess/mutation_pred.py and src/data/mutation_dataset.py
# 1) Configure model checkpoint
CKPT_PATH = "/path/to/NV-CodonFM-Encodon-1B-v1.ckpt" # change to your .ckpt
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
enc = EncodonInference(model_path=CKPT_PATH, task_type=TaskTypes.MUTATION_PREDICTION, ...
) # <- routes to predict_mutation
enc.configure_model() # loads Tokenizer + EncodonPL + weights
enc.to(DEVICE)
# 2) Prepare a mutation example and build a model batch
# Example CDS (coding DNA sequence) length must be divisible by 3 and in-frame.
cds = ("ATGCCGGCGGTCAAGAAGGAGTTCCCGGGCCGCGAGGACCTGGCCCTGGCTCTGGCCACGTTCCACCCGACC") # <--- replace with your full coding sequence (no introns, 5'->3')
# Choose a 0-based CODON index (not nucleotide index).
codon_idx = 10 #codon 10 in the CDS
# Define the ref codon present at that position and the alternate codon to test
ref_codon = "CGC"
alt_codon = "CGA" # e.g., a synonymous change
# --- Tokenize the full CDS into codon tokens using the tokenizer ---
tok = enc.tokenizer
context_length = 2048
cds = cds[:(context_length-2) * 3] # truncate the sequence to the context length
cds = '<cls>' + cds + '<sep>'
# Encode the full sequence to input IDs. The tokenizer works at the codon resolution.
input_ids = np.array(tok.encode(cds), dtype=np.int32)
attention_mask = np.ones_like(input_ids, dtype=bool)
# Add padding to the input_ids and attention_mask to the context length
input_ids = np.pad(input_ids, (0, context_length - len(input_ids)), ...)
attention_mask = np.pad(attention_mask, (0, context_length - len(attention_mask)), ...)
mutation_token_idx = codon_idx + 1 # the position in the tokenized sequence is shifted by 1 because of the <cls> token
设置 input_ids 和注意力遮罩后,获取 ref 和 alt tokens,对输入进行遮罩并完成批量处理,随后调用 predict_mutation 函数执行推理。
ref_tok = tok.convert_tokens_to_ids(ref_codon)
alt_tok = tok.convert_tokens_to_ids(alt_codon)
input_ids[mutation_token_idx] = tok.mask_token_id # replace the mutated codon with mask token
batch = {
MetadataFields.INPUT_IDS: torch.tensor(input_ids, dtype=torch.int32, device=DEVICE).unsqueeze(0),
MetadataFields.ATTENTION_MASK: torch.tensor(attention_mask, dtype=torch.bool, device=DEVICE).unsqueeze(0),
......
}
# 3) Run inference and interpret LLRs
out = enc.predict_mutation(batch, ids=["example_variant"])
print("IDs: ", out.ids[0])
print("log P(ref codon): ", out.ref_likelihoods[0])
print("log P(alt codon): ", out.alt_likelihoods[0])
print("LLR (ref - alt): ", out.likelihood_ratios[0])
在幕后,该脚本采用 Encodon 掩码语言模型,评估突变如何影响其序列上下文中的密码子概率。通过屏蔽目标密码子,模型分别预测参考序列与变异序列的概率,并计算其对数似然比(LLR)。正的 LLR 值表明突变在自然中较为少见,可能具有潜在破坏性;而负的 LLR 值则提示该突变可能被容忍,或在特定序列背景下更具优势。
mRNA 治疗设计
mRNA 设计正迅速成为现代疗法中的一种重要模态,有助于实现基因替换、蛋白质功能恢复以及可编程生物制剂的开发。该领域的一个关键挑战在于序列优化:即使是一个较小的蛋白质,也可以由大量同义的 mRNA 序列编码,而每个序列在表达效率、稳定性和免疫原性方面均可能产生不同的影响。
密码子的使用偏好以及序列上下文中的细微选择会显著影响翻译结果。CodonFM 为这些应用提供了高效的预测框架,在多种 mRNA 稳定性和表达水平的基准测试中表现出优异性能,能够实现蛋白质丰度的零样本预测和翻译效率预测,为优化 mRNA 设计奠定了坚实基础。
使用 CodonFM 进行微调
CodonFM 资源库实现了多种微调策略,帮助用户根据具体应用场景对预训练模型进行微调。这些策略包括:
- 低秩自适应 (LoRA):在微调预训练模型时,向每个 Transformer 层添加低秩适配器,以降低训练成本和内存占用。
- 纯头随机:随机初始化输出头并进行训练,同时冻结模型其余部分的参数。
- 预训练:使用预训练的输出头进行训练,同时冻结模型其余部分的参数。
- 完整:对模型的所有参数进行端到端的微调。
迈向可编程生物学
正如语言模型学会推理,蛋白质模型学会折叠,CodonFM 也学会了将 RNA 密码子与其生物学行为及蛋白质表达规律相联系。它使 RNA 从遗传信息的被动载体转变为一种可编程的语言,能够被解读、优化和重新设计。
这种读写生命形式语言的能力,构成了 NVIDIA 虚拟细胞计划的基石。通过发布 CodonFM 等开放、强大且可扩展的模型,研究人员和开发者得以构建不仅能够理解生物学,更能主动塑造生物学的 AI 系统。
我们诚挚邀请您携手 Arc Institute、Therna Biosciences、Greenstone Biosciences、Moonwalk Biosciences 以及斯坦福 RNA 医学计划的合作伙伴,共同探索并测试 CodonFM,推动生物智能的发展。
开始使用 CodonFM
CodonFM 基于为其他 Clara 开放模型提供支持的相同核心基础设施进行训练,利用 NVIDIA cuDNN 和 NVIDIA cuBLAS 库实现 GPU 原生加速,在基因组标记过程中对矩阵运算进行优化。输入数据集被转换为内存映射文件,以支持快速高效的数据流,同时采用 NVIDIA NeMo Run 作为核心训练配置与编排框架。
或者,通过 NVIDIA BioNeMo 框架 的 recipe 使用 Transformer 引擎,可将模型训练和微调速度提升高达 3 倍,同时准确性损失可忽略不计,从而保障可扩展性与计算效率。
准备好开启 CodonFM 之旅了吗?
- 在 NVIDIA-Digital-Bio/CodonFM GitHub 中查找完整代码
- 浏览 Hugging Face 和 NGC 上的模型检查点
- 深入了解 CodonFM
- 通过 BioNeMo 框架的内置 recipe 加速训练与推理
致谢
感谢以下人员对本篇博文的支持与贡献:Sajad Darabi、Fan Cao、Mohsen Naghipourfar、Hani Goodarzi、Sara Rabhi、Yingfei Wang、William Greenleaf、Yang Zhang、Cory Ye、Jonathan Mitchell、Timur Rvachov、T.J. Chen、Daniel Burkhardt 和 Neha Tadimeti。