Multi-Head Residual-Gated DeepONet for Coherent Nonlinear Wave Dynamics

Zhiwei Fan; Yiming Pan; Daniel Coca

cs.LGApr 13, 2026

Multi-Head Residual-Gated DeepONet for Coherent Nonlinear Wave Dynamics

Authors: Zhiwei Fan, Yiming Pan, Daniel Coca

Organizations: School of Engineering, Newcastle University, NE1 7RU, United Kingdom · School of Physical Science and Technology, ShanghaiTech University, Shanghai 200031, China

Abstract

Coherent nonlinear wave dynamics are often strongly shaped by a compact set of physically meaningful descriptors of the initial state. Traditional neural operators typically treat the input-output mapping as a largely black-box high-dimensional regression problem, without explicitly exploiting this structured physical context. Common feature-integration strategies usually rely on direct concatenation or FiLM-style affine modulation in hidden latent spaces. Here we introduce a different paradigm, loosely inspired by the complementary roles of state evolution and physically meaningful observables in quantum mechanics: the wave field is learned through a standard DeepONet state pathway, while compact physical descriptors follow a parallel conditioning pathway and act as residual modulation factors on the state prediction. Based on this idea, we develop a Multi-Head Residual-Gated DeepONet (MH-RG), which combines a pre-branch residual modulator, a branch residual gate, and a trunk residual gate with a low-rank multi-head mechanism to capture multiple complementary conditioned response patterns without prohibitive parameter growth. We evaluate the framework on representative benchmarks including highly nonlinear conservative wave dynamics and dissipative trapped dynamics and further perform detailed mechanistic analyses of the learned multi-head gating behavior. Compared with feature-augmented baselines, MH-RG DeepONet achieves consistently lower error while better preserving phase coherence and the fidelity of physically relevant dynamical quantities.

Figures & tables

Figure 1: Physical motivation and architecture of the MH-RG DeepONet framework. Left: For coherent nonlinear dynamics, the future evolution of the field $\psi(x,t)$ is often influenced not only by the full initial state $\psi(x,t_{0})$ , but also by a compact set of physically meaningful descriptors extracted from that initial condition. Examples include peak amplitude, total intensity, center position, spatial variance, spectral bandwidth and so on, which summarize the structured content of the initial field in real and spectral space. Right: The proposed MH-RG DeepONet combines these two sources of information. The full initial state is processed by the Branch network, while the coordinate query $(x,t)$ is processed by the Trunk network. In parallel, the descriptor vector $\phi$ modulates the operator through three residual conditioning pathways: a pre-branch residual modulator, a branch residual gate, and a trunk residual gate. The resulting conditioned branch and trunk embeddings are fused and passed to a multi-head readout to produce the prediction $\hat{\psi}(x,t)$ .

Model	Branch $B(\cdot)$	Trunk $T(\cdot)$	$p$
Vanilla	$(256,512,512,512,256)$	$(2,512,512,512,256)$	$256$
Concat	$(262,512,512,512,256)$	$(2,512,512,512,256)$	$256$
FiLM	$(256,512,512,512,256)$	$(2,512,512,512,256)$	$256$
RG	$(256,512,512,512,256)$	$(2,512,512,512,256)$	$256$
MH-RG ( $R$ heads)	$(256,512,512,512,256)$	$(2,512,512,512,256)$	$256$

Table 1: Network architecture specification (layer widths) for the NLSE study. Each tuple

(d_{0},\dots,d_{L})

denotes a feed-forward MLP with Linear layers

d_{0}\!\to\!\cdots\!\to\!d_{L}

and GELU activations between hidden layers. Here

2N_{x}=256

, output dimension

q=2

, and latent width

p=256

for all DeepONet variants. The selected conditioning descriptors are

\phi=[E_{0},A_{\max},x_{c},\sigma_{x}^{2},P_{0},\Delta k]\in\mathbb{R}^{6}

