🤖 machine learning

HyCOP: Hybrid Composition Operators for Interpretable Learning of PDEs

HyCOP is a modular, interpretable framework that learns parametric PDE solution operators by composing simple numerical and learned modules via a query-conditioned policy, achieving superior out-of-distribution generalization and transferability compared to monolithic neural operators.

Original authors: Jinpai Zhao, Nishant Panda, Yen Ting Lin, Eirik Valseth, Diane Oyen, Clint Dawson

Published 2026-05-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jinpai Zhao, Nishant Panda, Yen Ting Lin, Eirik Valseth, Diane Oyen, Clint Dawson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict how a complex system evolves over time, such as a flowing river, rising smoke, or spreading heat. In the world of science, these phenomena are referred to as Partial Differential Equations (PDEs).

Traditionally, scientists have used two main methods to solve these equations:

The old way (Numerical solvers): Like a very precise, slow, step-by-step calculator. It is accurate but requires long runtimes.
The new AI way (Monolithic Neural Operators): Like a "black box" superintelligence that looks at the beginning and end of a movie and guesses the middle. It is fast but often fails completely when it encounters a scene it has never seen before (such as a storm instead of a calm day), and it does not tell you why.

HyCOP is a new framework that attempts to combine the best of both worlds. Here is how it works, explained through simple analogies.

1. The "Lego" bricks versus the "Black box"

Imagine a complex physical system (like a river) as a huge, complicated machine.

The old AI approach tries to learn the entire machine as a single, indestructible block. If you change the water temperature or the shape of the riverbank, the AI gets confused because it has memorized the entire block, not the individual parts.
HyCOP treats the machine like a Lego set. It knows that the river consists of separate, understandable parts:
- Advection: The water flows downstream.
- Diffusion: The water spreads out and mixes.
- Reaction: Chemical changes occur in the water.
- Boundary conditions: How the water interacts with the riverbank.

Instead of learning one huge block, HyCOP learns a small "conductor" (a small guideline) that decides which Lego part to use and for how long, based on the current situation.

2. The "Conductor" analogy

Imagine an orchestra.

Monolithic AI is like a soloist trying to play the entire symphony alone. If the music becomes too complex or the key changes, they might stumble.
HyCOP is like a conductor with a small baton. The conductor does not play the instruments; they simply point to the violins (Advection) for 2 seconds, then to the drums (Diffusion) for 1 second, and finally to the brass (Reaction) for 3 seconds.

The "instruments" (the Lego parts) can be:

Classical mathematics: Precise, fast, and reliable (like a metronome).
AI models: Fast approximations of complex parts.
Hybrid: A mix of both.

The conductor (the guideline) is very small (only about 50–100 parameters). It learns to say: "Okay, right now the water is moving fast, so we spend more time on the Advection part." If the situation changes (e.g., a dam breaks), the conductor simply adjusts the schedule. It does not need to relearn how to play the violin; it only changes when it is played.

3. Why this is a big deal

The work claims that HyCOP solves three major headaches:

It is robust (The "stress test"): If you train the AI on calm rivers and then test it on a massive flood (a situation it has never seen), the "black box" AI usually crashes. HyCOP, however, simply adjusts its schedule. It says: "Oh, this is a flood! I need to focus more on the Advection part and less on the Diffusion part." The work shows that it is 10 times more accurate than standard AI models in handling these new, difficult situations.
It is interpretable (The "why" factor): When a black box AI fails, you do not know why. With HyCOP, you can look at the "program" it wrote. You can see: "It spent 80% of its time on the Reaction part and failed there." This tells the scientist exactly where the problem lies, so they can fix just that one part.
It is modular (The "swap" factor): Imagine you have a model for a river with smooth banks, but now you need a model for a river with a rocky wall.
- Old AI: You must discard the entire model and retrain from scratch.
- HyCOP: You simply swap the "bank" Lego part for a "wall" part. The small conductor immediately learns to use the new part, often without any further training. It is like changing the tires on a car without rebuilding the engine.

4. The "Compose-Diagnose-Enrich" loop

The work describes a workflow for cases where scientists do not know everything about a system:

Compose: Build a model with the parts you know.
Diagnose: Run the model. If it fails, the error "lights up" in the specific area where the missing part is (e.g., errors concentrate where a chemical reaction occurs).
Enrich: Add a new part (a "residual" or a new solver) to the Lego box to close that specific gap. You do not retrain the whole thing; you simply teach the conductor how to use the new part.

Summary

HyCOP is an intelligent, tiny manager that orchestrates a team of specialized tools (some math-based, some AI-based) to solve complex physical problems. Instead of trying to memorize the entire answer, it learns how to combine the right tools at the right time. This makes it faster, more accurate in new situations, and much easier to understand and fix than current AI methods.

1. Problem Statement

Scientific Machine Learning (SciML) currently faces a trade-off between monolithic neural operators (e.g., FNO, DeepONet) and classical numerical solvers:

Monolithic Neural Operators: Learn a single black-box mapping from data. Although they perform well within the training distribution, they often fail catastrophically under distribution shifts (OOD), lack interpretability, and provide no insight into why they failed or how specific physical mechanisms could be corrected.
Classical Solvers: Rely on fixed operator-splitting schemes (e.g., Lie-Trotter, Strang-Splitting). Although they are physically interpretable and robust, they use rigid, predefined schemes that do not adapt to specific query regimes (e.g., changing Péclet or Damköhler numbers).
The Gap: There is a lack of frameworks capable of learning the composition policy (i.e., which physical process is applied and for how long), conditioned on the query state, while simultaneously preserving the modularity and interpretability of the constituent physical mechanisms.

2. Methodology: HyCOP Framework

HyCOP (Hybrid Composition Operators) introduces a modular framework that learns PDE solution operators as query-conditioned compositions of reusable primitives.

