Transform Once: Efficient Operator Learning in Frequency Domain

Michael Poli, Stefano Massaroli, Federico Berto, Jinkyoo Park, Tri Dao, Christopher Re, Stefano Ermon

Keywords: neural operators, frequency, transform, differential equation, dynamics, turbulence, fluid flows, PDE, speedup, high-resolution

TL;DR: We present Transform Once (T1), a frequency-domain deep model to solve partial differential equations that is 3x to 10x faster than Fourier Neural Operators while improving in predictive accuracy.

Abstract: Spectrum analysis provides one of the most effective paradigms for information-preserving dimensionality reduction in data: often, a simple description of naturally occurring signals can be obtained via few terms of periodic basis functions. Neural operators designed for frequency domain learning are based on complex-valued transforms i.e. Fourier Transforms (FT), and layers that perform computation on the spectrum and input data separately. This design introduces considerable computational overhead: for each layer, a forward and inverse FT. Instead, this work introduces a blueprint for frequency domain learning through a single transform: transform once (T1). Our results significantly streamline the design process of neural operators, pruning redundant transforms, and leading to speedups of 3 x to 30 that increase with data resolution and model size. We perform extensive experiments on learning to solve partial differential equations, including incompressible Navier-Stokes, turbulent flows around airfoils, and high-resolution video of smoke dynamics. T1 models improve on the test performance of SOTA neural operators while requiring significantly less computation, with over $30\%$ reduction in predictive error across tasks.

Track: Original Research Track

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2211.14453/code)