Figure 2: NLSE benchmark: field-level, optimization, and physics-level comparisons. (a) Representative full-field predictions and absolute error maps for Vanilla, Concat, FiLM, RG, and MH-RG ( $R=18$ ). Rows show the predicted intensity $|\hat{\psi}|^{2}$ , the intensity error, the real-part error, and the imaginary-part error. (b) Distribution of per-sample full-field MSE on the test set. (c) Mean convergence curves in normalized target space. (d) Total intensity $E(t)=\int|\hat{\psi}|^{2}dx$ . (e) Complex-plane trajectory of $\hat{\psi}(x_{0},t)$ at $x_{0}=0$ . (f) Peak amplitude $\max_{x}|\hat{\psi}(x,t)|$ . (g) Center of intensity $x_{c}(t)$ . All quantities in (a), (d)-(g) are computed from de-normalized predictions on the full space-time grid. All models are trained using the same optimization protocol: Adam with a learning rate of $10^{-3}$ , batch size 64, 64 randomly sampled space–time queries per trajectory per epoch, and a total of 1500 epochs, gradient clipping with a maximum norm of 1.0.

(a) Full-field reconstruction
Model	Params (M)	Full-field MSE
Vanilla	1.447	$1.242\times 10^{-2}\pm 1.065\times 10^{-4}$
Concat	1.450	$5.202\times 10^{-3}\pm 9.443\times 10^{-5}$
FiLM	1.846	$4.985\times 10^{-3}\pm 6.338\times 10^{-5}$
OFormer adaptation	1.496	$7.059\times 10^{-3}\pm 2.871\times 10^{-3}$
RG	1.587	$3.110\times 10^{-3}\pm 2.914\times 10^{-5}$

Table 2: Comparison on the NLSE benchmark. Panel (a) reports the number of learnable parameters and the test-set full-field complex MSE. Panel (b) reports dataset-level trajectory errors in mass, momentum, peak amplitude, and centre of mass for the five representative models evaluated in the physical-fidelity study. For each quantity

q\in\{M,P,A_{\mathrm{max}},x_{c}\}

, the trajectory error for one test sample is defined as

\varepsilon_{q}=\frac{N_{t}^{-1}\sum_{n=1}^{N_{t}}\left|\widehat{q}(t_{n})-q(t_{n})\right|}{N_{t}^{-1}\sum_{n=1}^{N_{t}}\left|q(t_{n})\right|}.

All results are reported as the mean and standard deviation across three independent seed-level test-set means. Lower values are better.

Control configuration	Conditioning structure or context	Full-field MSE
(a) Complete $2^{3}$ ablation of conditioning placement
$\mathrm{P}_{0}\mathrm{B}_{0}\mathrm{T}_{0}$	All three conditioning pathways disabled	$1.242\times 10^{-2}$
$\mathrm{P}_{0}\mathrm{B}_{1}\mathrm{T}_{0}$	Branch-side residual conditioning only	$8.494\times 10^{-3}$
$\mathrm{P}_{0}\mathrm{B}_{0}\mathrm{T}_{1}$	Trunk-side residual conditioning only	$8.451\times 10^{-3}$
$\mathrm{P}_{0}\mathrm{B}_{1}\mathrm{T}_{1}$	Branch- and trunk-side residual conditioning	$6.912\times 10^{-3}$
$\mathrm{P}_{1}\mathrm{B}_{0}\mathrm{T}_{0}$	Pre-branch modulation only	$3.561\times 10^{-3}$

Table 3: Controlled mechanism studies on the NLSE benchmark. Panel (a) uses the single-head dense RG architecture. Here,

P

B

, and

T

denote the pre-branch modulator, branch residual gate, and trunk residual gate, respectively; a subscript of

1

indicates that the corresponding module is enabled, whereas

0

indicates that it is disabled. Panel (b) uses the low-rank gate parameterization of MH-RG to separate the effects of low-rank compression, multiple readouts, and head-specific conditioning. The learned-context and fixed-random-projection controls use the same head-specific

R=18

