利用 NVIDIA ALCHEMI 实现 AI 驱动材料发现的革命性突破

事实证明，AI 可发挥倍增效应，帮助创造一种未来，科学家可以设计全新材料，而工程师可以将这些设计无缝转换为生产计划，而这一切都无需涉足实验室。随着 AI 不断重新定义创新界限，这种曾经难以实现的愿景现在更触手可及。

认识到这种范式转变后，AI、化学和材料科学交叉领域的专业知识正在融合。AI 专家、化学家和材料科学家之间的这种跨学科合作旨在超越传统的、难以扩展的试错方法。

他们正在共同构建 AI 模型和数据集，将设计到生产的周期从十年缩短到几个月，志在实现从高效电池到可生物降解聚合物的下一代创新。

本文将介绍 NVIDIA ALCHEMI （AI 实验室用于化学和材料创新），旨在利用 AI 加速化学和材料的发现。

AI 加速的化学和材料发现工作流程

新型材料的发现大致可以分为四个阶段：

1. 假设生成
2. 解空间定义
3. 属性预测
4. 实验验证

假设生成 

• 见解合成：快速合成基于化学信息的大语言模型（LLM）（即根据化学文献训练或微调的 LLM），以合成大量化学文献。
• 假设表述：利用化学信息型语言模型作为思想合作伙伴来制定和完善假设，利用语言模型的能力在看似无关的概念之间建立联系。

解空间定义 

• 大型数据库搜索： 从现有化学数据库中识别出感兴趣的子空间。
• 有针对性的结构生成： 采用生成式 AI，根据所需的属性提出新的科学候选者。

属性预测 

• 分子表征学习： 利用习得的分子表征来预测候选项在解决方案空间中的属性。
• 原子模拟： 应用 AI 代理 (MLIP、GNN) 预测高保真属性。
• 密度泛函理论（DFT）验证： 通过 DFT 模拟验证预测属性。

实验验证 

• 候选推荐： 提出要在实验室中验证的候选（例如，通过贝叶斯优化来平衡已知化学的利用与白色空间的探索）。
• 实验自动化：利用支持人工智能的自动驾驶实验室和主动学习，加速实验室合成和测试。

加速化学和材料科学的研发

NVIDIA ALCHEMI 致力于借助 AI 的强大功能加速化学和材料科学的研发。ALCHEMI 将为开发者和研究人员提供 API 和 NVIDIA NIM 加速的推理微服务，从而实现：

• 创建和部署生成式 AI 模型，以探索广阔的材料宇宙，并根据所需的属性推荐新的潜在候选材料。
• 与密度泛函理论（Density functional theory，DFT）等传统的 从头算 模拟方法相比，开发和利用人工智能代理模型，在准确性和计算成本之间实现平衡。
• 可访问的化学和材料信息学工具和预训练的基础模型，用于将材料表示映射到属性，以进行快速筛选。
• 用于生成合成数据集的仿真工具，以便针对新用例训练和微调人工智能模型。

NVIDIA ALCHEMI 正在推出首款加速材料发现的 NIM，专注于加速一类名为机器学习原子间相互作用电位（MLIPs）的 AI 代理模型的常见但重要的应用。

机器学习原子间交互潜力 

在深入了解 NIM 之前，最好先介绍几个背景概念。传统的计算方法多种多样，从经典物理学到精确解，不一而足。每种方法都需要在准确性和计算成本之间进行权衡。精度更高的方法（例如 DFT）计算成本要高得多。例如，对具有 10 个原子的系统进行 DFT 计算需要几分钟，而对 100 个原子进行计算则需要几个小时。但是，对于原子数超过 1000 的系统，DFT（按原子数立方扩展）可能需要数周时间，因此对于大型复杂系统（例如具有 102 至 106 个原子的纳米粒子、具有 104 至 105 个原子的蛋白质）或作为高吞吐量模拟的一种方法，DFT 无法实现。

通过利用 AI 将原子结构映射到势能和原子力，MLIP 可在准确性和计算成本方面提供最佳性能（图 2）。MLIP 通常使用 GNN 将原子结构表示为图形，将原子表示为节点，将预设截止半径内的原子间距离表示为边缘。这使 GNN 能够轻松捕捉构成分子或材料的原子之间的空间关系。

从材料属性预测到大规模分子动力学模拟，MLIP 在化学、材料科学和生物学领域应用广泛。这与物理信息神经网络（PINNs）形成对比，后者嵌入物理方程中的知识，主要与计算流体动力学相关。

几何松弛 

在许多化学和材料发现工作流程中，一项关键任务是通过几何优化和能量比较来区分稳定材料和不稳定材料。这在涉及生成模型的工作流中尤为重要，因为生成的候选项本身可能并不稳定。

在几何松弛期间，材料的能量通过迭代评估每个原子上的力（推理）和调整原子在势能井底部方向的位置（优化）来最小化。每种候选材料可能需要数千个松弛步骤来实现此最低能耗。然后，通过将稳定材料的能量与其他候选材料进行比较来确定其优先级，以便根据所需属性进行进一步评估。

NVIDIA 批量几何松弛 NIM 

虽然 MLIPs 可显著加速几何松的能量和力计算 (与完整的 DFT 相比)，但实际实施仍然非常耗时。这就是为什么 NVIDIA 开发了 NVIDIA Batched Geometry Relaxation NIM，以加速几何松计算。NVIDIA NIM 只是一个公开 API 的容器。此 NIM 侧重于加速潜在电池材料的几何松，分别使用 MACE-MP-0 和 AIMNet2 模型在 Atomic Simulation Environment (ASE) 中。传统的基于 CPU 的模拟工作负载只能部分利用 MLIPs 的全部性能，并且在 CPU-GPU 数据移动中会产生严重的通信损失。此外，MLIPs 通常未充分利用 GPU，因为它们一次只处理一个系统。

