APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs

📄 arXiv: 2411.00180v1 📥 PDF

作者: Felix Koehler, Simon Niedermayr, Rüdiger Westermann, Nils Thuerey

分类: cs.LG

发布日期: 2024-10-31

备注: Accepted at Neurips 2024. The code is available at https://github.com/tum-pbs/apebench and APEBench can be installed via "pip install apebench"

💡 一句话要点

APEBench：用于偏微分方程自回归神经模拟器的基准测试平台

🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)

关键词: 偏微分方程 神经模拟器 基准测试 自回归模型 可微编程

📋 核心要点

  1. 现有偏微分方程求解的神经模拟器缺乏统一的评估标准和可比性，阻碍了该领域的发展。
  2. APEBench通过提供一个包含多种偏微分方程、可微求解器和评估指标的综合基准来解决这个问题。
  3. 实验表明，神经模拟器与传统数值模拟器存在相似性，为进一步研究神经模拟器的理论基础提供了方向。

📝 摘要（中文）

本文提出了自回归偏微分方程模拟器基准测试平台（APEBench），这是一个全面的基准测试套件，用于评估求解偏微分方程的自回归神经模拟器。APEBench基于JAX，并提供了一个无缝集成的可微模拟框架，该框架采用高效的伪谱方法，支持1D、2D和3D中的46种不同的偏微分方程。为了方便对学习到的模拟器进行系统分析和比较，我们提出了一种新颖的展开训练分类法，并为偏微分方程动力学引入了一个独特的标识符，该标识符直接关系到经典数值方法的稳定性标准。APEBench能够评估各种神经架构，并且与现有基准不同，其求解器的紧密集成支持可微物理训练和神经混合模拟器。此外，APEBench强调展开指标以了解时间泛化，从而深入了解模拟偏微分方程动力学的长期行为。在几个实验中，我们强调了神经模拟器和数值模拟器之间的相似之处。

🔬 方法详解

问题定义：现有的偏微分方程神经模拟器研究缺乏一个统一的、可复现的基准测试平台。不同的研究使用不同的数据集、求解器和评估指标，导致难以进行公平的比较和分析。此外，现有方法在长时间序列预测和泛化能力方面存在挑战。

核心思路：APEBench的核心思路是提供一个基于JAX的可微模拟框架，该框架集成了多种偏微分方程、高效的伪谱方法和标准化的评估流程。通过提供一个统一的平台，APEBench旨在促进神经模拟器的系统分析、比较和改进。此外，该基准强调展开指标，以评估模型的时间泛化能力。

技术框架：APEBench的整体架构包括以下几个主要模块：1) PDE定义模块：包含46种不同的1D、2D和3D偏微分方程。2) 可微求解器模块：使用JAX实现高效的伪谱方法，支持自动微分。3) 训练模块：提供多种训练策略，包括展开训练和神经混合训练。4) 评估模块：提供标准化的评估指标，包括时间泛化能力评估。5) 基准测试模块：提供统一的接口和流程，方便用户进行模型评估和比较。

关键创新：APEBench的关键创新点在于：1) 提供了一个全面的、可微的偏微分方程模拟框架。2) 引入了一种新颖的展开训练分类法，用于系统分析不同的训练策略。3) 为偏微分方程动力学引入了一个独特的标识符，该标识符与经典数值方法的稳定性标准相关。4) 强调时间泛化能力评估，关注模型在长时间序列上的表现。

关键设计：APEBench的关键设计包括：1) 使用JAX实现可微求解器，方便进行梯度计算和优化。2) 采用伪谱方法，提高求解器的效率和精度。3) 提供多种损失函数，包括L1损失、L2损失和对抗损失。4) 支持多种神经架构，包括MLP、CNN和Transformer。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

APEBench的实验结果表明，神经模拟器在某些情况下可以达到与传统数值模拟器相当的精度。通过对不同神经架构和训练策略的比较，可以发现哪些方法更适合于特定的偏微分方程。此外，实验还揭示了神经模拟器在长时间序列预测和泛化能力方面的优势和局限性。

🎯 应用场景

APEBench可应用于各种科学和工程领域，例如流体动力学、热传导、电磁学等。通过使用神经模拟器加速偏微分方程的求解，可以降低计算成本，提高仿真效率。此外，APEBench还可以用于设计新型材料、优化控制策略和预测复杂系统的行为。

📄 摘要（原文）

We introduce the Autoregressive PDE Emulator Benchmark (APEBench), a comprehensive benchmark suite to evaluate autoregressive neural emulators for solving partial differential equations. APEBench is based on JAX and provides a seamlessly integrated differentiable simulation framework employing efficient pseudo-spectral methods, enabling 46 distinct PDEs across 1D, 2D, and 3D. Facilitating systematic analysis and comparison of learned emulators, we propose a novel taxonomy for unrolled training and introduce a unique identifier for PDE dynamics that directly relates to the stability criteria of classical numerical methods. APEBench enables the evaluation of diverse neural architectures, and unlike existing benchmarks, its tight integration of the solver enables support for differentiable physics training and neural-hybrid emulators. Moreover, APEBench emphasizes rollout metrics to understand temporal generalization, providing insights into the long-term behavior of emulating PDE dynamics. In several experiments, we highlight the similarities between neural emulators and numerical simulators.