蛋白质设计本质上是一个搜索问题。为一条8个残基的蛋白质序列选择氨基酸,可能的组合就超过250亿种。而为一条100个残基的蛋白质序列选择氨基酸,可能性的数量将超过可观测宇宙中原子的数目。在这广阔的空间中,某处就存在着能够开启新疗法之路的蛋白质。由于组合数量过于庞大,无法进行穷举搜索,因此找到它取决于是否知道该去哪里寻找。

RFD4-Proteina是英伟达BioNeMo推出的一款开放模型,由华盛顿大学医学院蛋白质设计研究所联合开发,它利用从自然界中学习到的规律来在这个空间中进行导航。该模型能够协同设计蛋白质序列和全原子结构,并使用奖励引导的推理时搜索,将计算资源集中在最有潜力的候选方案上。

本教程将向计算生物学家、机器学习工程师和药物发现研究人员展示如何利用奖励引导推理时搜索,针对肿瘤学靶点PD-L1设计蛋白质结合候选物,并介绍该方法在多种设计任务中的实验验证结果。

用户可以通过两种方式访问RFD4-Proteina:

选择开放模型,用于训练、推理、方法开发,并直接控制采样代码;或者

选择NIM微服务,通过HTTP接口实现快速的推理部署。

RFD4-Proteina基于Proteina-Complexa的部分隐变量方法构建,将显式骨架坐标与隐表示、学习得到的序列和原子级变量相结合。其自编码器解码器能够同时生成序列和全原子结构,从而省去了与ProteinMPNN等模型之间单独的反向折叠衔接步骤。

前置条件

开放模型:需要支持的Linux GPU主机和Python 3.12环境;请遵循代码仓库的设置指南,并使用受支持的英伟达GPU。模型缓存方面,设置过程会自动下载Hugging Face模型包;只有在需要指定其他缓存目录时才需设置HF_HOME。

NIM:自行部署需要满足NIM前置条件的x86_64或arm64 Linux系统、Docker、英伟达驱动、英伟达容器工具包以及一个NGC API密钥。NIM API端点(不建议用于生产环境运行):可通过NVIDIA build.nvidia.com,借助Shell和Python脚本进行HTTP访问,无需本地GPU或模型安装。

关于RFD4-Proteina

上图(图1)所示的搜索循环,正是RFD4-Proteina集中计算资源的核心环节。

在推理时推理搜索过程中,模型对生成轨迹进行去噪以获得结构,对这些设计方案进行评分,保留并拓展有潜力的部分状态,并持续进行直至达到生成标准。

奖励可以用来衡量几何特征、相互作用或独立的预测置信度。这种额外的搜索机制能够针对特定目标优化设计方案,从而产生评分更高的设计结果。在本教程后续部分,用户将看到如何通过开放模型和NIM启用这一功能。

相较于RFD3,RFD4-Proteina在速度和内存使用方面均有提升,原因如下:

该模型基于残基级隐表示而非原子级表示进行运算,从而避免了内存使用量的激增。

扩散过程的工程优化将采样步数从Proteina-Complexa约400步缩减至60步,并借助模型的推理能力将计算资源集中于流程中最关键的部分。

得益于预测器(OF3以及诸如tmol的物理方法)采用CUDA加速,推理时搜索的速度也更快。

该模型及其核心计算单元与cuEquivariance团队及BioNeMo推理运行时(BIR)团队联合设计,以兼顾速度与可扩展性。

从分子设计构想到候选设计方案

上图(图3)所示的工作流程展示了如何利用RFD4-Proteina,从科学目标出发,生成可供实验测试的候选方案。智能体可以协调这些工具并跟踪结果,而科学家则负责定义目标,并利用实验证据指导下一轮迭代。

准备分子输入与运行环境

首先明确生成过程的标准与要求。例如,RFD4-Proteina需要用户定义哪些部分可以改变、哪些部分必须保持固定,以及计算层面和实验层面的成功标准。

