Partial Differential Equations (PDEs) are the bedrock for modern computational sciences and engineering, and inherently computationally expensive. While PDE foundation models have shown much promise for simulating such complex spatio-temporal phenomena, existing models remain constrained by the pretraining datasets and struggle with auto-regressive rollout performance, especially in out-of-distribution (OOD) cases. Furthermore, they have significant compute and training data requirements which hamper their use in many critical applications. Inspired by recent advances in ``thinking" strategies used in large language models (LLMs), we introduce the first test-time computing (TTC) strategy for PDEs that utilizes computational resources during inference to achieve more accurate predictions with fewer training samples and smaller models. We accomplish this with two types of reward models that evaluate predictions of a stochastic based model for spatio-temporal consistency. We demonstrate this method on compressible Euler-equation simulations from the PDEGym benchmark and show that TTC captures improved predictions relative to standard non-adaptive auto-regressive inference. This TTC framework marks a foundational step towards more advanced reasoning algorithms or PDE modeling, inluding building reinforcement-learning-based approaches, potentially transforming computational workflows in physics and engineering.
翻译:偏微分方程(PDEs)是现代计算科学与工程的基石,其求解本质上是计算密集型的。尽管PDE基础模型在模拟此类复杂时空现象方面展现出巨大潜力,但现有模型仍受限于预训练数据集,并且在自回归推演性能方面存在不足,尤其是在分布外(OOD)情况下。此外,它们对计算资源和训练数据的需求巨大,这阻碍了其在许多关键应用中的使用。受近期大型语言模型(LLMs)中“思维”策略进展的启发,我们首次为PDE提出了一种测试时计算(TTC)策略,该策略在推理阶段利用计算资源,以更少的训练样本和更小的模型实现更准确的预测。我们通过两种类型的奖励模型来实现这一目标,这些模型评估基于随机方法的模型在时空一致性方面的预测性能。我们在PDEGym基准的可压缩欧拉方程模拟上验证了该方法,结果表明,相较于标准的非自适应自回归推理,TTC能够捕获更优的预测结果。这一TTC框架标志着向更先进的PDE建模推理算法(包括构建基于强化学习的方法)迈出了基础性的一步,并可能彻底改变物理学与工程领域的计算工作流程。