KD-PINN: Knowledge-Distilled PINNs for ultra-low-latency real-time neural PDE solvers

TL;DR

This paper presents KD-PINN, a knowledge distillation framework that creates compact, ultra-low-latency neural PDE solvers with maintained or improved accuracy, suitable for real-time applications.

Contribution

It introduces a novel knowledge distillation method for PINNs, achieving significant speedups while preserving accuracy across various PDEs.

Findings

Inference speedups of 4.8x to 6.9x on benchmark PDEs.

Accuracy improved by approximately 1% with proper tuning.

Models achieve sub-10 ms latency, enabling real-time PDE solving.

Abstract

This work introduces Knowledge-Distilled Physics-Informed Neural Networks (KD-PINN), a framework that transfers the predictive accuracy of a high-capacity teacher model to a compact student through a continuous adaptation of the Kullback-Leibler divergence. In order to confirm its generality for various dynamics and dimensionalities, the framework is evaluated on a representative set of partial differential equations (PDEs). Across the considered benchmarks, the student model achieves inference speedups ranging from x4.8 (Navier-Stokes) to x6.9 (Burgers), while preserving accuracy. Accuracy is improved by on the order of 1% when the model is properly tuned. The distillation process also revealed a regularizing effect. With an average inference latency of 5.3 ms on CPU, the distilled models enter the ultra-low-latency real-time regime defined by sub-10 ms performance. Finally, this study examines how knowledge distillation reduces inference latency in PINNs, to contribute to the development of accurate ultra-low-latency neural PDE solvers.