Adaptive Mesh-Quantization for Neural PDE Solvers

TL;DR

This paper introduces Adaptive Mesh Quantization, a method that dynamically allocates computational resources in neural PDE solvers by adjusting bit-width based on local complexity, improving efficiency and performance.

Contribution

It presents a novel adaptive quantization strategy driven by a lightweight auxiliary model, enhancing resource allocation in neural PDE solvers for complex physical systems.

Findings

Up to 50% performance improvement at the same computational cost.

Effective integration with state-of-the-art models like MP-PDE and GraphViT.

Consistent Pareto improvements over uniform quantization baselines.

Abstract

Physical systems commonly exhibit spatially varying complexity, presenting a significant challenge for neural PDE solvers. While Graph Neural Networks can handle the irregular meshes required for complex geometries and boundary conditions, they still apply uniform computational effort across all nodes regardless of the underlying physics complexity. This leads to inefficient resource allocation where computationally simple regions receive the same treatment as complex phenomena. We address this challenge by introducing Adaptive Mesh Quantization: spatially adaptive quantization across mesh node, edge, and cluster features, 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.