近年来,人工智能在自然语言处理、计算机视觉、语音处理等领域取得了突破性进展,而将深度学习技术应用于物理、化学、生物、医学等传统科学领域,即AI for Science(科学智能),正逐渐成为新的交叉学科焦点,展现出巨大潜力并引发广泛关注。字节跳动研究院(ByteDance Research)也积极投身于AI for Science研究,聚焦机器学习与量子化学、大规模量子化学计算、AI制药等前沿领域,致力于推动该领域的持续发展。本文将简要介绍团队近两年在相关领域取得的研究成果,并期待与业界同仁展开更多交流合作。
在机器学习和量子化学方向,我们提出的LapNet算法在训练速度上比代表性模型FermiNet提升了10倍,同时能处理的化学体系规模和精度均达到当前领域最高水平。在大规模量子化学计算领域,我们开发的Periodic DMET算法采用经典与量子混合计算(实际为经典计算机模拟),仅用20个量子比特即可达到传统方法需近万个量子比特才能实现的精度。在AI制药领域,我们设计的LM-Design模型基于大量蛋白质序列数据及部分蛋白质结构与序列对应数据,实现了从蛋白质结构到序列转换的高精度模型,目前已成为蛋白质序列设计的顶尖方法。
量子化学物理学家狄拉克曾指出:”物理学和化学的基本定律已完全明了,但应用这些定律却异常困难,因为所得方程往往过于复杂难以求解。”量子化学正是基于量子力学原理研究化学现象的学科,其核心任务是通过计算方法求解分子或周期性体系(如固体)的电子薛定谔方程,从而推算出体系的基态能量、电极性等特性,即所谓的从头计算(ab initio)问题。传统方法如密度泛函理论(DFT)和耦合簇(CCSD)等,或因精度不足,或因规模受限,而近年来,机器学习方法为解决此类问题提供了新思路。
NN-VMC(Neural Network based Variational Monte Carlo)是其中一条重要路径。该方法利用神经网络近似薛定谔方程的波函数,通过随机采样获取体系中电子的空间分布,进而计算基于薛定谔方程的体系能量。通过最小化能量上界并迭代优化神经网络参数,最终得到近似最优的波函数和体系基态能量。图1展示了NN-VMC的基本原理。其关键在于神经网络设计和学习算法优化。DeepMind和ICL于2019年提出的FermiNet是NN-VMC领域的代表性方法,随后多家机构又提出了改进方案。自2021年起,我们与北京大学合作开展了一系列研究,提出了包括NN-VMC+ECP、NN-DMC、LapNet等新方法。
NN-VMC+ECP结合了NN-VMC和赝势ECP(Effective Core Potential)技术,通过定量表示原子内侧轨道电子的势能,大幅降低计算量,有效提升效率和体系规模。NN-DMC则将神经网络与扩散蒙特卡洛法(DMC)结合,通过虚时演化计算体系基态能量,相比FermiNet等现有方法能显著提高精度和规模。LapNet采用前向拉普拉斯算子优化神经网络学习过程,直接计算动能和哈密顿算子,省去黑塞矩阵计算环节,实现约10倍的加速效果。我们将ECP、DMC和前向拉普拉斯技术结合,进一步探索更大规模体系的计算可能性。此外,NN-VMC方法还成功应用于固体薛定谔方程求解、分子体系力场、电极化计算等问题,证明其广泛实用性。图2展示了当前NN-VMC方法的精度和规模对比,我们提出的LapNet在更高精度下实现了更大体系计算,最大体系包含116个电子。
由于直接求解薛定谔方程的计算规模限制,量子嵌入方法(Quantum Embedding Method)成为解决大规模量子化学计算的有效途径。密度矩阵嵌入理论(DMET)通过分而治之和多精度计算实现体系规模扩展。其基本思路是将体系划分为若干Fragment,对每个Fragment及其环境进行高精度计算,其他部分采用低精度计算,并根据需要并行处理,最终合并结果并迭代逼近原始体系。图3示意DMET方法流程:首先划分体系为Fragment组,对关注Fragment及其环境采用高精度方法(如CCSD)计算,其他部分用低精度方法(如Hartree–Fock)处理,并行计算所有Fragment,通过迭代优化使低精度解逐渐逼近高精度结果。
我们基于经典计算机和量子计算机两种范式实现了DMET及其变体SIE,推动大规模量子化学体系计算。其中,量子计算化学方法利用量子计算机实现DMET的高精度计算部分,经典计算机完成低精度计算部分。我们开发了DMET-ESVQE(分子体系)和Periodic DMET(周期性体系)方法,分别仅需16个和20个量子比特,即可达到传统方法需144个和近万个量子比特的计算精度。图4展示了混合计算机上实现的DMET Periodic方法过程:输入晶体,划分体系并行计算每个Fragment,关注Fragment及其环境在量子计算机上用U-CCSD求解,其他部分在经典计算机上用Hartree–Fock法处理,最终合并结果并迭代优化。
量子蒙特卡洛法(包括VMC、DMC)是量子化学的核心算法之一。我们创新性地将量子计算与量子蒙特卡洛法结合,充分发挥量子计算优势,部分解决符号问题。在AI制药领域,基于AI技术的药物设计已成为新范式。我们开发了多种小分子药物候选生成方法,包括MARS和DESERT。MARS从种子分子开始,通过马尔可夫链蒙特卡洛法(MCMC)和打分函数(评估亲和性、成药性、无毒性等)自动生成候选分子,利用图神经网络(MPNN)预测分子编辑操作,仅需打分函数和分子数据库即可生成全新候选分子。DESERT通过Sketching(基于靶点形状采样互补形状)和Generating(形状到分子式生成)两步生成候选分子,图5展示了其过程。图6示意Shape2Mol模型架构,编码器将3D分子形状编码为中间表示,解码器根据该表示生成分子式。DESERT在2022年小分子药物候选生成领域表现最佳。
近年来,我们聚焦大分子药物设计,特别是蛋白质设计。蛋白质设计包括抗体和多肽药物设计,其核心是从功能出发设计蛋白质结构,再确定对应序列。我们开发的LM-Design模型实现了从蛋白质结构到序列的高精度转换。该模型由结构编码器(已训练的图神经网络)和序列解码器组成,解码器基于预训练的蛋白质语言模型(Protein Language Model),并加入结构适配器进行序列改写。图7展示LM-Design模型架构。其学习过程类似BERT的掩码语言建模,通过多次还原被掩码序列中的氨基酸符号,充分利用蛋白质序列进化关系,大幅提升预测准确率。图8显示LM-Design是目前效果最佳的蛋白质序列生成模型。
致谢:感谢任维络、吕定顺、顾全全、吴凯、郑在翔、周奕、罗曼平、张震宇为本文撰写提供的宝贵建议和帮助。
