Efficient Neural PDE-Solvers using Quantization Aware Training
This work addresses the practical applicability issue for researchers and engineers using neural PDE solvers, though it is incremental as it adapts existing quantization methods to this domain.
The paper tackles the computational cost bottleneck in neural PDE solvers by applying quantization-aware training to reduce inference costs while maintaining performance, achieving results across four PDE datasets and three architectures with up to three orders of magnitude FLOPs reduction.
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.