Compositional Neural Operators for Multi-Dimensional Fluid Dynamics

Abstract: Partial differential equations (PDEs) govern diverse physical phenomena, yet high-fidelity numerical solutions are computationally expensive and Machine Learning approaches lack generalization. While Scientific Foundation Models (SFMs) aim to provide universal surrogates, typical encoding-decoding approaches suffer from high pretraining costs and limited interpretability. In this paper, we propose Compositional Neural Operators (CompNO) for 2D systems, a framework that decomposes complex PDEs into a library of Foundation Blocks. Each block is a specialized Neural Operator pretrained on elementary physics. This modular library contains convection, diffusion, and nonlinear convection blocks as well as a Poisson Solver, enabling the framework to address the pressure-velocity coupling. These experts are assembled via an Adaptation Block featuring an Aggregator. This aggregator learns nonlinear interactions by minimizing data loss and physics-based residuals driven from governing equations. The proposed approach has been evaluated on the Convection-Diffusion equation, the Burgers' equation, and the Incompressible Navier-Stokes equation. Our results demonstrate that learning from elementary operators significantly improves adaptability, enhances model interpretability and facilitates the reuse of pretrained blocks when adapting to new physical systems.

• The paper demonstrates that pre-training modular neural operator blocks, when composed and fine-tuned with physics-informed loss, achieves significantly reduced errors in fluid dynamics simulations.
• It uses a 'pre-train, assemble, fine-tune' approach to build specialized PFNOs for PDE operators, yielding acceleration up to 35× over traditional solvers.
• The approach improves interpretability and scalability in multi-dimensional CFD, enabling rapid adaptation to new physical systems with fewer trainable parameters.

Compositional Neural Operators for Multi-Dimensional Fluid Dynamics: A Technical Summary

Motivation and Context

The accurate solution of parametric partial differential equations (PDEs), especially in computational fluid dynamics (CFD), remains computationally intensive. Recent advances using neural operators and scientific foundation models (SFMs) have shown promise in learning surrogate mappings between function spaces. However, current approaches—whether via monolithic neural architectures or broad-distribution SFMs—encounter significant limitations, including high pretraining costs, limited interpretability, and suboptimal generalization across diverse physics.

This paper introduces Compositional Neural Operators (CompNO), extending the compositional paradigm to multi-dimensional settings—specifically 2D fluid dynamics. The core idea is to construct a modular library of Foundation Blocks, each pretrained on elementary physics, and aggregate these experts using an Adaptation Block governed by physics-informed loss, thereby addressing system coupling and nonlinearities with structural interpretability.

Architectural Formulation

The CompNO framework is defined by a "pre-train, assemble, fine-tune" methodology. Foundation Blocks specialize in operators such as convection, diffusion, nonlinear convection, and pressure (Poisson), each realized as Parametric Fourier Neural Operators (PFNOs) for discretization invariance and parameterization.

Figure 1: Overview of the CompNO model architecture, highlighting modular Foundation Blocks and the Aggregator for compositional assembly.

Each block is pretrained to solve its designated elemental PDE operator, and its weights are subsequently frozen. The aggregation process concatenates encoded latent features from selected blocks, which are then processed by a multilayer perceptron to realize coupled operator dynamics.

Figure 2: PFNO architecture, illustrating parameter-influenced input transformation for operator generalization.

The final assembly, crucially, utilizes a physics-informed loss that includes both data loss (MAE) and residual-based constraints derived from the governing equations, ensuring long-term consistency and adherence to physical laws during autoregressive inference.

Empirical Results

The evaluation benchmarks CompNO against standard PFNO baselines across several representative fluid systems: Convection-Diffusion, Viscous Burgers’, and Incompressible Navier-Stokes (INS). Pretraining establishes strong accuracy of the Foundation Blocks, with MAEs in the range of $10^{-5}$ to $10^{-4}$, providing robust operator specialization.

When queried for scalar and vectorial system rollouts, CompNO demonstrates superior performance on both velocity and pressure prediction metrics relative to PFNOs. Notably:

• For scalar Convection-Diffusion ($Pe \in [1, 250]$): CompNO yields a velocity MAE of $3\times10^{-4}$, compared to PFNO’s $9\times10^{-3}$, while the physics loss is $0.05$ versus PFNO’s $24$.
• For vectorial INS ($Re \in [1, 100]$): CompNO attains velocity MAE of $4\times10^{-3}$ and pressure MAE of $6\times10^{-3}$, outperforming PFNO's $10^{-4}$0 and $10^{-4}$1 respectively.

These results are achieved with substantially fewer trainable parameters and improved numerical stability across iterative rollouts.

Figure 3: Comparative MAE for INS variables, highlighting error stability during autoregressive inference.

CompNO's autoregressive predictions display linear (rather than exponential) error accumulation, indicating numerical stability and physical validity even in long-term predictions.

Regarding inference speed, CompNO achieves 12–35× acceleration over classical solvers, especially notable for INS problems where typical iterative solution steps are reduced to single-pass modular aggregation.

Practical and Theoretical Implications

The modularity and compositionality of CompNO directly address several core challenges in computational physics ML:

• Generalization: By assembling pretrained operators, CompNO enables rapid adaptation to new PDE systems without retraining the entire model, broadening the applicability of SFMs beyond their current limitations.
• Interpretability: Each block corresponds to a physically interpretable mechanism, permitting transparent reasoning about solution structure and operator interactions.
• Scalability: The approach is extensible—additional physical regimes are tractable via inclusion of new Foundation Blocks—potentially facilitating deployment to broader, coupled multi-physics domains.

Moreover, CompNO's physically grounded aggregation via residual minimization maintains adherence to fundamental laws (e.g., incompressibility, conservation of momentum), a crucial property for industrial surrogate modeling and scientific discovery tasks.

Future Outlook

The CompNO framework paves the way for extensible, reusable scientific foundation models that can be incrementally augmented to cover diverse classes of governing PDEs. Expected future advancements include:

• Expansion to higher-dimensional domains and more complex geometries.
• Integration of more elaborate operator libraries for broad-spectrum multi-physics applications.
• Automated assembly strategies for block selection and adaptation, possibly guided by symbolic or data-driven analysis of target equations.
• Enhanced interpretability via visualization and probing of latent embeddings for physical mechanism discovery.

Conclusion

Compositional Neural Operators represent a robust, scalable, and interpretable SFMs framework for multi-dimensional fluid dynamics (2605.11691). By decomposing complex systems into specialized neural operators and employing physics-informed aggregation, CompNO achieves strong accuracy, parameter efficiency, and inference speed while maintaining physical fidelity. The outlined approach holds promise for advancing surrogate modeling, enabling practical real-time simulation, and facilitating theoretical exploration across scientific domains.