使用AtomWorks来生成或处理输入数据。AtomWorks是用于为RFD4-Proteina准备结构和设计注释的分子数据层。其Python和PyMOL工作流程可用于准备固定原子、可设计残基、热点区域以及分子上下文信息。在设计周围蛋白质结构时,应保留所提供的配体或模体化学结构。

开放模型和NIM均接受带注释的CIF文件,或是条件规格说明加上其对应的源结构。下方的规格说明固定了靶点,并定义了新的结合物。

代码仓库中来自PDB 5O45的裁剪版PD-L1结构作为靶点。固定靶点的序列和骨架坐标,允许其侧链重新排列,并将结合物生成引导至Tyr56、Met115和Tyr123这几个热点原子。新的结合物包含100到150个残基,序列与坐标均为自由变量。

克隆代码仓库,然后在其根目录下运行示例。这会生成pdl1.json文件;两条路径均使用同一个靶点CIF文件和条件设置。

(代码部分略,保持原样)

开放模型设置

在受支持的Linux/CUDA系统上,使用当前版本的uv(已在0.12.24版本上验证通过)安装代码仓库环境。如果已安装uv,请先更新以确保其能够识别CUDA 13.2。设置过程会缓存Hugging Face模型包,包括配套的自编码器。由于模型仓库为私有仓库,需使用具有访问权限的账户通过hf auth login进行身份验证。

(代码部分略)

NIM环境设置

要部署NIM,请在安装有Docker和英伟达容器工具包的受支持GPU主机上,于Bash环境中运行以下代码块。

(代码部分略)

生成序列与结构

开放模型生成方式

将以下内容保存为代码仓库根目录下的design.py文件。同一脚本同时支持生成和搜索功能。Proteina.from_pretrained("nvidia/RFD4-Proteina")会加载当前Hugging Face模型及配套的自编码器检查点。

(代码部分略)

在代码仓库根目录下生成一个候选方案:

(代码部分略)

输出结果将保存在runs/pdl1-open/inference_samples/目录下。添加num_replicates=10参数即可生成十个候选方案,且每次运行需选择不同的out_dir目录。Hugging Face会缓存下载内容;若需复现特定版本的模型,可在from_pretrained中添加revision="<commit>"参数。

NIM生成方式

将以下内容保存为nim_client.py文件。该脚本会将PD-L1条件规格说明和靶点CIF文件发送至NIM,然后保存生成的结构及响应指标。

(代码部分略)

安装requests库并在代码仓库根目录下运行客户端脚本。请将pdl1.json文件与靶点CIF文件放在同一目录下。

(代码部分略)

开放模型与NIM的性能优化

为加快设计迭代速度,在六种经过测试的Hopper GPU配置下,BioIR后端相比默认的编译后端实现了1.37至1.73倍的加速(见下方图4)。在H100 PCIe上,整套测试的总耗时从161秒降至111秒,运行时间减少约31%。该基准测试使用批次大小为1,并统计了15种序列长度(128至1024个Token)下的预热端到端耗时总和。

借助外部预测器进行推理时搜索

使用OpenFold3进行开放模型搜索

通过运行./scripts/download_example_data.sh ./data_path安装可选的OpenFold3运行时环境及权重文件。按照脚本输出的路径设置OF3_IMAGE、OF3_MODEL_DIR和OF3_JAX_CACHE_DIR环境变量,然后运行mkdir -p "$OF3_JAX_CACHE_DIR"。GPU主机上必须已安装Apptainer。

此示例使用了采用开源Apache 2.0许可证发布的AlphaFold3代码库,以及同样采用Apache 2.0许可证发布的OpenFold3模型权重。

使用搜索预设重新运行design.py:

(代码部分略)

该预设使用两个分支、宽度为3的束搜索,并在第20、40、50和60步进行评分。它以所设计链的OpenFold3 pLDDT值(oracle.plddt_coi)作为奖励,设定最低阈值为50,并将奖励值达到80的合格推演结果保存下来。此指标衡量的是对结合物的置信度,而非结合证据本身;用户应单独检查界面和重折叠相关指标。输出结果保存在runs/pdl1-open-search/目录下。

