Reformulating Neural Operators in $d+1$ Dimensions for Embedding Evolution

This paper introduces a novel $d+1$ dimensional Neural Operator framework that models embedding evolution via an auxiliary function dimension, achieving state-of-the-art accuracy and robustness across diverse physical benchmarks while avoiding the computational costs of traditional embedding-scaling approaches.

The Gist

Imagine you are trying to teach a computer to predict how a complex physical system changes over time, like how heat spreads through a metal plate or how water swirls in a storm. In the world of artificial intelligence, these systems are often described by mathematical rules called Partial Differential Equations (PDEs).

For a long time, AI models designed to solve these problems (called Neural Operators) have relied on a strategy similar to "brute force." If the model wasn't accurate enough, the engineers would simply make the model "fatter" by adding more internal channels or layers. It's like trying to carry more water by using a wider bucket, even if the bucket is already heavy and clumsy.

This paper introduces a smarter way to carry the water. Instead of just making the bucket wider, the authors propose adding a new dimension to the bucket itself.

The Core Idea: The "Shadow" Dimension

Think of the physical world (like a 2D map of a city) as a flat sheet of paper. Traditional AI models try to learn the patterns on this sheet by looking at it from above, layer by layer.

The authors, Haoze Song and his team, suggest we shouldn't just look at the paper; we should imagine the paper has a shadow or a ghost dimension attached to it. They call this an "auxiliary dimension" (let's call it the "p-dimension").

• The Old Way: Imagine trying to understand a 3D object by looking at a 2D photo and just squinting harder (adding more pixels) to see the details.
• The New Way (SKNO): Imagine you have a 2D photo, but you also have a special "shadow projector" that casts a shadow of that photo onto a wall next to it. By studying both the photo and its shadow together, you can understand the 3D shape much better without needing a bigger photo.

In this paper, they create a model called SKNO (Schrödingerised Kernel Neural Operator). It treats the data as if it exists in a space with one extra dimension. It doesn't just update the data on the physical map; it updates the data on the map and its shadow simultaneously.

How It Works: The "Two-View" Strategy

The magic of SKNO lies in how it updates this extra dimension. The authors use a clever trick inspired by quantum physics (specifically the Schrödinger equation, though they use it just as a design blueprint, not a physics simulation).

They update the "shadow" data in two different ways at the same time:

1. The Raw View: Looking at the data exactly as it is (like reading a book in normal text).
2. The Fourier View: Looking at the data as a mix of waves and frequencies (like reading the book as a musical score of sound waves).

By combining these two "views" of the shadow dimension, the model can capture complex patterns much more efficiently. It's like having a translator who speaks both "Normal English" and "Poetic English" at the same time; they can understand the nuance of a sentence much better than someone who only speaks one.

The Results: Faster, Smaller, and More Accurate

The team tested this new model on over ten different challenging physics problems, ranging from simple heat equations to highly chaotic 3D fluid explosions (Rayleigh–Taylor instability).

Here is what they found:

• Lower Errors: SKNO consistently made fewer mistakes than the best existing models (like FNO, Transolver, and DeepONet).
• Efficiency: It achieved these results without needing to be "fatter" or more expensive. In fact, it was often faster to train and required less computing power.
• Robustness: Even when the model was tested on data it had never seen before (like predicting weather patterns for a day it wasn't trained on, or at a much higher resolution), it held up better than the competition. It didn't get confused when the "grid" of the data changed size.

The Takeaway

The paper argues that instead of just making AI models bigger and heavier to solve hard physics problems, we should change how they look at the data. By adding a "shadow dimension" and updating the data through two different mathematical lenses (raw and frequency-based), the model learns the underlying rules of physics more naturally.

It's a shift from "throwing more resources at the problem" to "finding a better angle to look at the problem." The result is a model that is not only more accurate but also more elegant and efficient.

Technical Summary

Technical Summary: Reformulating Neural Operators in d + 1 Dimensions for Embedding Evolution

Problem Statement

Neural Operators (NOs) are designed to learn mappings between function spaces, particularly for solving Partial Differential Equations (PDEs). While recent advances have focused on refining kernel parameterizations over the $d$-dimensional physical domain, the evolution of lifted embeddings remains underexplored. Existing architectures typically compensate for insufficient embedding expressivity by brute-force scaling (enlarging embedding width or adding heads). However, this strategy incurs high computational costs: dense channel mixing scales quadratically with embedding width, and head-wise factorization only partially mitigates this by inducing block-diagonal structures that weaken cross-head coupling. The paper identifies a gap in directly designing how embeddings evolve, rather than merely increasing their capacity.

Methodology

The authors propose reformulating the Neural Operator pipeline in $d + 1$ dimensions by introducing an auxiliary function dimension $p$. Instead of evolving embeddings solely over the physical domain $D_x$, the proposed framework evolves latent scalar functions over the product domain $D_x \times D_p$.

The General Framework

1. Lifting: The input field $a(x)$ is lifted to a scalar latent function $v_0(x, p)$ on the product domain. This is achieved via a lifting operator $P$, often implemented as a separated linear map $v_0(x, p) = w^\top(p)a(x)$.
2. $(d+1)$-Dimensional Evolution: The latent function is evolved through a sequence of learnable linear operators $\mathcal{L}$ and nonlinear maps $\sigma$. The core component is a kernel integral operator $\mathcal{K}$ acting on both physical coordinates $x$ and the auxiliary coordinate $p$:
   $$\mathcal{K}_l[v_l](x, p) = \int_{D_x} \int_{D_p} \kappa_l(x, y, p, p') v_l(y, p') \, dp' \, dy$$
3. Recovery: The evolved function $v_L(x, p)$ is mapped back to the output domain via a recovery operator $Q$, typically an integration over $p$: $u_{pred}(x) = \int_{D_p} \chi(p) v_L(x, p) \, dp$.

The Schrödingerised Kernel Neural Operator (SKNO)

The paper instantiates this framework with a Fourier-based model named SKNO. Key design choices include:

• Basis-Diversified Auxiliary Evolution: For each spatial location, the signal along the auxiliary dimension $p$ is updated using two distinct coordinate views:
  1. Raw $p$-coordinate mixing: A linear mixing in the spatial domain of $p$.
  2. Fourier-$p$-coordinate mixing: A spectral mixing in the Fourier domain of $p$.
     This dual-branch structure ($F_p^{-1} \tilde{A}_l F_p + B_l$) allows the model to capture features from both views without simply duplicating the same channel-mixing path.
• Physical Domain Propagation: SKNO employs $(L-1)$ global propagators using Spectral Convolution Operators (diagonalized in the Fourier domain of $x$) and one final local propagator using differential operators to capture local information lost by global spectral methods.
• Residual Connections: The linear blocks include residual connections to facilitate training and stability.

Key Contributions

1. Operator-Level Reformulation: The authors reformulate the NO pipeline to evolve latent functions via kernel integrals over both physical and auxiliary coordinates, establishing an explicit operator-based mechanism for embedding evolution.
2. SKNO Architecture: They propose the Schrödingerised Kernel Neural Operator, which utilizes basis-diversified auxiliary evolution (mixing raw and Fourier-$p$ coordinates) to improve expressivity without brute-force scaling.
3. Comprehensive Evaluation: The model is evaluated on over ten benchmarks ranging from 1D linear equations to highly nonlinear 3D instabilities.
4. Controlled Analysis: The paper provides rigorous comparisons against scaled and ablated baselines to demonstrate that performance gains stem from the architectural design (basis diversity) rather than mere parameter count increases.

Experimental Results

Across benchmarks including the 1D Heat/Advection equations, 1D Burgers, 2D Darcy Flow, 2D Gray-Scott, 2D/3D Navier-Stokes, and 3D Rayleigh-Taylor instability, SKNO consistently achieves the lowest relative $L_2$ error among evaluated baselines (DeepONet, FNO, Transolver, CNO).

• Performance Gains: On 2D incompressible Navier-Stokes ($\nu=10^{-5}$), SKNO reduces relative $L_2$ error by approximately 37.1% compared to FNO. On 2D Gray-Scott, the reduction is 42.1%. On 3D Rayleigh-Taylor, SKNO achieves a 14.3% error reduction.
• Capacity Efficiency: Controlled experiments show that SKNO (A+B) outperforms systematically scaled FNO variants and parallel stacked FNOs with fewer parameters and FLOPs. A "B+B" variant (duplicating the raw-$p$ branch) fails to match the performance of the basis-diversified "A+B" variant, confirming the value of the dual-coordinate view.
• Robustness: SKNO demonstrates superior resolution invariance, maintaining low error under mixed-resolution training and zero-shot super-resolution inference (e.g., training on 128 grids and testing on 8192). It also exhibits strong zero-shot generalization to unseen temporal regimes.
• Efficiency: Despite the added dimension, SKNO maintains competitive training times, often outperforming Transformer-based models like Transolver which suffer from quadratic complexity in embedding size.

Significance and Claims

The paper claims that auxiliary-domain operator evolution is a promising alternative to brute-force embedding scaling. By applying the operator-design principle along an auxiliary coordinate, the model improves expressivity and approximation capabilities without the prohibitive computational costs associated with widening embeddings.

The authors emphasize that the "Schrödingerised" naming serves as design inspiration for structured operator evolution along the auxiliary coordinate, rather than claiming a direct classical numerical acceleration mechanism for the PDEs themselves. The results suggest that the proposed $d+1$ dimensional design offers a more direct and efficient path to improving Neural Operator performance, supported by empirical evidence of lower error, better resolution robustness, and superior capacity efficiency.

The paper concludes by noting that future work should focus on developing quantitative criteria to compare neural operators beyond final test error, specifically investigating how different aggregation designs affect optimization trajectories and the selection of local minima in high-dimensional error landscapes.