model and differ only in how the six-dimensional context vector is constructed.

Figure 3: Mechanistic analysis of the MH-RG model on the NLSE benchmark. (a) Per-head output components for a representative test trajectory. (b) Head-output correlation matrix on the representative sample. (c) Head-ablation analysis, where the leave-one-head-out effect is defined as $\Delta_{r}=\mathrm{MSE}(\widehat{\psi}^{(-r)},\psi)-\mathrm{MSE}(\widehat{\psi},\psi)$ , with $\widehat{\psi}^{(-r)}$ denoting the prediction obtained after removing head $r$ . Positive values indicate that the corresponding head improves the full-field prediction. (d) Sensitivity of each head output to the six initial-state descriptors. (e) Sensitivity of the branch-side gates to the same descriptors. (f) Sensitivity of the trunk-side gates to the same descriptors. (g) Overlap map between the learned low-rank branch and trunk gate subspaces, computed from the shared upsampling matrices.

Figure 4: Noise robustness on the NLSE benchmark. Complex Gaussian noise is added to the observed initial condition at inference time, with perturbation level $\epsilon$ . The same noisy input is used for both the sensor representation and the extracted physical descriptors. Shaded regions indicate small, moderate, strong, and severe noise regimes. Left: Test-set full-field MSE versus noise level. Right: Relative degradation, measured by the ratio between the noisy-input mean MSE and the clean-input mean MSE of each model.

(a) Pre-branch gate sweep with $\alpha_{B}=\alpha_{T}=0.5$
$\alpha_{\mathrm{pre}}$	Full-field MSE
0.5	$2.133\times 10^{-3}\pm 1.956\times 10^{-4}$
1.0	$\mathbf{1.406\times 10^{-3}\pm 1.854\times 10^{-4}}$
1.5	$1.505\times 10^{-3}\pm 7.057\times 10^{-5}$
(b) Latent gate sweep with $\alpha_{\mathrm{pre}}=1.0$ and $\alpha_{B}=\alpha_{T}$
$\alpha_{B}=\alpha_{T}$	Full-field MSE

Table 4: Sensitivity of MH-RG (

R=18

) to the gate factors on the NLSE benchmark. Reported values are test-set full-field MSE aggregated over three independent random seeds (

\mathrm{mean}\pm\mathrm{std}

across seed means). Lower is better.

Model	Params (M)	Full-field MSE
Vanilla	8.719	$1.188\times 10^{-2}\pm 2.978\times 10^{-4}$
Concat	8.722	$1.165\times 10^{-2}\pm 5.714\times 10^{-4}$
FiLM	8.921	$5.976\times 10^{-3}\pm 1.742\times 10^{-4}$
RG	8.854	$4.540\times 10^{-3}\pm 2.495\times 10^{-4}$
MH-RG ( $R=2$ )	8.829	$3.650\times 10^{-3}\pm 8.398\times 10^{-5}$
MH-RG ( $R=4$ )	8.837	$3.180\times 10^{-3}\pm 1.566\times 10^{-4}$

Table 5: Model comparison on the 2D damped Gross-Pitaevskii shift-breather benchmark. All models use the same global-normalization protocol and are trained with Adam and StepLR. We report the number of learnable parameters and the test-set full-field MSE aggregated over three independent random seeds (

\mathrm{mean}\pm\mathrm{std}

across seed means). Lower MSE is better.

Figure 5: 2D damped Gross-Pitaevskii shift-breather benchmark. (a) Representative density snapshots and density-error maps for one test trajectory. The first row shows the ground-truth density $|\psi_{GT}(x,y,t)|^{2}$ . The remaining rows show $\bigl||\hat{\psi}_{\mathrm{pred}}|^{2}-|\psi_{\mathrm{GT}}|^{2}\bigr|$ for Concat, FiLM, and MH-RG ( $R=16$ ) at selected times. (b) Test-set full-field MSE across models. (c) Mean convergence of the normalized test loss across seeds. (d-g) Physical diagnostics computed from the reconstructed full field: total mass $N(t)=\iint|\hat{\psi}(x,y,t)|^{2}\,dx\,dy,$ peak amplitude, RMS radius $R_{\mathrm{rms}}(t)=\left[\frac{\iint(x^{2}+y^{2})|\hat{\psi}(x,y,t)|^{2}\,dx\,dy}{\iint|\hat{\psi}(x,y,t)|^{2}\,dx\,dy}\right]^{1/2},$ and the center-of-mass coordinates.