Core Architecture

Dictionary of Primitives ( $\mathcal{D}$ ): A collection of sub-flow operators representing different physical mechanisms (e.g., advection, diffusion, reaction, boundary handling). These primitives can be:
- Numerical: Classical solvers (e.g., RK4, spectral methods).
- Learned: Pre-trained neural surrogates (e.g., FNO) for specific processes.
- Hybrid: A mixture of both.
Learned Policy ( $\pi_\theta$ ): A small neural network (approx. 50–100 parameters) that receives a query $x = (u_0, \mu, b, \Omega, t)$ $x = (u_{0}, μ, b, Ω, t)$ and outputs a program.
- Input: The policy is conditioned on dimensionless regime features (e.g., Péclet, Damköhler, Froude numbers) and coarse state statistics, ensuring resolution-independent results.
- Output: A sequence of $(primitive\_index, duration)$ pairs.
- Execution: The solution is computed by sequentially composing these sub-flows: $\hat{u}(t) = \hat{\Phi}_{\tau_k}^{(j_k)} \circ \dots \circ \hat{\Phi}_{\tau_1}^{(j_1)} (u_0)$ .
Key Distinction: Unlike "Mixture of Experts" (which mixes outputs) or autoregressive rollouts (which advance in fixed time steps), HyCOP composes flows sequentially with learned, variable durations.

Optimization

Algorithm: The framework uses Evolutionary Strategies (ES) for training.
Rationale: Since primitives can be inherited black-box solvers (non-differentiable), standard backpropagation is impossible. ES treats program execution as a black box and optimizes the policy parameters $\theta$ by sampling perturbations and ranking performance.
Efficiency: The policy is low-dimensional, enabling efficient parallelization and convergence with fewer training iterations compared to monolithic models.

Theoretical Foundation

The paper provides a learning-theoretic analysis of compositional operator flows:

Expressivity: It is proven that finite compositions of sub-flows are dense in the space of PDE solution operators (universal approximation with structure).
Error Decomposition: A structured error bound is derived:
$\text{Total Error} \leq \underbrace{\text{Splitting Error}}_{\text{Policy/Scheme}} + \underbrace{\text{Primitive Error}}_{\text{Module Accuracy}}$
This decomposition enables process-level diagnostics: If the error is high, it can be identified whether the issue lies in the scheduling policy or a specific primitive module.

3. Main Contributions

Framework: A modular architecture that learns PDE operators as query-conditioned compositions, supports hybrid dictionaries (numerical + learned), and enables multi-step predictions without autoregressive rollouts.
Theory: The first learning-theoretic analysis of learned, query-conditioned compositions, providing an error decomposition that separates composition errors from module errors.
Interpretability & Diagnosis: The framework generates interpretable programs (e.g., "apply advection for 20%, diffusion for 6%..."). It enables a "Compose–Diagnose–Enrich" workflow: if a model fails, errors can be localized to specific missing mechanisms, allowing scientists to add specific closures or swap modules without retraining the entire system.
Modularity: Supports zero-shot transfer via dictionary updates (e.g., swapping a periodic boundary primitive for a wall boundary primitive) and adaptation to missing physics by adding residual closures.

4. Experimental Results

The authors evaluated HyCOP on five benchmarks: 2D compressible Navier-Stokes equations (PDEBench), 2D shallow water equations (SWE), 2D advection-diffusion-reaction (ADR), and 1D systems.

Robustness under Out-of-Distribution (OOD):
- HyCOP achieves order-of-magnitude improvements over state-of-the-art monolithic baselines (FNO, DeepONet, PINO) under OOD shifts (parameter extrapolation, boundary changes, initial condition shifts).
- Example: In the 2D SWE dam-break transfer (shocks, wall boundaries, unseen parameters), HyCOP reduced the error by 10-fold compared to the best monolithic baseline. Swapping the boundary primitive further reduced the error by 4-fold without retraining.
Long-term Stability:
- HyCOP remains stable over long time horizons (20-step rollouts), whereas autoregressive baselines accumulate errors rapidly (often failing with Rel. L2 > 1.0). HyCOP composes at the operator level, avoiding stepwise error accumulation.
Efficiency:
- Training: Requires 25 times fewer training iterations and 10 times shorter wall-clock time compared to U-Net/AR baselines.
- Inference: Competitive speed, often faster than autoregressive methods for long time horizons.
Adaptation Scenarios (Regime B):
- Unknown Missing Physics: When trained on Advection-Diffusion (AD) but tested on Advection-Diffusion-Reaction (ADR), HyCOP localized errors in reaction-dominated regions. Adding a learned residual closure and relearning only the policy restored accuracy.
- Heterogeneous Primitives: Successful composition of a pre-trained FNO (for AD) with a textbook numerical solver (for reaction) without retraining either component, outperforming monolithic models trained end-to-end on the entire system.

5. Significance

Paradigm Shift: Shifts SciML from the "black-box replacement" of solvers to the "orchestration" of heterogeneous components. It treats the solver as a composition of known physics and learned corrections.
Trustworthiness: By providing process-level error diagnosis, HyCOP aligns with emerging verification and validation (V&V) frameworks that require interpretable, modular models for scientific trust.
Scientific Workflow: It formalizes the "Compose-Diagnose-Enrich" loop used by human modelers, enabling AI to assist in discovering missing physics or refining models by precisely identifying which mechanism failed.
Scalability: The policy is resolution-agnostic; only the cost of the primitives scales with grid size, making the framework potentially scalable to production resolutions.

In summary, HyCOP bridges the gap between the flexibility of neural operators and the physical rigor of classical solvers, offering a robust, interpretable, and modular approach to learning PDEs that excels in out-of-distribution scenarios and complex multi-physics systems.