Multi-Grid Tensorized Fourier Neural Operator for High- Resolution PDEs
Authors: Jean Kossaifi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Anima Anandkumar
TMLR 2024 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We demonstrate superior performance on the turbulent 2D Navier-Stokes equations where we achieve less than half the error with over 150 compression compared to the FNO baseline. The tensorization combined with the domain decomposition, yields over 150 reduction in the number of parameters and 7 reduction in the domain size without losses in accuracy. In this section, we first introduce the data, experimental setting and implementation details before empirically validating our approach through thorough experiments and ablations. |
| Researcher Affiliation | Collaboration | Jean Kossaifi EMAIL NVIDIA, Nikola Kovachki EMAIL NVIDIA, Kamyar Azizzadenesheli EMAIL NVIDIA, Anima Anandkumar EMAIL Caltech |
| Pseudocode | No | The paper describes the methodology using mathematical equations and prose (Section 3: Methodology) but does not include any explicitly labeled pseudocode or algorithm blocks. |
| Open Source Code | Yes | We release an open-source implementation of our proposed method to run all variations of FNO and MGTFNO, as well as the Navier-Stokes data used in this paper, along with the paper. Our code will be released under the permissive MIT license, as a Python package that is well-tested and comes with extensive documentation, to encourage and facilitate downstream scientific applications. It will be made available along withe the final version of the manuscript. |
| Open Datasets | Yes | We release an open-source implementation of our proposed method to run all variations of FNO and MGTFNO, as well as the Navier-Stokes data used in this paper, along with the paper. |
| Dataset Splits | Yes | We experiment on a dataset of 10K training samples and 2K test samples of the two-dimensional Navier Stokes equation with Reynolds number 500. ... We use 8K samples for training and 2K for testing. |
| Hardware Specification | Yes | All experiments are done on a NVIDIA Tesla V100 GPU. ... To experimentally measure the gains, we created an FNO keeping all the modes, and we benchmarked two network widths, 64 (Tables 4.5 and 4.5) and 128 (Tables 4.5 and 4.5), for a batch size of 16. In both cases, we measured the memory usage of both a regular FNO and our proposed Tucker TFNO, on a single RTX 3080 GPU, for a single forward pass (inference) and forward-backward pass (training)... |
| Software Dependencies | No | We use Py Torch Paszke et al. (2017) for implementing all the models. The tensor operations are implemented using Tensor Ly Kossaifi et al. (2019) and Tensor Ly-Torch Kossaifi (2021). |
| Experiment Setup | Yes | We train all models via gradient backpropagation using a mini-batch size of 16, the Adam optimizer, with a learning rate of 1e-3, weight decay of 1e-4, for 500 epochs, decreasing the learning rate every 100 epochs by a factors of 1/2. The model width is set in all cases to 64 except when specified otherwise (for the Trimmed FNO)... |