Adaptive Mesh Quantization for Neural PDE Solvers

TMLR Paper5469 Authors

Abstract:

Physical systems commonly exhibit spatially varying complexity, presenting a significant challenge for neural PDE solvers. In traditional numerical methods, adaptive mesh refinement addresses this challenge by increasing node density in dynamic regions, thereby allocating more computational resources where needed. However, for graph neural operators, this is not always a feasible or optimal strategy. We therefore introduce a novel approach to this issue: rather than modifying grid resolution, we maintain a fixed mesh while dynamically adjusting the bit-width used by a quantized model. We propose an adaptive bit-width allocation strategy driven by a lightweight auxiliary model that identifies high-loss regions in the input mesh. This enables dynamic resource distribution in the main model, where regions of higher difficulty are allocated increased bit-width, optimizing computational resource utilization. We demonstrate our framework's effectiveness by integrating it with two state-of-the-art models, MP-PDE and GraphViT, to evaluate performance across multiple tasks: 2D Darcy flow, large-scale unsteady fluid dynamics in 2D, steady-state Navier–Stokes simulations in 3D, and a 2D hyper-elasticity problem.
Our framework demonstrates consistent Pareto improvements over uniformly quantized baselines, yielding up to 50\% improvements in performance at the same cost.

Submission Length: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Fred_Roosta1

Submission Number: 5469