Extending machine learning model for implicit solvation to free energy calculations

本文介绍了 Lambda 溶剂化神经网络（LSNN），这是一种基于图神经网络的隐式溶剂模型，其在力和炼金变量导数上进行了联合训练，旨在实现与显式溶剂模拟相当的自由能预测精度，同时为药物发现应用提供显著的计算加速。

要点

想象一下，你正在试图弄清楚一把特定的钥匙（药物分子）与一把特定的锁（蛋白质）的契合程度。为了准确做到这一点，你需要理解这把钥匙在水环境中的行为，因为在人体中，一切事物都浸泡在由水分子构成的海洋里。

本文介绍了一种名为LSNN（Lambda-溶剂化神经网络）的新工具，它帮助科学家比以往的方法更快、更准确地计算这种“水行为”。

以下是关于问题、旧解决方案和新修复方案的简单故事：

问题：“拥挤的房间”与“幽灵”

为了理解药物如何起作用，科学家使用计算机模拟。

• “金标准”（显式溶剂）： 想象一下，试图模拟一把钥匙在一个房间里，而你必须追踪围绕它移动的每一个人（水分子）。你必须计算钥匙如何撞到 A 人，然后是 B 人，接着是 C 人。这极其准确，但就像试图数清海滩上的每一粒沙子一样。它需要巨大的计算能力和时间。
• “快速”方法（隐式溶剂）： 为了节省时间，科学家过去曾假装水不是由个体组成的，而是一种平滑、不可见的雾气。他们使用一个简单的数学公式来猜测雾气如何推动钥匙。这非常快，但这层“雾气”只是一个粗略的猜测。它经常搞错细节，导致关于药物是否有效的预测不准确。

旧的“机器学习”修复方案（以及它为何失败）

最近，科学家尝试使用人工智能（特别是神经网络）来让这层“雾气”变得更聪明。他们通过向 AI 展示水如何推动钥匙（即力）来训练它。

• 缺陷： 这就像只教某人如何转动方向盘来教他们开车，却从不告诉他们速度有多快或使用了多少汽油。AI 学会了向正确的方向推动钥匙，但它无法计算将钥匙从一个地方移动到另一个地方所需的总“努力”（能量）。因此，旧的 AI 模型无法用于比较不同药物的总能量。

新解决方案：LSNN

作者创建了LSNN，这是该 AI 的更智能版本。他们不仅教它如何推动（力），还教它当慢慢“打开”或“关闭”药物与水之间的相互作用时，能量是如何变化的。

类比：
想象你正在试图测量一个背包的重量。

• 旧 AI： 你能感觉到肩带拉肩膀有多重（力），但你无法判断背包重 10 磅还是 20 磅，因为秤坏了。
• LSNN： 他们修好了秤。现在，AI 不仅能感觉到拉力，还能通过观察当你慢慢向包里添加或移除物品时拉力如何变化，来计算出确切的总重量。

他们如何测试它

该团队在包含约30 万个小分子的庞大库上训练了这个新 AI。他们将其与“金标准”（缓慢的数沙子方法）和旧的“雾气”方法进行了测试。

结果：

1. 速度： LSNN 是一名短跑运动员。它计算结果大约只需20 秒。“金标准”耗时近28 分钟（约 1600 秒）。旧的“雾气”方法也很快（约 15–22 秒）。
2. 准确性：
   • “金标准”最准确（得分为 1 分中的 0.86 分）。
   • LSNN 排名第二，得分为0.73。这比得分低得多（0.48 到 0.63）的旧“雾气”方法有了巨大改进。
   • 本质上，LSNN 达到了“金标准”级别的准确性，但运行速度却达到了“雾气”级别。

那更大的东西呢？（蛋白质）

本文还尝试使用 LSNN 来预测药物如何粘附在大蛋白质上（这是药物发现的终极目标）。

• 结果： 它显示出希望，但尚未完美。当他们尝试将其用于完整的蛋白质系统时，准确性下降了。作者认为，这是因为 AI 主要是在简单的小分子上训练的，可能会“过度思考”大蛋白质中复杂的相互作用。然而，它仍然显示出清晰、一致的模式，表明它可以得到改进。

底线

本文提出了一种新的“智能雾气”（LSNN），它解决了以前 AI 模型的最大缺陷：无法计算总能量。

• 它快（像旧的简单数学）。
• 它准确（更接近缓慢、昂贵的模拟）。
• 它对于比较不同药物是可靠的。

作者得出结论，该工具为药物发现的未来奠定了坚实的基础，使科学家能够以更快的速度筛选数百万种潜在药物，而不会牺牲找到真正疗法所需的准确性。

技术摘要

技术摘要：将隐式溶剂化机器学习模型扩展至自由能计算

