Timezone: »

 
Speeding up Fourier Neural Operators via Mixed Precision
Renbo Tu · Colin White · Jean Kossaifi · Kamyar Azizzadenesheli · Gennady Pekhimenko · Anima Anandkumar

@
in The Synergy of Scientific and Machine Learning Modelling (SynS & ML) Workshop »
Event URL: https://openreview.net/forum?id=qBIYVLWFf8 »

The Fourier neural operator (FNO) is a powerful technique for learning surrogate maps for partial differential equation (PDE) solution operators. For many real-world applications, which often require high-resolution data points, training time and memory usage are significant bottlenecks. While there are mixed-precision training techniques for standard neural networks, those work for real-valued datatypes and therefore cannot be directly applied to FNO, which crucially operates in (complex-valued) Fourier space. On the other hand, since the Fourier transform is already an approximation (due to discretization error), we do not need to perform the operation at full precision. In this work, we (i) profile memory and runtime for FNO with full and mixed-precision training, (ii) conduct a study on the numerical stability of mixed-precision training of FNO, and (iii) devise a training routine which substantially decreases training time and memory usage (up to 27%), with little or no reduction in accuracy, on the Navier-Stokes and Darcy flow equations. Combined with the recently proposed tensorized FNO (Kossaifi et al., 2023), the resulting model has far better performance while also being significantly faster than the original FNO.

Keywords: mixed precision · neural operators

Author Information

Renbo Tu (University of Toronto)
Colin White (Caltech)
Jean Kossaifi (NVIDIA)
Kamyar Azizzadenesheli (NVIDIA)
Gennady Pekhimenko (Department of Computer Science, University of Toronto)
Anima Anandkumar (Caltech and NVIDIA)

Anima Anandkumar is a Bren Professor at Caltech and Director of ML Research at NVIDIA. She was previously a Principal Scientist at Amazon Web Services. She is passionate about designing principled AI algorithms and applying them to interdisciplinary domains. She has received several honors such as the IEEE fellowship, Alfred. P. Sloan Fellowship, NSF Career Award, Young investigator awards from DoD, Venturebeat’s “women in AI” award, NYTimes GoodTech award, and Faculty Fellowships from Microsoft, Google, Facebook, and Adobe. She is part of the World Economic Forum's Expert Network. She has appeared in the PBS Frontline documentary on the “Amazon empire” and has given keynotes in many forums such as the TEDx, KDD, ICLR, and ACM. Anima received her BTech from Indian Institute of Technology Madras, her PhD from Cornell University, and did her postdoctoral research at MIT and assistant professorship at University of California Irvine.

More from the Same Authors