Efficient Neural PDE-Solvers using Quantization Aware Training

TL;DR

This paper explores the use of quantization-aware training to significantly reduce the computational costs of neural PDE solvers while maintaining their performance, demonstrating broad applicability across datasets and architectures.

Contribution

It introduces quantization-aware training as an effective method to lower inference costs in neural PDE solvers, outperforming classical resolution-based approaches.

Findings

Quantization reduces inference FLOPs by up to three orders of magnitude.

Quantization-aware training maintains performance across multiple datasets and architectures.

Pareto-optimal cost-performance trade-offs are achieved mainly through quantization.

Abstract

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.