🔢 mathematics

Multi-stage neural operator learning with application for convolutions

This paper introduces two multi-stage neural operator learning frameworks, Deep Collocation Neural Operator (DCNO) and Deep Galerkin Neural Operator (DGNO), which iteratively refine operator approximations through supervised and unsupervised training respectively to achieve machine-precision accuracy and superior efficiency in solving convolution integrals and related multi-input problems.

Original authors: Zhiping Mao, Zhenye Wen, Yong Zhang, Xiaofei Zhao

Published 2026-08-20

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhiping Mao, Zhenye Wen, Yong Zhang, Xiaofei Zhao

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

In the vast landscape of modern science and engineering, many problems boil down to understanding how one thing influences another across a distance. Imagine a crowd of people where every individual exerts a subtle pull or push on everyone else, creating a complex web of interactions. In physics, this might be the gravitational pull between stars, the electric charge between particles, or the way a signal ripples through a medium. Scientists call this a "convolution," a mathematical operation that sums up these distant influences to predict a final outcome, such as the shape of a gravitational field or the strength of an electric potential. For decades, calculating these interactions has been a heavy computational burden. Traditional methods can solve the problem for a specific set of points, but they struggle when asked to provide answers for new scenarios or for points scattered anywhere in space. They are like a mapmaker who can only draw roads on a fixed grid; if you ask for a location between the lines, they must guess, often losing precision.

Researchers have recently turned to artificial intelligence to solve this, training computer programs to learn the rules of these interactions directly. However, standard AI approaches often hit a wall of accuracy, producing results that are good but not precise enough for the most demanding scientific tasks. They are like a student who learns a concept quickly but makes small, persistent errors that add up. A team of scientists from institutions in China has now developed a new way to teach these machines, allowing them to reach a level of precision that was previously out of reach. Their work introduces two distinct learning strategies that act like a master craftsman refining a sculpture: they start with a rough shape and then repeatedly chip away at the imperfections, stage by stage, until the result is nearly perfect.

The researchers, led by Zhiping Mao and colleagues, focused on the specific challenge of teaching neural networks to handle these long-range interactions. They realized that trying to learn the entire solution in one go was the bottleneck. Instead, they proposed a multi-stage approach where the computer learns the main pattern first, and then, in subsequent rounds, it learns only the mistakes, or "residuals," left behind by the previous attempt. Think of it as a painter who first sketches the broad outlines of a landscape, then paints the details, and finally adds the tiny highlights and shadows that bring the image to life. By breaking the task into these progressive steps, the computer builds a richer and more accurate understanding of the physics involved.

They developed two versions of this method to suit different types of problems. The first, which they call the Deep Collocation Neural Operator, is a supervised approach. It works like a student with a textbook and an answer key. The computer is fed pairs of inputs and their correct outputs, generated by traditional, slower solvers. It learns to predict the answer, checks its work against the correct answer, and then trains a new, specialized network to fix the specific errors it made. This cycle repeats, with each new network focusing solely on the mistakes of the previous one, until the total error becomes vanishingly small. The second method, the Deep Galerkin Neural Operator, is designed for situations where the underlying physics can be described by a specific set of governing equations, but where the correct answers are too expensive to generate for training. This version is unsupervised; it does not need an answer key. Instead, it learns by checking if its predictions satisfy the fundamental laws of physics encoded in those equations. It constantly adjusts its internal logic to ensure that the laws are obeyed, refining its solution until it aligns perfectly with the physical reality, even without ever seeing a single pre-calculated example.

To test these ideas, the team applied them to a variety of realistic scenarios, including the calculation of gravitational and electric potentials. In one experiment involving a two-dimensional gravitational field, they found that their multi-stage method could reduce errors to a level so small it is nearly indistinguishable from the limits of the computer's own precision. In single-precision arithmetic, which is the standard for many high-speed calculations, the error dropped to a tiny fraction of a percent, effectively reaching the machine's limit. This was a dramatic improvement over standard methods, which often stalled at much higher error rates. The researchers also tested the system on problems where the interaction rules were complex and did not have a simple equation to describe them, proving that the supervised version could handle these difficult cases with the same high accuracy.

