Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

• URL: http://arxiv.org/abs/2307.15034v3
• Date: Sun, 5 May 2024 04:01:31 GMT
• Title: Guaranteed Approximation Bounds for Mixed-Precision Neural Operators
• Authors: Renbo Tu, Colin White, Jean Kossaifi, Boris Bonev, Nikola Kovachki, Gennady Pekhimenko, Kamyar Azizzadenesheli, Anima Anandkumar,
• Abstract summary: We build on intuition that neural operator learning inherently induces an approximation error. We show that our approach reduces GPU memory usage by up to 50% and improves throughput by 58% with little or no reduction in accuracy.
• Score: 83.64404557466528
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Neural operators, such as Fourier Neural Operators (FNO), form a principled approach for learning solution operators for PDEs and other mappings between function spaces. However, many real-world problems require high-resolution training data, and the training time and limited GPU memory pose big barriers. One solution is to train neural operators in mixed precision to reduce the memory requirement and increase training speed. However, existing mixed-precision training techniques are designed for standard neural networks, and we find that their direct application to FNO leads to numerical overflow and poor memory efficiency. Further, at first glance, it may appear that mixed precision in FNO will lead to drastic accuracy degradation since reducing the precision of the Fourier transform yields poor results in classical numerical solvers. We show that this is not the case; in fact, we prove that reducing the precision in FNO still guarantees a good approximation bound, when done in a targeted manner. Specifically, we build on the intuition that neural operator learning inherently induces an approximation error, arising from discretizing the infinite-dimensional ground-truth input function, implying that training in full precision is not needed. We formalize this intuition by rigorously characterizing the approximation and precision errors of FNO and bounding these errors for general input functions. We prove that the precision error is asymptotically comparable to the approximation error. Based on this, we design a simple method to optimize the memory-intensive half-precision tensor contractions by greedily finding the optimal contraction order. Through extensive experiments on different state-of-the-art neural operators, datasets, and GPUs, we demonstrate that our approach reduces GPU memory usage by up to 50% and improves throughput by 58% with little or no reduction in accuracy.

Related papers

• Neural Precision Polarization: Simplifying Neural Network Inference with Dual-Level Precision [0.4124847249415279]
  A floating-point model can be trained in the cloud and then downloaded to an edge device. Network weights and activations are directly quantized to meet the edge devices' desired level, such as NF4 or INT8. We show that neural precision polarization enables approximately 464 TOPS per Watt MAC efficiency and reliability.
  arXiv  Detail & Related papers  (2024-11-06T16:02:55Z)
• Just How Flexible are Neural Networks in Practice? [89.80474583606242]
  It is widely believed that a neural network can fit a training set containing at least as many samples as it has parameters. In practice, however, we only find solutions via our training procedure, including the gradient and regularizers, limiting flexibility.
  arXiv  Detail & Related papers  (2024-06-17T12:24:45Z)
• Linearization Turns Neural Operators into Function-Valued Gaussian Processes [23.85470417458593]
  We introduce a new framework for approximate Bayesian uncertainty quantification in neural operators. Our approach can be interpreted as a probabilistic analogue of the concept of currying from functional programming. We showcase the efficacy of our approach through applications to different types of partial differential equations.
  arXiv  Detail & Related papers  (2024-06-07T16:43:54Z)
• Spectral-Refiner: Fine-Tuning of Accurate Spatiotemporal Neural Operator for Turbulent Flows [6.961408873053586]
  We propose a new Stemporal Neural Operator (SFNO) that learns maps between Bochner spaces, and a new learning framework to address these issues. This new paradigm leverages wisdom from traditional numerical PDE theory and techniques to refine the pipeline of commonly adopted end-to-end neural operator training and evaluations. Numerical experiments on commonly used benchmarks for the 2D NSE demonstrate significant improvements in both computational efficiency and accuracy, compared to end-to-end evaluation and traditional numerical PDE solvers.
  arXiv  Detail & Related papers  (2024-05-27T14:33:06Z)
• Training with Mixed-Precision Floating-Point Assignments [8.5323697848377]
  We generate precision assignments for convolutional neural networks that use less memory. We evaluate our technique on image classification tasks by training convolutional networks on CIFAR-10, CIFAR-100, and ImageNet.
  arXiv  Detail & Related papers  (2023-01-31T08:01:35Z)
• Precision Machine Learning [5.15188009671301]
  We compare various function approximation methods and study how they scale with increasing parameters and data. We find that neural networks can often outperform classical approximation methods on high-dimensional examples. We develop training tricks which enable us to train neural networks to extremely low loss, close to the limits allowed by numerical precision.
  arXiv  Detail & Related papers  (2022-10-24T17:58:30Z)
• Refining neural network predictions using background knowledge [68.35246878394702]
  We show we can use logical background knowledge in learning system to compensate for a lack of labeled training data. We introduce differentiable refinement functions that find a corrected prediction close to the original prediction. This algorithm finds optimal refinements on complex SAT formulas in significantly fewer iterations and frequently finds solutions where gradient descent can not.
  arXiv  Detail & Related papers  (2022-06-10T10:17:59Z)
• How Low Can We Go: Trading Memory for Error in Low-Precision Training [52.94003953419242]
  Low-precision arithmetic trains deep learning models using less energy, less memory and less time. We pay a price for the savings: lower precision may yield larger round-off error and hence larger prediction error. We borrow ideas from meta-learning to learn the tradeoff between memory and error.
  arXiv  Detail & Related papers  (2021-06-17T17:38:07Z)
• Robust Implicit Networks via Non-Euclidean Contractions [63.91638306025768]
  Implicit neural networks show improved accuracy and significant reduction in memory consumption. They can suffer from ill-posedness and convergence instability. This paper provides a new framework to design well-posed and robust implicit neural networks.
  arXiv  Detail & Related papers  (2021-06-06T18:05:02Z)
• Learning Frequency Domain Approximation for Binary Neural Networks [68.79904499480025]
  We propose to estimate the gradient of sign function in the Fourier frequency domain using the combination of sine functions for training BNNs. The experiments on several benchmark datasets and neural architectures illustrate that the binary network learned using our method achieves the state-of-the-art accuracy.
  arXiv  Detail & Related papers  (2021-03-01T08:25:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.