通过利用 NVIDIA Warp (用于编写 GPU 加速模拟代码的 Python 开发者框架)，您可以编写常规的 Python 函数，并让 Warp 在运行时将其编译为高效的 GPU 内核代码。您无需一次进行一次几何放松，而是可以启动批量几何放松模拟，从而高效并行运行数百项能量最小化操作，并最大限度地利用可用 GPU 资源。这一点至关重要，因为每个候选者都需要数千个放松步骤。

表 1 显示了在单个 NVIDIA H100 80 GB GPU 上使用 MACE-MP-0 模型 和 ASE FIRE 几何优化器 的 2,048 个中小型无机晶体系统（每个周期单元 20-40 个原子）的结果。对于不使用 NVIDIA Batched Geometry Relaxation NIM 的 2,048 个样本，使用 NIM 大约需要 15 分钟，而使用 NIM 大约需要 36 秒，加速约为 25 倍。将批量大小从 1 增加到 128，可进一步将几何松弛加速至 9 秒，代表约 100 倍的加速。

表 2 显示了使用 AIMNet2 和 ASE FIRE 对 GDB-17 数据库中的 851 种中小型有机分子 (~20 个原子/分子) 的可比加速情况。启用 NIM 后，加速度约为原来的 60 倍，从 11 分钟缩短到 12 秒。通过将批处理大小从 1 增加到 64 (表示加速约为 800 倍)，可以观察到进一步的加速。观察到的加速差异 (使用 MACE-MP-0 时为 100 倍，使用 AIMNet2 时为 800 倍) 源于建模系统 (周期晶体与小分子)，而非部署的模型。

SES AI 是领先的锂金属电池开发商，正在探索将 NVIDIA ALCHEMI NIM 微服务与 AIMNet2 模型结合使用，以加速识别电动汽车所用的电解质材料。

SES AI 首席执行官 Qichao Hu 表示：“SES AI 致力于通过 AI 加速的材料发现来推进锂电池技术的发展，使用我们的 Molecular Universe 项目来探索和识别有潜力的锂金属电解质发现候选产品。“使用 ALCHEMI NIM 微服务与 AIMNet2 结合，能够显著提高我们绘制分子属性的能力，从而大幅节省时间和成本，并加速创新。”

SES AI 最近在半天内绘制出 100,000 个分子，有可能使用 ALCHEMI 在一小时内实现这一目标，这表明微服务将对材料筛选效率产生变革性影响。

展望未来，SES AI 的目标是在未来几年内绘制多达 10 亿个分子的属性地图，推动 AI 驱动的高吞吐量发现。

开始使用 NVIDIA 批量几何放松 NIM

要开始使用批处理几何图形松弛（Batched Geometry Relaxation）NIM，您需要以下内容：

• 熟悉 Python 和 ASE
• 具备运行 Docker 容器的知识
• 需要 MACE-MP-0（材料）或 AIMNet2（分子）模型

程序 

请务必理解，NVIDIA NIM 只是一个公开 API 的容器。与 NIM 交互涉及两个关键步骤：

首先，启动 NIM 容器：

docker run --rm -it --gpus all 
-p 8003:8003 \
<NIM container address>

请注意，--gpus all 指示 NIM 将在哪些设备上加载模型和优化器的单独实例。如果存在多台设备，NIM 将在可用的模型实例之间分发请求。其次，我们公开端口转发 8003:8003，将本地端口 8003 映射到容器端口 8003，这是 NIM API 的默认通信端口。客户端的请求将被路由到此端口。

接下来，启动并实例化容器后，您现在可以提交请求：

import os, requests, json
import numpy as np 
from ase import Atoms
from ase.io.jsonio import MyEncoder, decode

atoms: list[Atoms] = # This is your ase.Atoms input molecules

# Define the url of the NIM
# below is a typical local IP address and port
url: str = 'http://localhost:8003/v1/infer'

# Prepare input atoms by converting to json
data = {"atoms": json.dumps(atoms, cls = MyEncoder)}
headers = {'Content-Type': "application/json"}

# Submit request to NIM
response = requests.post(url, headers = headers, json=data)
response.raise_for_status()

# Convert from json output back to ase.Atoms
optimized_atoms = decode(response.json())

注意事项：

• 鉴于 ASE 是一种热门的分子建模 Python 包，输入数据以 ase.Atoms 的形式表示，这是描述材料系统的一种灵活方式。
• The user must provide the URL of the NIM. This URL, http://localhost:8003/v1/infer, comprises three components:
  • IP 地址：本例中的 localhost
  • 端口：8003 已映射到 Docker 命令中
  • NIM 推理 API 端点：v1/infer
• 用户使用 ase 实用程序将 ase.Atoms 转换为 .json 字符串，并使用 requests 库提交请求。这个 POST 请求将数据传输到“批量几何松弛 NIM”，并从客户端代码异步启动松弛过程。当分子或材料收敛到松弛状态时，它们会被传回客户端，并可以转换回 ase.Atoms，如上面的代码片段所示。

结束语 

使用 NVIDIA Batched Geometry Relaxation NIM 后，MLIP 计算速度提高了 800 倍。这种加速接近三个数量级，为数百万候选项的高吞吐量模拟打开了大门，从而使使用高质量数据训练的新一代基础模型成为可能，并提高下游属性预测能力。它还支持对更复杂和更逼真的系统进行仿真，从而释放新的化学和应用。这仅仅是开始。通过 ALCHEMI，NVIDIA 旨在端到端加速 AI 赋能的化学和材料研发工作流，并开创突破性发现的新时代，助力打造更可持续、更健康的未来。

注册以在 NVIDIA 批量几何松弛 NIM 可供下载时接收通知 。