The benefits of this approach extend far beyond just being more accurate. Once the training is complete, the new system can predict the outcome for any point in space almost instantly, without needing to recalculate the entire grid. In a direct comparison with a traditional, highly optimized solver, the new method was thousands of times faster when evaluating the results for a large number of different scenarios. While the initial training took some time, the payoff came in the speed of application. For scientists who need to run thousands of simulations or explore how a system changes as parameters shift, this speedup is transformative. It turns a process that might take hours or days into one that takes seconds, opening the door to exploring complex systems that were previously too computationally expensive to study in detail.

The researchers also demonstrated that their methods could handle even more complex situations where multiple factors change at once. In one test, they trained the system to understand how the outcome changed when both the density of the material and the nature of the interaction kernel varied simultaneously. The system successfully learned to generalize across these changing conditions, maintaining its high accuracy. This suggests that the approach is robust and flexible enough to be applied to a wide range of scientific problems, from modeling the behavior of quantum particles to simulating the flow of fluids.

The work represents a significant step forward in the field of scientific computing, showing that artificial intelligence can be pushed to perform with a level of rigor that matches traditional mathematical methods. By moving away from the "one-shot" learning style that has dominated the field, and embracing a process of iterative refinement, the researchers have shown that machines can learn to solve complex physical problems with near-perfect precision. The findings suggest that the future of scientific simulation may lie not in faster hardware alone, but in smarter ways of teaching computers to learn, allowing them to build knowledge layer by layer until the picture is complete.

Technical Summary: Multi-stage Neural Operator Learning with Application for Convolutions

Problem Statement
Convolution integrals are fundamental in scientific and engineering applications, including electrostatics, gravitational simulations, and signal processing. While the mapping from a density function $\rho$ to a potential $\phi$ is linear, the nonlocal nature of the convolution operator makes direct evaluation computationally expensive, typically requiring $O(N \log N)$ operations for $N$ grid points using fast solvers (e.g., Fast Multipole Methods, hierarchical matrices). Although neural operator learning methods (such as Fourier Neural Operators and DeepONets) offer the promise of fast, mesh-free predictions for arbitrary input functions, existing approaches rely on "one-shot" supervised training. The paper notes that such standard training often yields limited accuracy, making these learned operators less competitive than traditional fast solvers. Furthermore, generating high-fidelity training data for supervised learning incurs significant computational costs. The authors aim to develop high-accuracy operator learning frameworks that can approach machine precision and offer efficiency gains for many-query or parametric variation scenarios.

Methodology
The paper proposes two general multi-stage neural operator learning frameworks that iteratively refine operator approximations by constructing basis operators through multiple training stages guided by residuals. Both methods utilize the Deep Operator Network (DeepONet) architecture as the underlying neural structure.

Deep Collocation Neural Operator (DCNO):
- Paradigm: Supervised learning.
- Mechanism: DCNO is designed for scenarios where input-output data pairs $(\rho, \phi)$ are available. It constructs a sequence of basis operators $\{G_{\theta_j}\}$ iteratively. Starting with a zero operator, the method trains a new DeepONet at each stage $s$ to approximate the residual operator $e_{s-1} = G - G_{s-1}$ , where $G_{s-1}$ is the cumulative approximation from previous stages.
- Training: The basis operator at stage $s$ is trained by minimizing a collocation-based least-squares loss between the normalized residual and the network output. The final approximation is a linear combination of the trained basis operators weighted by their respective residual norms.
- Theoretical Basis: The method relies on the universal approximation capability of DeepONets to progressively reduce the error in the $L^2$ norm, with theoretical guarantees showing strictly decreasing error bounds at each stage.
Deep Galerkin Neural Operator (DGNO):
- Paradigm: Unsupervised learning.
- Mechanism: DGNO is applicable when the target operator can be represented by a Partial Differential Equation (PDE). It leverages the weak formulation of the PDE residual to train basis operators without requiring labeled solution data.
- Training: Similar to DCNO, it builds a sum of basis operators. However, the training objective maximizes the normalized PDE residual functional (based on the weak form) over the input space. The coefficients for the linear combination of basis operators are determined via a Galerkin projection (solving a linear system) at inference time.
- Extension: For whole-space convolution problems, the authors integrate a moment-matching strategy to truncate the domain and convert the problem into a bounded PDE with zero boundary conditions, enabling the use of DGNO.
Multi-Input Extension:
- Both frameworks are extended to handle multi-input scenarios (e.g., varying both the density $\rho$ and the kernel $K$ ) by employing the Multiple-Input Operator Network (MIONet) architecture, which adds additional branch networks to the DeepONet structure.

