HyCOP: Hybrid Composition Operators for Interpretable Learning of PDEs

Abstract: We introduce HyCOP, a modular framework that learns parametric PDE solution operators by composing simple modules (advection, diffusion, learned closures, boundary handling) in a query-conditioned way. Rather than learning a monolithic map, HyCOP learns a policy over short programs - which module to apply and for how long - conditioned on regime features and state statistics. Modules may be numerical sub-solvers or learned components, enabling hybrid surrogates evaluated at arbitrary query times without autoregressive rollout. Across diverse PDE benchmarks, HyCOP produces interpretable programs, delivers order-of-magnitude OOD improvements over monolithic neural operators, and supports modular transfer through dictionary updates (e.g., boundary swaps, residual enrichment). Our theory characterizes expressivity and gives an error decomposition that separates composition error from module error and doubles as a process-level diagnostic.

• The paper introduces a hybrid composition framework that combines learned policies with process-specific primitives for modular and interpretable PDE surrogates.
• It demonstrates strong in-distribution and out-of-distribution performance with reduced training time and error compared to monolithic neural operators.
• The framework enables actionable error decomposition and zero-shot modular transfer, advancing adaptive and diagnostic capabilities in scientific machine learning.

HyCOP: Hybrid Composition Operators for Interpretable Learning of PDEs

Motivation and Background

Hybrid scientific machine learning (SciML) increasingly emphasizes the integration of neural surrogates and classical numerical methods, especially for parametric PDE surrogates in fluid dynamics, climate science, and multi-physics. Traditional monolithic neural operators (e.g., FNO, DeepONet) learn a single black-box functional map, discarding the process decomposition that scientists use for model design, debugging, and extension. Classical operator splitting exploits mechanistic modularity (e.g., advection, diffusion, reaction), but with fixed, non-adaptive schedules. This dichotomy yields two unsatisfying extremes: interpretable yet inflexible fixed splitting, and flexible but opaque neural surrogates.

HyCOP addresses this gap by introducing a modular architecture that learns how to compose a heterogeneous dictionary of process-specific primitives (numerical or learned modules), orchestrated by a small policy network that outputs regime-conditioned composition "programs." This enables interpretable, robust, and modular PDE surrogates that adapt to distribution shift, support zero-shot transfer via dictionary updates, and enable systematic failure diagnosis.

Framework Overview

HyCOP's architecture comprises a dictionary of process primitives (advection, diffusion, reaction, boundary solvers, etc.), each potentially being a numerical integrator, a black-box neural surrogate, or a learned closure. Given a query—defined by initial state, PDE parameters, boundaries, and target time—the policy $\pi_\theta$ (with $\sim$50–100 parameters) emits a finite program: a sequence of primitive-and-duration pairs. This program is executed to predict the solution at arbitrary query times in a single evaluation, circumventing the need for fine-grained autoregressive rollout.

Figure 1: A small learned policy $\pi_\theta$ maps a query to a regime-adaptive composite program of (primitive, duration) pairs, flexibly allocating prediction time across physical processes.

Distinctively, the composition policy conditions on regime descriptors such as Peclet/Froude/Damköhler numbers and high-level statistics, rather than grid-level raw data. Time allocation and primitive scheduling thus adapt smoothly to underlying physics—even far outside the training regime.

Theoretical Contributions

HyCOP introduces compositional operator flows as a new hypothesis class: mappings realized by finite sequences of process-specific sub-flows, parameterized end-to-end by a policy. Importantly, the theory demonstrates:

• Expressivity (Universality with Structure): The composite flow class is universally expressive on compact query sets for stable split systems. Any PDE solution operator (with decomposable processes) can be approximated arbitrarily well by an appropriately configured HyCOP program.
• Error Decomposition: The total surrogate error decomposes into (i) splitting error (from composition policy), (ii) primitive error (from inaccuracies in sub-flows), and (when the physics is incomplete) (iii) dictionary mismatch. This structuring yields actionable diagnostics and modularity.
• Policy Existence and Learnability: For any loss tolerance, there exists a (measurable, continuous) near-optimal policy, attainable by neural networks or Evolution Strategies (ES).
• Out-of-Distribution Robustness: The error decomposition enables OOD generalization guarantees (as a function of the input distribution shift) under Lipschitz continuity assumptions on the query-parameterized solution operator.

Algorithm and Training

HyCOP's policy network, conditioned on query features (parameters, coarse statistics, target time), outputs a sequence of primitive choices (via softmax logits), durations (positive via softplus), and program length (up to a maximum). Execution corresponds to physical integration using the indicated composition, with all primitives dynamically selected. The expected surrogate loss (prediction error vs. trusted solver output) is minimized by ES, leveraging the low-parameter policy to enable practical black-box optimization even with non-differentiable or legacy numerical primitives.