问题陈述
隐式溶剂模型通过以数学近似平均力替代离散溶剂分子，为分子模拟提供了一种计算高效的框架。然而，其精度往往落后于显式溶剂模型，限制了其在绝对自由能比较等精确热力学计算中的应用。尽管近期基于机器学习（ML）的方法通过在力匹配数据上训练神经网络改进了隐式溶剂描述，但仍存在一个关键局限：仅靠力匹配只能将势能确定至一个任意常数。因此，这些模型无法在不同化学物种间提供有意义的绝对自由能比较。此外，传统隐式模型（如 GBSA、PBSA）依赖简化的溶剂可及表面积（SASA）项来处理非极性贡献，这些项容易产生显著误差。

方法
作者引入了$\lambda$-溶剂化神经网络（LSNN），这是一种基于图神经网络（GNN）的隐式溶剂模型，旨在克服标准力匹配的局限性。

• 架构：LSNN 建立在 Katzberger 和 Riniker 的基础工作之上，后者利用在标准 GBSA 参数上训练的三层不变 GNN；LSNN 将相互作用 GNN 与多层感知机（MLP）集成，以处理非线性依赖关系。
• 训练目标：与仅最小化预测力与参考力之间差异的先前方法不同，LSNN 将炼金变量（alchemical variables）的导数纳入损失函数。具体而言，模型被训练以匹配：
  1. 溶质原子上的平均施加力（MAFs）。
  2. 关于静电耦合因子（$\lambda_{elec}$）的导数。
  3. 关于位阻耦合因子（$\lambda_{steric}$）的导数。
• 损失函数：修改后的均方误差（MSE）损失函数定义为：
  $$L = w_F \left( \langle \frac{\partial U_{solv}}{\partial r_i} \rangle - \frac{\partial f}{\partial r_i} \right)^2 + w_{elec} \left( \langle \frac{\partial U_{solv}}{\partial \lambda_{elec}} \rangle - \frac{\partial f}{\partial \lambda_{elec}} \right)^2 + w_{steric} \left( \langle \frac{\partial U_{solv}}{\partial \lambda_{steric}} \rangle - \frac{\partial f}{\partial \lambda_{steric}} \right)^2$$
  其中权重经经验调整（比例为 1:1:1.2）。这确保了模型学习一个保守向量场，从而使标量势能能够近似真实的平均力势（PMF）。
• 数据集与训练：该模型在来自 BigBind 数据集的约 280,000 个小中性分子数据集上进行训练。数据按 80:10:10 的比例划分为训练集/验证集/测试集，并设有特定约束，确保与 FreeSolv 数据集中相似的分子被保留用于测试。力和相互作用导数是使用 OpenMM 结合 GAFF 力场在 0.5 纳秒模拟中计算的。
• 实现：该模型利用 PyTorch Autograd 进行导数计算。为确保在完全解耦状态下总能量为零，能量项会乘以其对应的 $\lambda$ 值。

关键结果
LSNN 框架针对 FreeSolv 数据集（647 个中性小分子）的实验水合自由能进行了基准测试，并与显式溶剂（TIP3P）及传统隐式模型（OBC2、GBn2）进行了比较。

• 精度：LSNN 与实验值的相关系数（$R^2$）达到 0.73，显著优于传统隐式模型（GBn2：$R^2$ 0.48；OBC2：$R^2$ 0.63），并接近显式溶剂模拟的精度（TIP3P：$R^2$ 0.86）。
• 计算效率：与显式溶剂方法相比，LSNN 表现出显著的速度提升。LSNN 每分子的平均计算时间为 20.47 秒，而 TIP3P 为 1658.54 秒（约 27.6 分钟）。LSNN 的速度与 GBn2（15.82 秒）和 OBC2（21.81 秒）相当。
• 结合亲和力初步结果：在蛋白质 - 配体复合物的初步测试中（使用 MM-LSNN，即用 LSNN PMF 替换 GBSA 溶剂化项），该模型显示出与实验值的线性相关性（完整蛋白质系统的 $R^2$ 为 0.44）。然而，作者指出，由于训练域局限于小分子，导致长程相互作用被高估，因此目前在完整蛋白质系统上的独立性能仍受限制。

意义与主张
本文主张，LSNN 通过将训练从简单的力匹配扩展到包含炼金导数，代表了基于机器学习的可迁移势能的根本性转变。该方法使得绝对自由能的计算成为可能，而这一能力此前受限于力匹配中的任意常数问题。

作者断言，LSNN 成功捕捉了配体去溶剂化趋势，并在不同配体间保持了排序的一致性，提供了一个平衡显式溶剂模拟精度与隐式模型计算效率的框架。虽然当前迭代主要针对小分子的热力学一致自由能计算进行了优化，而非大生物分子的综合构象采样，但该框架为药物发现领域的未来应用奠定了基础，包括向带电配体的扩展以及蛋白质 - 配体相互作用能量估算的潜力。