Adaptive Correction for Ensuring Conservation Laws in Neural Operators

Abstract

Physical laws, such as the conversation of mass and momentum, are fundamental principles in many physical systems. Neural operators have achieved promising performance in learning the solutions to those systems, but often fail to ensure conservation. Existing methods typically enforce strict conservation via hand-crafted post-processing or architectural constraints, leading to limited model flexibility and adaptability. In this work, we propose a novel plug-and-play adaptive correction approach to ensure the conservation of fundamental linear and quadratic quantities for neural operator outputs. Our method introduces a lightweight learnable operator to adaptively enforce the target conservation law during training. This method allows the model to flexibly and adaptively correct its output to guarantee strict conservation. We provide a theoretical result showing that our correction method does not hamper the expression ability of neural operators and can potentially achieve lower reconstruction loss than their conservation-constrained counterparts. Our method is evaluated across multiple neural operator architectures and representative PDEs. Extensive experiments show that incorporating our correction method into baseline models significantly improves both accuracy and stability. In addition, the experimental results demonstrate that our approach consistently achieves superior performance over widely used conservation-enforcement techniques on various PDE benchmarks.