APEBench

APEBench是由慕尼黑工业大学的机器学习中心创建的一个综合性基准套件，旨在评估自回归神经网络在求解偏微分方程（PDE）中的表现。该数据集包含46种不同的PDE动态，涵盖从一维到三维的空间维度，数据量庞大且多样化。数据集的创建过程结合了高效的伪谱方法和可微分模拟框架，确保了数据生成的高效性和准确性。APEBench主要应用于神经网络架构的评估和优化，旨在解决复杂PDE模拟中的精度和速度问题。APEBench is a comprehensive benchmark suite developed by the Machine Learning Center of the Technical University of Munich, aimed at evaluating the performance of autoregressive neural networks in solving partial differential equations (PDEs). This dataset encompasses 46 distinct PDE dynamics, spanning spatial dimensions from 1D to 3D, and features a large-scale and highly diversified dataset. The creation process of this dataset integrates efficient pseudospectral methods and differentiable simulation frameworks, ensuring the efficiency and accuracy of data generation. Primarily utilized for the evaluation and optimization of neural network architectures, APEBench targets addressing the accuracy and speed challenges in complex PDE simulations.

原始信息汇总

APEBench 数据集概述

数据集简介

APEBench 是一个基于 JAX 的工具，用于评估周期域上 1D、2D 和 3D 偏微分方程（PDE）的自回归神经模拟器。该工具包含一个基于谱方法的高效参考模拟器，用于程序化数据生成，无需下载大型数据集。由于该模拟器还可以嵌入到模拟器训练中（例如，用于“解算器在环”校正设置），这是第一个支持可微物理的基准套件。

数据集特点

• 支持维度：1D、2D 和 3D 周期域。
• 数据生成：基于谱方法的参考模拟器，用于程序化数据生成。
• 训练支持：支持“解算器在环”校正设置，适用于可微物理。

安装要求

• Python 版本：3.10 及以上。
• JAX 版本：0.4.12 及以上。

快速开始

• 示例代码：提供了一个训练 ConvNet 以模拟 1D 对流的示例，展示了训练损失、测试误差度量滚动和样本滚动。
• Colab 示例：提供了一个 Google Colab 示例，链接为 Open In Colab。

文档

• 文档链接：tum-pbs.github.io/apebench/。

背景

• 自回归神经模拟器：用于高效预测与微分方程相关的瞬态现象。
• 参考模拟器：基于 Exponax 的 ETDRK 方法。
• 神经网络架构：使用 PDEquinox 实现常见的架构，如卷积 ResNet、U-Net 和 FNO。
• 训练框架：使用 Trainax 进行监督滚动训练。

引用

• 论文：APEBench 是 NeurIPS 2024 会议的论文，链接为 arxiv.org/abs/2411.00180。
• BibTeX 引用： bibtex @article{koehler2024apebench, title={{APEBench}: A Benchmark for Autoregressive Neural Emulators of {PDE}s}, author={Felix Koehler and Simon Niedermayr and R{"u}diger Westermann and Nils Thuerey}, journal={Advances in Neural Information Processing Systems (NeurIPS)}, volume={38}, year={2024} }

许可证

• MIT 许可证：详细信息见 LICENSE.txt。

搜集汇总

数据集介绍

构建方式

APEBench is meticulously crafted as a comprehensive benchmark suite designed to evaluate autoregressive neural emulators for solving partial differential equations (PDEs). The dataset is constructed using JAX, a high-performance computing library, and integrates a differentiable simulation framework that employs efficient pseudo-spectral methods. This framework enables the simulation of 46 distinct PDEs across 1D, 2D, and 3D dimensions. The construction process involves a novel taxonomy for unrolled training and introduces a unique identifier for PDE dynamics that directly relates to the stability criteria of classical numerical methods.

特点

APEBench stands out for its tight integration of a highly efficient pseudo-spectral solver, which is used both for procedural data generation and as a differentiable solver that networks can dynamically interact with during training. This unique feature supports differentiable physics training and neural-hybrid emulators, a capability not found in existing benchmarks. Additionally, APEBench emphasizes rollout metrics to understand temporal generalization, providing insights into the long-term behavior of emulating PDE dynamics.

使用方法

APEBench is designed to facilitate the evaluation of diverse neural architectures and training methodologies. Users can install the benchmark via pip and access it through a seamless integration with JAX. The benchmark includes recipes that rapidly adapt to new architectures and training methodologies, allowing for quick modification of the underlying physics. For visual analysis of the emergent structures, the benchmark is accompanied by a fast volume visualization module, which seamlessly interfaces with the PDE dynamics to provide immediate feedback for emulator development in 2D and 3D.

背景与挑战

背景概述

APEBench, introduced in 2024 by researchers from the Technical University of Munich, is a comprehensive benchmark suite designed to evaluate autoregressive neural emulators for solving partial differential equations (PDEs). Developed under the auspices of the Munich Center for Machine Learning, APEBench leverages the JAX framework and employs efficient pseudo-spectral methods to simulate 46 distinct PDEs across 1D, 2D, and 3D dimensions. The primary objective of APEBench is to facilitate systematic analysis and comparison of learned emulators, proposing a novel taxonomy for unrolled training and introducing a unique identifier for PDE dynamics that directly relates to the stability criteria of classical numerical methods. This benchmark enables the evaluation of diverse neural architectures and supports differentiable physics training and neural-hybrid emulators, emphasizing rollout metrics to understand temporal generalization.

当前挑战

The creation of APEBench presents several challenges. Firstly, the integration of a differentiable simulation framework requires precise handling of numerical methods to ensure stability and accuracy. Secondly, the benchmark must address the challenge of evaluating neural emulators across a wide array of PDE dynamics, each with unique characteristics and difficulty levels. Additionally, the benchmark faces the challenge of ensuring that the neural emulators can generalize over time, which necessitates the use of rollout metrics to assess long-term behavior. Finally, the tight integration of the solver with the neural emulator training process introduces complexities in managing the interplay between classical solvers and neural emulators, particularly in scenarios involving unrolled training and neural-hybrid approaches.

常用场景

经典使用场景

APEBench 是一个用于评估自回归神经偏微分方程（PDE）模拟器的综合基准套件。其经典使用场景包括在1D、2D和3D空间中模拟46种不同的PDE动态，通过高效的伪谱方法实现无缝集成的可微分模拟框架。APEBench 支持对多样化的神经网络架构进行评估，并强调通过可微物理训练和神经混合模拟器来理解PDE动态的长期行为。

实际应用

APEBench 在实际应用中具有广泛的前景，特别是在需要高效模拟复杂物理现象的领域，如气候建模、流体动力学和材料科学。其可微分物理训练和神经混合模拟器的支持，使得在实际问题中能够更准确地预测和模拟PDE动态。

衍生相关工作

APEBench 的推出催生了一系列相关工作，包括对不同神经网络架构的深入研究、可微分物理训练方法的改进以及神经混合模拟器的应用扩展。此外，APEBench 还促进了与其他基准（如PDEBench和PDEArena）的比较研究，进一步推动了神经PDE求解器领域的发展。