Adaptive Correction for Ensuring Conservation Laws in Neural Operators

TL;DR

This paper introduces a flexible, learnable correction method for neural operators that ensures the conservation of physical quantities like mass and momentum, improving accuracy and stability without limiting model flexibility.

Contribution

A novel adaptive correction approach that enforces conservation laws in neural operators through a lightweight learnable operator, enhancing performance and adaptability.

Findings

Significantly improves accuracy and stability of neural operators.

Outperforms traditional conservation enforcement techniques.

Maintains model expressiveness while ensuring conservation.

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.