Dual-branch neural operator for enhanced out-of-distribution generalization

Abstract

Neural operators, which learn mappings between function spaces, offer an efficient alternative for solving partial differential equations. However, their generalization to out-of-distribution (OOD) parameters often falls short, with accuracy rapidly decreasing outside the training domain. To tackle this issue, we propose a dual-branch neural operator architecture. In this setup, the in-distribution branch performs supervised learning from training data and transfers its learned knowledge to the OOD branch. The OOD branch then uses pseudo-solutions provided by the in-distribution branch to enhance the generalization capability of the neural operator. To ensure the OOD branch prioritizes more reliable information, we introduce a weighting method to control the contributions of different pseudo-solutions. Experimental results show that our method can improve OOD accuracy of neural operators while maintaining in-distribution prediction performance with a few or even no additional OOD observations. The code is publicly available at https://github.com/JcLimath/DB-NO.