Empirical Results

Surrogate Fitting and Generalization

Across multiple benchmarks—including 2D compressible Navier–Stokes (PDEBench), 2D shallow-water equations (SWE), and 2D advection–diffusion–reaction (ADR)—HyCOP demonstrates:

• Strong ID and OOD Performance: On challenging OOD setups (unseen parameter regimes, initial and boundary condition shifts, longer rollout horizons), HyCOP improves $L^2$ error by over an order of magnitude relative to state-of-the-art monolithic baselines, with especially large OOD robustness and graceful degradation under shift.
• Efficient Training: HyCOP requires $25\times$ fewer forward passes and $10\times$ shorter wall-clock training time than conventional neural operators (e.g., U-Net, AR baselines), due to its shallow policy search and operator-level prediction.
• Conservation and Interpretability: When using conservation-preserving numerical primitives, HyCOP inherits machine-precision invariants (e.g., mass conservation in SWE/NS benchmarks), unlike monolithic neural surrogates.

Boundary and Modularity Transfer

(Dam-break transfer, Figure 2)

Figure 2: HyCOP rapidly adapts to dam-break (shock OOD) via policy adaptation, and compositionally swaps a single boundary primitive for order-of-magnitude further improvement, highlighting modular transfer.

Zero-shot transfer is demonstrated on dam-break SWE: HyCOP, trained only on smooth periodic-bounded flows, achieves $10\times$ lower error than any baseline on discontinuous, wall-bounded dam-breaks (never seen during training). Swapping the boundary primitive alone (periodic $\rightarrow$ wall) without any retraining yields a further $4\times$ gain, localizing the source of error to boundary physics.

Trajectory Prediction

On both ADR and SWE, HyCOP maintains low long-horizon error (both in-distribution and OOD), while neural autoregressive baselines accumulate error and become unstable:

Figure 3: HyCOP maintains accurate, stable predictions over 1/5/10/20-step horizons; monolithic and autoregressive baselines degrade rapidly OOD.

Figure 4: Example long-horizon trajectory for 2D SWE; HyCOP preserves key physical structures and error remains bounded across extended rollouts.

Adaptation, Closure Discovery, and Heterogeneous Integration

HyCOP operationalizes "compose–diagnose–enrich" workflows foundational to physical discovery:

• Unknown-Physics Diagnosis: When transfer reveals systematic errors—e.g., AD (advection-diffusion) pretrained on ADR (advection-diffusion-reaction) data—HyCOP's process-aligned surrogate localizes errors to domains dominated by the missing reaction, guiding targeted closure augmentation.
• Hybrid Dictionaries: When components are known but come from heterogeneous sources (e.g., a pretrained FNO-based AD surrogate and a hand-coded reaction operator), HyCOP orchestrates their interplay via the learned policy, with no retraining required for either module.

Ablations and Robustness

Systematic ablations confirm:

• Policy Size/Efficiency: Even with very small policies, HyCOP robustly outperforms monolithic networks under shift.
• Feature Conditioning: Physics-informed regime descriptors (dimensionless numbers/statistics) are critical for OOD generalization.
• Resolution Adaptivity: Programs learned at low resolution generalize to higher-resolution grids; only primitive cost increases.
• Redundant/Missing Primitives: The policy suppresses redundant modules without loss, and error localization guides interventions for dictionary incompleteness.

Implications, Limitations, and Prospects

HyCOP structurally encodes process decomposition, establishing a rigorous error decomposition that not only quantifies and localizes generalization failures but also enables modular updates and integration of arbitrary (including legacy) components. This reflects, in machine learning operator surrogates, the philosophy of composability long favored in scientific computing—finite-element assembly, multi-physics splitting, programmatic coupling.

Practical limitations include the requirement of process-level domain knowledge to define the primitive dictionary (though superfluous entries are generally harmless), and the present reliance on trustworthy numerical or robust learned primitives. The robustness of neural sub-solvers under extreme shifts remains a critical area for development.

Further integration with symbolic program synthesis, expanded application to multi-physics foundation models, and advances in universally robust learned primitives represent promising directions for extending HyCOP's framework.

Conclusion

HyCOP advances the interpretability, modularity, and adaptability of neural PDE solution operators by learning explicit query-conditioned composition policies over dictionaries of numerical and/or learned primitives. The compositional operator flow architecture and associated theory provide practical and diagnostic benefits—order-of-magnitude OOD robustness, zero-shot modular transfer, and actionable error decompositions—setting a new direction for scientific ML surrogates rooted in the workflows and verification standards of computational science.