走向“AI物理学家“:做科学发现和仿真 | 未来光锥AI For Science社群分享回顾

2023年7月12日,未来光锥AI For Science社群分享第四期邀请到西湖大学工学院助理教授、人工智能与科学仿真发现实验室PI吴泰霖博士,介绍 走向“AI物理学家”:做科学发现和仿真。以下为未来光锥对吴泰霖博士分享内容,以及部分观众提问进行的简要整理。
图片
分享嘉宾:吴泰霖
西湖大学人工智能与科学仿真发现实验室PI
今天我将介绍AI在用于科学仿真和发现的学科前沿进展以及开放问题,主要包含我在这些方面做的一些工作。我们的目标是最终能够实现更加自动化的AI科学家,帮助人们更快更准确地进行科学仿真、设计和科学发现,当然目前离这个目标还有很大的距离。
首先,对于大部分与AI有关的问题,都可以认为由以下几部分组成,首先是数据(data),我们需要发现数据的规律和结构等,从而进行建模;若该系统是一个动力学系统,我们还可以对系统进行仿真;随后,我们也会关心这个系统的预测的准确性和鲁棒性,以及它如何受到压缩的影响。最后,我们也会思考,这样的模型如何可以反过来设计或控制系统。这四个问题是AI领域的重要问题。
图片
图1 AI中的四个基础问题|图片由吴泰霖提供
如果我们把图1蓝色文本框中的“Data”改成“Scientific data”,那么就有了AI For Science的问题。无论在物理、化学、材料,还是生命科学等领域,我们都需要发现复杂系统的简单规律,进行仿真,以及基于此进行反向设计。这些在工程领域中同样重要,比如飞机的设计、可控核聚变的仿真、化工中的设计等。
我的研究也可以用AI+Science来概括。研究一方面涉及通过AI加速科学发现和仿真,另一方面是science for AI,思考科学问题给AI带来了哪些挑战和启发,如此也可以反过来提高AI算法的鲁棒性和泛化性。
今天的分享主要围绕AI在科学发现以及科学仿真中的一些问题。
第一部分,AI for 科学发现,这个问题比较复杂,但是可以进行简化,也就是用一种简单但必要的设置来进行研究。比如说,科学家在进行研究时会遇到很多不同场景、环境下的问题,每个环境中可能会包含不同的因素、各种条件等。对于这样的问题,一种方式是通过标准的深度学习方式,在每一个场景中学习一个对应的模型。这样的方式虽然可行,但没有好的普适性,可解释性也相对较低。另一方面,科学家的研究方式与以上的做法有所不同,他们会尝试在不同的环境之中找出较为简单、普适的规律。虽然这样做只能描述其中的某些方面,但是这种做法非常简单,而且泛化性很强。
由此,我们可以总结出科学发现的几个主要部分,第一个是从高维复杂的系统中提取重要概念,比如,物理学中的熵、能量等概念(提取概念,concept identification)。第二步是基于这些提取出来的概念,学习这些概念满足的方程(理论学习,theory learning)。第三步是概念的泛化,将特定场景中出现的概念泛化成更加普适的概念(概念泛化,concept generalization)。我认为,这三步便是科学发现的必要环节(当然也有别的重要环节,比如实验设计等)。
图片
图2 科学发现的必要环节|图片由吴泰霖提供
过去我在AI用于科学发现的工作也是从这几方面展开的,今天主要介绍理论学习这部分。这一部分的设置如下:
Environment: a dynamic system|不同的环境,比如可以是一个动力学系统(E.g., moving particles, fluid, etc.)
Environment state|环境的状态,可以是(粒子)系统中一个点的速度和位置(E.g., particle positions)
Trajectory: A sequence of (particle, fluid) states|系统不同状态的合集(E.g., positions of particles over time)
Domain: a subspace of one or more environments where the same physical laws apply|系统满足某物理定律的区域,不同的区域可能满足不同的物理定律
AI物理学家
在这个设置中,我们输入的是系统在不同时间的状态。我们希望可以学习出一个能够根据系统的过去预测将来的模型,并且这个模型具有较强的可解释性。基于这个问题,我和Max Tegmark提出了AI物理学家(AI Physicist)的架构(Physical Review E, 2019[1])。它最核心的部分是提出,我们要学习不同的理论。每个理论是一个2-tuple(f, c),包含两个对象,一个是预测函数(prediction function),能够根据过去预测未来;另一个是区域分类器(domain classifier),可以给定一个数据来判断这个数据是否应该被理论解释。这两个函数都可以用神经网络或者通过符号(方程)来表示。
那么,如何在不同的环境中,学习能解释这些环境中数据的简单、普适的理论?我们总结了物理学家常用的四个归纳偏置(inductive biases,表1),并将之用于我们的模型之中。通过这样的方法,我们便可以学到一个可解释性更强、更准确、学习更快的模型。
图片
表1 AI物理学家测试的几种策略|Wu & Tegmark, 2019[1]
具体而言,图3-a是AI物理学家的架构。蓝色大方框包含是AI物理学家的模型架构,它会遇到不同的环境,在每个环境中,它会通过可微的分而治之(divide and conquer)的方法从数据中学习不同的通过神经网络表示的理论,每个理论通过竞争的方式尝试解释所有数据,而我们的损失函数使得对于每个数据,越好的理论学得越快,使得最后每个理论能够专业化解释一部分的数据。学好的理论会加入到理论中心(theory hub)中。当它遇到一个新的环境,不会重新学习所有东西,而是先通过用已有的理论来尝试解释数据。只有对于数据中部分不能被解释的部分,它才会学习新的理论。
此外,在理论中心中,我们还用奥卡姆剃刀(Occam’s Razor)将神经网络简化成符号表示的方程。同时,我们也采用统一(unification)的方式,将结构类似的方程统一成同一个方程,只不过再加一些参数。最后,是终身学习(lifelong learning),将新学习到的理论用于新数据的解释(图3-b)。整个AI物理学家优化的目标是最小化表示长度(minimal description length),这个表示长度表示包括模型本身的复杂程度和我们的模型在进行预测时的误差。这四个策略走下来,就形成了一个闭环,这样的一个系统就可以不断尝试解释新的数据并能不断总结简化已经学习的理论。
图片
图3-a|图片来自Wu & Tegmark, 2019[1]
图片
图2|图片由吴泰霖提供
AI费曼2.0上线
我们又进一步做了AI费曼2.0(AI Feynman 2.0,Udrescu et al.,NeurIPS 2020[2]),它可以重新发现《费曼物理学讲义》中的最著名的100个物理方程(举例见表2),并具有很高的鲁棒性(当数据噪音达到10%时,仍能发现73%的方程)。
图片
表2 | 图片来自AI Feynman 2.0,Udrescu et al.,2020[2]
具体来说,我们可以先训练一个神经网络,有了神经网络后,就可以进行很多不同的测试,比如测试这个神经网络是否有某一些结构,有没有某一些对称性等。比如,我们可以交换两个输入,看结果是否有变化。一旦有某些对称性,我们就可以简化神经网络,并把数据集变成子数据集。如此,我们可以递归地(recursively)把数据进行简化。简化到最核心的部分是简化成能够进行暴力搜索(brute-force search)得到方程。
图片
图3 AI费曼2.0通过在训练好的神经网络中进行测试并发现其内部的对称性,不断递归,最终发现从输入映射到输出的方程|图片来自Udrescu et al. 2020[2]
我认为在AI用于科学发现领域中有很多开放性的问题。比如,在概念发现(concept identification)工作中,我们如何能够从高维的数据中发现比较简单的概念?还有一个我认为很有意思的开放性问题是,我们如何构建一个真正普适的AI科学家?图4来自今年发表在《自然·通讯》上的一篇文章(Cornello et al., 2023[3]),列出了AI可能可以进行工作的环节。我认为,如果AI用于科学发现能真正实现自动化的话,能够很好地加速各类科学发现。
图片
图4|图片来自Cornello et al., 2023[3]
AI for 科学仿真
接下来是第二部分,AI用于科学仿真。AI用于科学仿真在物理、能源、机械、天气、材料等众多领域都有广阔的应用前景。我之前在这方面的工作主要开发基于神经网络的代理模型,加速偏微分方程的仿真,所做的工作应用在激光等离子体仿真(模拟激光和等离子体的耦合的动力学,并且能够反过来设计激光的参数)、地下流体仿真、计算机图形学中的仿真中等。图5是我们所做的这个AI用于科学仿真的任务介绍。
ut代表是该系统在t时间的状态;这个状态可以通过多种数据结构来表示,比如网格、粒子,或者是一个无限维的函数(u (t, x),为时间t和空间x的函数)。a是系统不随时间变化的参数,可以是偏微分方程的参数,也可以是系统随空间变化但不随时间变化的参数。表示的是系统的边界条件,包括系统的形状以及边界上的方程。
图5 AI for 科学仿真的任务介绍|图片由吴泰霖提供
我们的任务是,给定数据是否可以学一个代理模型,它可以根据这个系统过去某时间的状态ut以及参数a、边界条件,预测它下一个时间的状态ut+1。根据系统不同的数据结构,我们也可以构建不同类型的神经网络。
科学仿真是一个困难的问题,我觉得至少有以下三个主要难点。第一,它是一个大规模的系统。一般来说,这个系统可能会有上百万甚至上亿个不同的粒子或者网络的单元格。第二,它是一个多尺度、多分辨率的系统。多尺度是指系统的动力学往往会横跨多个时空尺度,比如从蛋白质到细胞,到组织再到生物等,又比如等离子体仿真、催化剂的仿真中,系统的动力学横跨微观、介观、宏观。多分辨率指的是,系统的一小部分的动力学非常剧烈,需要非常精细的分辨率才能精确模拟,它们的数量可能只占这个系统的很小一部分,但能量可能占据系统的一半,而系统的其他部分比较静态,只需要比较粗粒化的网格就可以模拟。第三个难点是长时间预测的准确度。因此,降低长时预测的误差累积也是很重要的。
传统的科学仿真比较准确,可解释性较高,同时往往有误差方面的保证。但是,它的核心问题在于非常慢,因为它往往需要很小的时间度和空间度来满足它数值的稳定性。并且,对于一部分问题需要用一些隐式(implicit)的方法来演化。而且,对于一些多分辨率的系统而言,它需要精细到分辨率最细的地方,非常浪费计算成本。近几年来,基于神经网络的代理模型有了一系列非常好的工作。这种方式的好处在于它非常的快。因为它可以用更大的时间和空间分辨率,并且可以用显式的方法进行演化,也可以学习更好的表示。
在这里,我介绍一个我最近在这方面的工作:学习自适应多分辨率物理仿真(Wu et al. ICLR 2023[4])。在这里,我们针对很多物理系统都具有的多分辨率的问题。比如天气预报、很多偏微分方程的模拟中,往往系统的某些部分变化剧烈需要精确的模拟,而其他部分变化不那么剧烈,不需要精细的模拟。目前的神经网络代理模型主要是就系统对未来的预测进行优化,并没有优化计算成本。
在这项工作中,我们提出了首个可以同时演化系统,并同时优化计算成本的神经网络架构。它基于两个图神经网络(graph neural networks, GNNs),一个演化系统的状态,一个作为策略网络(policy network),在每个时间步,在每一条边上决定是否要粗化或者细化网格。我们通过强化学习训练这个模型,其奖励包含对未来的预测误差和计算成本,并通过一个的β参数进行可控权衡。这个β也作为模型的输入,使得在推理的时候,模型能够根据这项的权重,可控地决定在预测未来时,对减小预测误差和计算成本的重视程度。模型在一维的Burgers’ equation和二维的基于网格的仿真中进行了测试,超越了之前的SOTA模型。
AI 用于科学设计
当我们对科学的系统有了较好的仿真后,我们也可以进行一个好的设计(design)。这里的任务是,给定一个目标函数,通过AI的方法优化系统的参数或者组分,以优化目标函数。AI(更具体地说deep learning)之所以能做到这些,一个核心的insight是因为它不是通过一步,而是通过多步映射得到结果,其中的每一步只需要学习一些很简单的变化。
基于这个核心的insight,我认为AI在用于科学设计方面,有三类非常重要的方法。第一类方法是反向传播(backpropagation)。基于已学好的神经网络代理模型,我们可以通过一步步的回传来优化系统的参数。第二类是扩散模型(diffusion models),将高斯分布映射到数据分布,其中每一步的学习也都非常简单。还有第三类方法,强化学习(reinforcement learning)。比如在设计分子时,一开始只有一个原子,通过强化学习学习一个策略网络每一步加一个原子,这样以来,每一步的决策也都很简单,就这样一步步地得到最后的结果。这三类方法目前都有一定的工作,也有很大的扩展空间。另外,还有直接学习映射、通过贝叶斯的方式优化等方法。
图6 AI for scientific design的主要的三种方式|图片来自Xu et al. 2022[5], You et al. 2018[6]
参考文献
[1] Wu, Tailin, and Max Tegmark. "Toward an artificial intelligence physicist for unsupervised learning." Physical Review E 100.3 (2019): 033311.
[2] Udrescu, Silviu-Marian, Andrew Tan, Jiahai Feng, Orisvaldo Neto, Tailin Wu, and Max Tegmark. "AI Feynman 2.0: Pareto-optimal symbolic regression exploiting graph modularity." Advances in Neural Information Processing Systems 33 (2020): 4860-4871.
[3] Cornelio, Cristina, et al. "Combining data and theory for derivable scientific discovery with AI-Descartes." Nature Communications 14.1 (2023): 1777.
[4] Wu, Tailin, Takashi Maruyama, Qingqing Zhao, Gordon Wetzstein, and Jure Leskovec. "Learning Controllable Adaptive Simulation for Multi-resolution Physics." In International Conference on Learning Representations (ICLR 2023)
[5] Xu, Minkai, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. "Geodiff: A geometric diffusion model for molecular conformation generation." In International Conference on Learning Representations (ICLR 2022)
[6] You, Jiaxuan, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. "Graph convolutional policy network for goal-directed molecular graph generation." Advances in Neural Information Processing Systems 31 (2018).
观众提问
观众1
吴老师好,请问第一部分的工作相当于是用机器学习来反演出数据背后的非线性关系?
吴泰霖
是的,我们经常通过机器学习来反演出数据背后的一些关系,然后最后能通过一些方程来表示。
观众2
您在第一部分讲的东西是假设它最后有一个非常简单直白的公式的表达,但是比如说像一元六次方程、一元多次方程这种没有解析解的情况,如何用机器学习的方法来做?因为你假设的这些物理规律最后都能有一个被人类的公式描绘出来的写法,假设这种东西不存在怎么办?
吴泰霖
很好的问题。这里涉及到一个问题,一个东西没有解析表述,是不是等同于它不能被学习?或者是,它是不是可以通过某些另外的方式,通过图灵机来描述。举个例子,比如说一个一元六次方程,我能够有很多的数据。这些数据它的输入是方程参数,输出是它的解,它的复数解,或者说实数解,或者说没有解。假如获得了上百上千个这样的输入-输出组合,我觉得是有可能可以通过神经网络来fit一个这样变化的映射的。然后。再通过AI费曼来对学到的神经网络进行简化。我不确定以上的方法是否能成功,但值得尝试。