来字节跳动实习,还能有机会发 Nature 子刊?
没错,ByteDance Research 与北京大学物理学院陈基课题组合作的一项研究近期登上了国际顶级刊物 Nature Communications。作者之中,共同一作伟中就是字节跳动的实习生同学,来自北大物理学院陈基课题组,正在攻读凝聚态理论方向的博士。
论文名为《 Towards the ground state of molecules via diffusion Monte Carlo on neural networks 》,作者将神经网络与扩散蒙特卡洛方法结合,大幅提升神经网络方法在量子化学相关任务上的计算精度、效率以及体系规模,成为最新 SOTA(State Of The Art,当前最高水平)。
可以通俗的理解为,在一个分子中,我们知道原子核的位置坐标,但是并不知道电子的位置坐标。但利用参数众多的大规模神经网络,输入原子核的位置,就可以计算出电子出现在不同坐标的概率,推算出电子在空间中的分布情况。
有了这样的研究,进一步就可以预测分子化学性质,预测反应发生的速率、催化效率,为分子动力学研究做基础。有了这些科学知识积累,我们才能进一步在生物制药、化工、材料等领域做出创新。
作者将基于神经网络的试探波函数运用于固定节点面的扩散蒙特卡洛方法 (Diffusion Monte Carlo, or DMC) ,用以精确计算具有不同电子特性的原子以及分子系统。
扩散蒙特卡洛方法是量子化学领域精确计算分子和材料基态能量的常用方法之一。通过与扩散蒙特卡洛方法结合,作者显著提高了量子化学中神经网络 SOTA 方法的计算精度与效率。此外作者还提出了一种基于经验线性关系的外推方法,大幅改善了分子结合能计算。总体而言,该计算框架作为求解量子多体问题的高精度方法,为化学分子性质的深入理解提供了更强大的工具。
2018年以来,多个研究小组将神经网络运用于变分蒙特卡洛方法 (Variational Monte Carlo, or VMC) 中[1,2,3],借助神经网络强大的表达能力,得到了更为精确的分子基态能量。本工作于2022年公开时,基于神经网络的变分蒙特卡洛方法中的 SOTA 工作是 DeepMind 于2019年提出的 FermiNet [2],能够在规模较小的体系上得到非常精确的结果。然而变分蒙特卡洛方法的精度受限于神经网络的表达能力,在处理较大体系时会有越来越明显的精度问题。此外该类方法在处理较大体系时收敛非常缓慢,对计算资源提出了巨大挑战。
扩散蒙特卡洛方法作为量子化学领域的经典高精度算法之一,具有精度高、可并行性好、适合进行大规模计算等良好的特性。此外扩散蒙特卡洛可以突破神经网络的表达能力限制,利用投影算法超越变分蒙特卡洛方法的精度。
本工作中,作者将 SOTA 的神经网络 (FermiNet) 作为试探波函数与扩散蒙特卡洛方法结合。新的计算方法相比于 FermiNet 显著提升了精度并减少了所需的计算步数。本工作中所设计并实现的扩散蒙特卡洛软件具有神经网络友好、GPU友好、并行友好的特点,可以与广泛的神经网络波函数结合,自动提升其精度与效率。
使用神经网络对大型分子体系进行量子蒙特卡洛计算时,由于算力限制,所能使用的神经网络的表达能力也会受到一定限制。为了模拟这一场景,作者使用了仅仅两层的神经网络来研究第二、三排的原子。计算结果显示随着体系变大,变分蒙特卡洛方法的精度愈来愈差,而扩散蒙特卡洛方法所带来的精度提升也愈来愈明显。
作者在一系列分子体系上也验证了基于神经网络的扩散蒙特卡洛方法的有效性,包括氮气分子,环丁二烯以及双水分子。在所测试的体系上均观察到了明显的计算精度提升。
本工作公开前,量子化学领域中基于变分蒙特卡洛的神经网络波函数方法只处理过30电子以内的小型分子。本工作首次将神经网络波函数方法应用于42~84个电子的体系,即苯环与双苯环。计算结果显示,扩散蒙特卡洛方法在精度上显著优于变分蒙特卡洛方法,同时可以用少一个数量级的计算步数达到相同或更优的精度。
作者在考察神经网络的不同训练阶段所对应的能量时,在很多体系上均发现变分蒙特卡洛与扩散蒙特卡洛的计算结果具有经验性的线性关系(下左图)。使用该线性关系对双苯环的解离能计算进行外推,显著提升了计算精度,得到了吻合于化学实验的结果(下右图)。
本工作表明,基于神经网络的扩散蒙特卡洛方法在精度与效率上均优于变分蒙特卡洛方法。作者开源的扩散蒙特卡洛代码可以与量子化学领域不断推陈出新的神经网络[4,5]快速结合,实现对研究社区的赋能。此外扩散蒙特卡洛方法也可以与处理真实固体的周期性神经网络[6]、带赝势的神经网络[7]等一系列方法结合,在相应任务上提升计算效果。
参考文献
[1] Han, J., Zhang, L., & Weinan, E. (2019). Solving many-electron Schrödinger equation using deep neural networks. Journal of Computational Physics, 399, 108929.
[2] Pfau, D., Spencer, J. S., Matthews, A. G., & Foulkes, W. M. C. (2020). Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3), 033429.
[3] Hermann, J., Schätzle, Z., & Noé, F. (2020). Deep-neural-network solution of the electronic Schrödinger equation. Nature Chemistry, 12(10), 891-897.
[4] Gerard, L., Scherbela, M., Marquetand, P., & Grohs, P. (2022). Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?. In Advances in Neural Information Processing Systems.
[5] von Glehn, I., Spencer, J. S., & Pfau, D. (2023). A Self-Attention Ansatz for Ab-initio Quantum Chemistry. The Eleventh International Conference on Learning Representations.
[6] Li, X., Li, Z., & Chen, J. (2022). Ab initio calculation of real solids via neural network ansatz. Nature Communications, 13(1), 7895.
[7] Li, X., Fan, C., Ren, W., & Chen, J. (2022). Fermionic neural network with effective core potential. Physical Review Research, 4(1), 013021.
论文发表之后,技术范儿采访了伟中同学,来看看他在字节跳动做研究的体验吧。
技术范儿:
你是怎样参与到这项工作中来的?
伟中:
我学校的导师和 ByteDance Research 一直保持密切的合作,我大四的时候就每两周跟着导师来字节交流一次,探讨前沿技术。
字节的氛围和学校科研很像,大家交流起来就像平常做学术一样,有很多博士毕业、有丰富科研经验的同学,各种问题都能找到一个特别精通的人去请教,收获很大。
本科毕业、博士入学后,我就正式来字节实习了。
技术范儿:
你在这项工作中,都做了哪些贡献?
伟中:
我入职的时候,我的 mentor(也是论文的一作)已经把代码都写好了。
在此基础之上,我加了一些电子移动的新算法,做了一个波函数节点面的可视化,理清了一些 bug。
最重要的是设计、执行并分析了大量计算实验,以证明算法是有效的,用真实的测算来和算法推演结果对比,验证算法的有效性。
技术范儿:
完成这项工作后,
你收获了哪些方面的成长?
伟中:
一方面,在科研过程中,我调研了本领域内大量的文献,也提升了自己的代码水平;
另一方面,我学会了如何把一个好的科研成果写成论文,如何更好的展示相关成果,提升了自己的学术水平。
技术范儿:
在字节做科研工作是怎样的状态?
伟中:
非常专注,大家都在一个空间里工作,有一种沉浸式的氛围让我专注在科研项目中。
效率非常高,我 mentor 是斯坦福的数学博士,懂的特别多,我可以随时请教;当我灵光一现的时候,我可以直接飞书定个会议室、随手找个白板就可以开始讨论新的想法。
而且,我们每周都会有一次文献讨论会,会有很专业的同学给大家讲解最新的学术论文并组织讨论,这个过程中会有很大的收获。
另外就是字节跳动计算资源丰富,可以方便地申请很多 GPU 资源,是我之前能用到的算力的很多倍了。
技术范儿:
你们是一个怎样的团队?
伟中:
应该是北京最豪华的量子化学团队之一了吧。
我们很多团队成员来自中科院、清华、北大,还有很多世界各地量子化学前沿研究团队的博士生也会来这里实习一段时间。
虽然团队成立距今只有两年,但已经发表了两篇 Nature Communications,还有许多其他基础学科研究的论文成果,都有实习同学的参与。
在这里,我们经常可以和世界最前沿的量子化学团队交流,不仅有邮件往来,我们还会邀请他们来开分享会,分享他们在公开论文里看不到的一些想法。
当然,做学术之外,大家日常相处也很轻松,会一起去环球影城团建,也经常一起外出聚餐。