Deep Operator Networks

Deep Operator Networks learn mappings between functions, enabling fast surrogates for parameterized differential equations and other complex systems.

See all 50 papers

FAST-DeepONet: Factor-Augmented Branch Representations for High-Dimensional PDE Inputs in the Small-Sample Regime
Aug 15, 2026Jiyong Kwon, Bongseok Kim, Guang LinPDE Operator Learning Deep Operator Networks
School of Mechanical Engineering, Purdue University, 585 Purdue Mall, West Lafayette, 47907, IN, USA · Department of Mathematics, Purdue University, 150 N. University Street, West Lafayette, 47907, IN, USA
Deep Embedded Multiplicative DMD for Algebra-Preserving Koopman Learning
Jun 3, 2026Kelan Gray, Finlay Brown, Nicolas Boullé +1Koopman Operator Learning Latent Dynamics Modeling
Department of Mathematics, Imperial College London, London, SW7 2AZ, UK. · Department of Applied Mathematics and Theoretical Physics, University of Cambridge, Cambridge, CB3 0WA, UK.
The Error of Deep Operator Networks Is the Sum of Its Parts: Branch-Trunk and Mode Error Decompositions
Feb 25, 2026Alexander Heinlein, Johannes TarazDeep Operator Networks Neural Operators
Delft Institute of Applied Mathematics, Delft University of Technology, Delft, 2628CD, The Netherlands

Aug 15, 2026J/K move · Enter open · S save

FAST-DeepONet: Factor-Augmented Branch Representations for High-Dimensional PDE Inputs in the Small-Sample Regime

Jiyong Kwon, Bongseok Kim, Guang Lin

School of Mechanical Engineering, Purdue University, 585 Purdue Mall, West Lafayette, 47907, IN, USA · Department of Mathematics, Purdue University, 150 N. University Street, West Lafayette, 47907, IN, USA

PDE Operator Learning Deep Operator Networks PDE Surrogate Modeling Operator Learning

Open paper AskPDF arXiv

Deep operator networks can become statistically unstable when partial differential equation inputs are observed at thousands of strongly correlated sensors but only a small number of operator samples is available. We introduce FAST-DeepONet, a branch representation combining a fixed spectral path with a regularized projection of the orthogonal residual, in which the directional penalty acts on the effective residual map after each of its rows is normalized. On Navier--Stokes flow a plain DeepONet degrades from

0.0394

0.1556

mean relative

L_2

error as the branch grows from

129

8193

coordinates, while FAST-DeepONet stays near

0.04

, so the sensor grid can be refined without a statistical penalty. Across independent test sets for Navier--Stokes flow, Darcy flow, and signed terminal wavefield prediction it lowers mean relative

L_2

error by

4.7\%

37.0\%

with three to seven times fewer trainable parameters. A spectral-only branch sharing the same basis separates the two paths: the fixed spectral path carries the improvement on Navier--Stokes and Darcy, while terminal wave prediction requires the residual path together with its directional penalty. FAST-DeepONet targets coordinate-query architectures and trains on solution values alone.

Abstract

Deep Operator Networks

Explore similar work

FAST-DeepONet: Factor-Augmented Branch Representations for High-Dimensional PDE Inputs in the Small-Sample Regime

Deep Embedded Multiplicative DMD for Algebra-Preserving Koopman Learning

The Error of Deep Operator Networks Is the Sum of Its Parts: Branch-Trunk and Mode Error Decompositions