CompNO: A Novel Foundation Model approach for solving Partial Differential Equations

TL;DR

CompNO tackles the high cost and limited transferability of monolithic Scientific Foundation Models for PDEs by introducing a compositional framework that pretrains a library of Foundation Blocks for elementary operators, then assembles them with lightweight Adaptation Blocks and a Boundary Condition Operator. This pretrain–assemble–fine-tune paradigm yields data-efficient, interpretable solvers that generalize across parametric regimes and discretizations, while enforcing Dirichlet boundaries exactly. Empirically, CompNO achieves state-of-the-art relative $L^2$ error on linear parametric PDEs and competitive performance on nonlinear Burgers' flows, with robust behavior across varying Pe and Re numbers. The approach offers a scalable pathway toward reusable PDE foundation models, extendable to higher dimensions and more complex multi-physics problems.

Abstract

Partial differential equations (PDEs) govern a wide range of physical phenomena, but their numerical solution remains computationally demanding, especially when repeated simulations are required across many parameter settings. Recent Scientific Foundation Models (SFMs) aim to alleviate this cost by learning universal surrogates from large collections of simulated systems, yet they typically rely on monolithic architectures with limited interpretability and high pretraining expense. In this work we introduce Compositional Neural Operators (CompNO), a compositional neural operator framework for parametric PDEs. Instead of pretraining a single large model on heterogeneous data, CompNO first learns a library of Foundation Blocks, where each block is a parametric Fourier neural operator specialized to a fundamental differential operator (e.g. convection, diffusion, nonlinear convection). These blocks are then assembled, via lightweight Adaptation Blocks, into task-specific solvers that approximate the temporal evolution operator for target PDEs. A dedicated boundary-condition operator further enforces Dirichlet constraints exactly at inference time. We validate CompNO on one-dimensional convection, diffusion, convection--diffusion and Burgers' equations from the PDEBench suite. The proposed framework achieves lower relative L2 error than strong baselines (PFNO, PDEFormer and in-context learning based models) on linear parametric systems, while remaining competitive on nonlinear Burgers' flows. The model maintains exact boundary satisfaction with zero loss at domain boundaries, and exhibits robust generalization across a broad range of Peclet and Reynolds numbers. These results demonstrate that compositional neural operators provide a scalable and physically interpretable pathway towards foundation models for PDEs.

Figures (6)

• Figure 1: Illustration of CompNO architecture. The input consists of the initial function state $F(x, t_0)$, the physical parameter vector $\gamma$ (e.g., velocity $\beta$, viscosity $\nu$), and the Boundary Conditions (BCs). The Foundation Blocks are pre-trained Neural Operators that independently predict the time-evolution corresponding to specific elementary operators (e.g., $\nabla u$, $\Delta u$). The Aggregator is a neural network (linear or non-linear MLP) that combines these embeddings. Finally, the Constraint BC Operator enforces the Dirichlet values $u_{BC}$ exactly at the domain boundaries before producing the final output.
• Figure 2: The graph illustrates the extrapolation of convection equation.
• Figure 3: L2 relative error measured for different values of Pe and Re numbers.
• Figure 4: Visualization of the model prediction with and without BCs.
• Figure 5: Comparison between the Convection-Diffusion solutions and predictions for various Pe at different time steps.