使用OpenFold3进行NIM搜索

将以下内容保存为nim_search_request.py文件。该脚本复用PD-L1输入数据,并提交一个奖励引导的OpenFold3搜索请求。

(代码部分略)

启用NIM预测器

按照预测器指南中的说明,挂载OpenFold3的reward.py回调脚本。该脚本会报告整个复合物的oracle_of3_nvcf.overall_plddt指标;结合物界面需单独评估。

将预测器指南中的OpenFold3 NVCF回调脚本保存为reward.py文件,并导出带有NVCF调用凭证的OF3_NIM_API_KEY环境变量。使用该回调脚本重新启动服务:

(代码部分略)

运行搜索脚本。生成的设计方案和响应指标将保存在runs/pdl1-nim-search/目录下。

(代码部分略)

该请求使用两个分支、检查点[0, 20, 40, 50, 60],以及宽度为3的束搜索。它设定了50的奖励下限,并保存得分至少为80的中间设计方案。这些均为针对OpenFold3整体pLDDT(即整个复合物的预测置信度指标)设定的示例阈值。搜索阶段使用的奖励标签与最终报告所用的标签是两套独立的标签体系。

搜索输出的数量会存在差异。在最终选定候选方案之前,应检查生成的CIF文件及保留的各项指标。在进行计算量比较时,应将重折叠和评分过程也纳入考量。

筛选与排序候选方案

检查约束条件、模体保留情况、原子碰撞、界面情况、重折叠一致性以及结构多样性指标。OpenFold3和RoseTTAFold3可提供预测选项;TMol则可提供受Rosetta启发的全原子评分与优化功能。

常用指标包括:来自OpenFold3等蛋白质结构预测模型的ipTM(界面预测模板建模得分)、ipSAE(基于对齐误差的交互预测评分)和RMSD(均方根偏差)评分,以及来自TMol等基于物理原理模型的相互作用能量数据。

实验测试与迭代优化

上方表1展示了RFD4-Proteina在各类分子任务中所具备的广泛能力。配套的研究论文不仅证实该模型支持如此多样的任务类型,更表明在输出质量和处理效率两方面,该模型在这一广泛的分子任务范围内均达到了业界领先水平。

将这种处理效率转化为科学进展,需要为合成、表达及特定任务检测环节挑选出多样化的候选方案。再将实验结果反馈至目标设定、约束条件以及下一轮迭代的排序标准中。

借助AI智能体开展研究项目

使用rfd4-proteina-campaign技能来协调生成、搜索与评估流程,同时记录输入数据、模型版本及计算资源使用情况。rfd4-proteina-package技能则能够生成包含结构数据、各项指标以及筛选报告的结果归档文件。

使用rfd4-proteina-campaign技能,基于pdl1.json文件设计100至150个残基的PD-L1结合物,同时保留其靶点和热点约束条件。通过from_pretrained方法调用Hugging Face模型,并结合OpenFold3使用。确定评判标准和计算资源预算后,对比常规生成方式与奖励引导搜索方式的效果,并通过rfd4-proteina-campaign运行研究项目。最终返回最多十个具有多样性的候选方案及其评分和计算资源使用情况;并使用rfd4-proteina-package对结果进行归档打包。

将该工作流程应用于你的设计问题

RFD4-Proteina将快速生成、智能搜索与智能体可编程性整合于同一工作流程之中。序列与全原子结构的联合生成简化了候选方案的创建过程,而奖励引导搜索则能将计算资源集中用于实现特定任务目标。共享的条件规格说明将这两种使用形式与通用评估工具相连接,而NIM的HTTP API则使智能体能够提交请求并查看结果。

将开放模型用于方法开发,将NIM用于应用集成(可选配BioIR加速),从而将计算探索与实验反馈融入一个迭代式的设计研究流程中。

立即开始使用RFD4-Proteina NIM微服务,在GitHub上探索该开放模型,或者让你的智能体直接调用BioNeMo智能体工具包中的相关技能。

NVIDIA