KD-PINN: Knowledge-Distilled PINNs for ultra-low-latency real-time neural PDE solvers
It addresses the need for ultra-low-latency real-time neural PDE solvers, which is incremental as it applies knowledge distillation to an existing method.
This work tackled the problem of reducing inference latency in physics-informed neural networks (PINNs) for solving partial differential equations (PDEs), achieving speedups of 4.8x to 6.9x with an average latency of 5.3 ms while preserving or slightly improving accuracy by about 1%.
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.