Key Contributions

Novel Frameworks: The introduction of two general multi-stage frameworks, DCNO (supervised) and DGNO (unsupervised), which enrich the approximation space by progressively constructing neural operator bases.
Theoretical Analysis: Provision of theoretical error estimates demonstrating that the approximation error strictly decreases at each stage, with bounds converging geometrically under sufficient training and sampling.
High-Accuracy Application: Successful application to convolution operator learning, demonstrating that these methods can achieve accuracy approaching machine precision (single-precision float) for both standard and near-singular kernels.
Efficiency: Demonstration that while training incurs overhead, the trained operators provide substantial efficiency gains (orders of magnitude speedup) for inference compared to traditional solvers, particularly for many-query settings.

Numerical Results
The authors validate the methods through extensive numerical experiments:

Poisson Potential (2D): Both DCNO and DGNO were tested on a 2D Poisson problem. DCNO achieved a mean relative $L^2$ error of $5.47 \times 10^{-7}$ , and DGNO achieved $6.17 \times 10^{-7}$ , both approaching single-precision machine limits. The results showed that multi-stage construction significantly outperformed standard one-shot DeepONet training (which yielded errors around $10^{-3}$ ).
Robustness to Noise: DCNO demonstrated robustness when trained on noisy data, with prediction errors remaining below the noise level across various noise magnitudes.
Multiquadrics (1D): For a convolution with a multiquadric kernel (lacking a local PDE formulation), DCNO successfully handled near-singular regimes ( $a \to 0$ ), maintaining high accuracy ( $~10^{-6}$ ) across different regularization scales.
Efficiency Comparison: In a comparison with the traditional Sum-of-Exponentials (SOE) method for 1D convolution, DCNO reduced evaluation time from 670.44s to 0.021s for 1000 test cases, despite a 202s training cost.
Schrödinger-Poisson System: DGNO was applied to replace the Poisson solver in a time-dependent Schrödinger-Poisson simulation. The method maintained stability, mass conservation, and provided accelerated evaluation of the gravitational potential compared to finite element methods.
Multi-Input Scenarios: Experiments with MIONet showed that the multi-stage approach significantly improved accuracy for operators depending on both density and kernel variations, reducing errors from $10^{-2}$ (one-shot) to $10^{-5}$ (multi-stage).

Significance and Claims
The paper claims that the proposed multi-stage frameworks address the accuracy limitations of standard one-shot operator learning. By iteratively refining the approximation space through residual-driven training, the methods achieve accuracy levels comparable to high-precision numerical solvers while retaining the inference speed and mesh-free flexibility of neural operators.

DCNO is presented as a robust solution for problems where labeled data is available but high accuracy is required, or where no governing PDE exists.
DGNO is highlighted as an efficient unsupervised alternative for problems with known PDE formulations, eliminating the need for expensive data generation.
The authors emphasize that these methods offer a competitive addition to state-of-the-art solvers for convolution problems, particularly in scenarios involving parametric variations or numerous queries where traditional solvers become computationally prohibitive.

More like this