Late Fusion Neural Operators for Extrapolation Across Parameter Space in Partial Differential Equations

TL;DR

This paper introduces Late Fusion Neural Operators, a novel architecture that improves the extrapolation ability of neural operators across unseen parameter regimes in PDEs by disentangling state and parameter learning.

Contribution

The work proposes a new neural operator architecture that separates state dynamics from parameter effects, enhancing generalization in PDE predictions beyond training regimes.

Findings

Outperforms Fourier Neural Operator and CAPE-FNO on four benchmark PDEs.

Achieves an average RMSE reduction of over 70% both in-domain and out-domain.

Demonstrates strong generalization across unseen parameter regimes.

Abstract

Developing neural operators that accurately predict the behavior of systems governed by partial differential equations (PDEs) across unseen parameter regimes is crucial for robust generalization in scientific and engineering applications. In practical applications, variations in physical parameters induce distribution shifts between training and prediction regimes, making extrapolation a central challenge. As a result, the way parameters are incorporated into neural operator models plays a key role in their ability to generalize, particularly when state and parameter representations are entangled. In this work, we introduce the Late Fusion Neural Operator, an architecture that disentangles learning state dynamics from parameter effects, improving predictive performance both within and beyond the training distribution. Our approach combines neural operators for learning latent state representations with sparse regression to incorporate parameter information in a structured manner. Across four benchmark PDEs including advection, Burgers, and both 1D and 2D reaction-diffusion equations, the proposed method consistently outperforms Fourier Neural Operator and CAPE-FNO. Late Fusion Neural Operators achieve consistently the best performance in all experiments, with an average RMSE reduction of 72.9% in-domain and 71.8% out-domain compared to the second-best method. These results demonstrate strong generalization across both in-domain and out-domain parameter regimes.