Quantization-Aware Preconditioning for Physics-Informed Neural Networks

Ill-conditioned physics-informed neural networks (PINNs) residual objectives can slow or destabilize optimization; residual preconditioning reshapes this geometry, while quantization reduces memory and arithmetic cost. These two design choices are usually treated independently, even though the same residual transform acts on both the physical error and the perturbations induced by low-bit weights and activa- tions. We introduce a quantization-aware framework for residual preconditioning in physics-informed neural networks, in which the preconditioner is designed jointly from the physical tangent geometry and measured quantization-induced residual statistics. The resulting quantization-aware preconditioning (QAP) framework combines isotropic and covariance-aware regularization with scalable low-rank estimates of the quantization geometry. Our analysis separates optimization con- traction from quantization injection and shows why a preconditioner that is effective in full precision can become ineffective after quantization. Across ten-seed end- to-end quantization-aware training (QAT) experiments, QAP reduces deployed solution error by up to 2.07× on heterogeneous one-dimensional diffusion and by 1.45× on a 1000:1 two-dimensional checkerboard-diffusion problem relative to no preconditioning; in the two-dimensional setting, the stabilized directional variant also outperforms isotropic QAP. These results identify quantization-induced residual geometry as a preconditioner-design variable. On the evaluated diffusion problems, QAP improves deployed solution accuracy over quantized PINNs trained without preconditioning, at the same weight and activation bitwidths