# Efficient Neural PDE-Solvers using Quantization Aware Training

arXiv: https://arxiv.org/abs/2308.07350 · HTML: https://mlsanity.com/paper/2308.07350/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 28 references that are in the map.
- First posted on arXiv 14 Aug 2023.

## Builds on (most cited first, 17 of 17)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [Learned Step Size Quantization](https://mlsanity.com/paper/1902.08153/) (2019)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2208.05419/) (2022)
- [Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast](https://mlsanity.com/paper/2211.02556/) (2022)
- [Clifford Neural Layers for PDE Modeling](https://mlsanity.com/paper/2209.04934/) (2022)
- [Trained Quantization Thresholds for Accurate and Efficient Fixed-Point Inference of Deep Neural Networks](https://mlsanity.com/paper/1903.08066/) (2019)
- [Partial Differential Equations Meet Deep Neural Networks: A Survey](https://mlsanity.com/paper/2211.05567/) (2022)
- [Neural Network Quantization with AI Model Efficiency Toolkit (AIMET)](https://mlsanity.com/paper/2201.08442/) (2022)

## Built on by (most cited first, 6 of 6)

- [Feasibility Study on Active Learning of Smart Surrogates for Scientific Simulations](https://mlsanity.com/paper/2407.07674/) (2024)
- [AI for Scientific Discovery is a Social Problem](https://mlsanity.com/paper/2509.06580/) (2025)
- [Ricci flow regularization in latent spaces for the forward learning of partial differential equations](https://mlsanity.com/paper/2401.14591/) (2024)
- [Adaptive Mesh-Quantization for Neural PDE Solvers](https://mlsanity.com/paper/2511.18474/) (2025)
- [Tensor-Compressed and Fully-Quantized Training of Neural PDE Solvers](https://mlsanity.com/paper/2512.09202/) (2025)
- [Hierarchical Spatial Algorithms for High-Resolution Image Quantization and Feature Extraction](https://mlsanity.com/paper/2510.08449/) (2025)

## Abstract (arXiv)

In the past years, the application of neural networks as an alternative to classical numerical methods to solve Partial Differential Equations has emerged as a potential paradigm shift in this century-old mathematical field. However, in terms of practical applicability, computational cost remains a substantial bottleneck. Classical approaches try to mitigate this challenge by limiting the spatial resolution on which the PDEs are defined. For neural PDE solvers, we can do better: Here, we investigate the potential of state-of-the-art quantization methods on reducing computational costs. We show that quantizing the network weights and activations can successfully lower the computational cost of inference while maintaining performance. Our results on four standard PDE datasets and three network architectures show that quantization-aware training works across settings and three orders of FLOPs magnitudes. Finally, we empirically demonstrate that Pareto-optimality of computational cost vs performance is almost always achieved only by incorporating quantization.