Technical Summary: Adaptive Coordinate Transforms for Neural Operators

Problem Statement

Neural operators have demonstrated significant success in solving partial differential equations (PDEs) by learning mappings between function spaces. However, a fundamental limitation persists: most existing architectures (e.g., Fourier Neural Operators, ConvNeXt-UNet, Transolver) are built upon fixed Eulerian coordinates. This static coordinate system creates a spatial misalignment with the evolving physical structures characteristic of many PDEs, such as propagating waves, moving vortices, and advecting interfaces.

This misalignment forces the neural operator to learn highly non-local and complex input-output mappings, significantly increasing learning complexity and hindering generalization. Furthermore, these fixed-coordinate operators exhibit a "smoothness preference," often failing to preserve the correct regularity of solutions and blurring high-frequency physical phenomena like shocks or sharp transitions. Classical PDE analysis addresses these issues through adaptive coordinate transformations (e.g., Arbitrary Lagrangian-Eulerian formulations, moving mesh methods) that dynamically reshape the computational grid to track structures and compress around sharp gradients.

Methodology: The Adaptive Coordinate Transform (ACT) Block

Inspired by classical adaptive techniques, the authors propose the Adaptive Coordinate Transform (ACT) block, a plug-and-play module designed to equip neural operators with data-driven geometric realignment capabilities.

Core Mechanism

The ACT block operates within the neural operator learning pipeline to learn adaptive coordinate systems layer-wise. Instead of explicitly learning a globally invertible coordinate transform, the block predicts an incremental coordinate offset relative to the identity mapping based on the current feature map.

The process involves three steps:

1. Coordinate Offset Prediction: The input feature map is projected and split into independent value tensors. A lightweight multi-scale convolutional module predicts a $d$-dimensional offset field ($\Delta_h$) conditioned on the current features. To ensure stability, raw offsets are bounded by a maximum displacement magnitude $\alpha$ using a $\tanh$ function.
2. Grid Resampling: The predicted offsets define a sampling map $\phi(\xi) = \xi + \Delta(\xi)$. Intermediate features are resampled via differentiable grid sampling at these new locations. This operation preserves the underlying signal values while altering their spatial representation, effectively expressing the same physical quantity in a transformed coordinate system. The authors implement physics-informed boundary conditions (e.g., periodic wrapping) to handle out-of-domain coordinates, preventing boundary artifacts.
3. Feature Merging: The resampled features are concatenated, projected, and merged with the original representation via a residual connection. This ensures stable, incremental updates.

Architectural Integration

ACT blocks are interleaved layer-wise with foundational operator backbones (e.g., FNO, CNextU, Transolver). The hidden features evolve as $h_{l+1} = T^l_\theta \circ G^l_\theta (h_l)$, where $G^l_\theta$ is the latent operator update and $T^l_\theta$ is the ACT block. This continuous geometric realignment prevents spatial misalignment from accumulating as features propagate through the network depth.

Autonomous Role Decoupling

The paper identifies a distinct division of labor among ACT blocks based on their topological location:

• Intermediate Blocks: Act as progressive spatial trackers. They realign evolving features to neutralize spatial misalignment and maintain a stable latent space, effectively simplifying the operator mapping for subsequent layers.
• Terminal Block: Acts as a regularity modulator. By implicitly compressing the local coordinate space, it leverages the Jacobian multiplier to explicitly amplify spatial gradients. This counteracts the inherent smoothness preference of neural networks, aiding in the recovery of sharp physical transitions.

Key Contributions

1. Proposal of the ACT Block: A learnable, geometric realignment module inspired by classic PDE techniques, specifically designed to overcome spatial misalignment and smoothness preferences in fixed-coordinate neural operators.
2. Empirical Validation: Extensive experiments across diverse foundational architectures (CNextU, Transolver, FNO) and six PDE benchmarks (including 2D and 3D datasets like Navier-Stokes, Shallow Water, and MHD turbulence) demonstrate that ACT-augmented models significantly improve predictive accuracy.
3. Mechanistic Insight: The discovery of an autonomous role decoupling phenomenon where intermediate ACTs track advection for latent alignment, while the terminal ACT reconstructs sharp gradients via Jacobian-driven spatial compression.

Experimental Results

The authors evaluated ACT-augmented models against vanilla baselines across six datasets. The results show consistent and significant improvements in Normalized Root Mean Squared Error (NRMSE):

• Transolver: Average relative error reduction of 76.4%.
• FNO: Average relative error reduction of 38.4%.
• CNextU: Average relative error reduction of 30.8%.

Ablation studies on the Navier-Stokes equation confirmed that these gains are not merely due to increased model capacity. A larger FNO with double the parameters showed limited improvement compared to the vanilla model, whereas adding ACT blocks (even a single terminal block) yielded substantial gains. Layer-wise application of ACT provided further improvements, validating the hypothesis that coordinate adaptation, rather than parameter scaling, drives the performance.

Significance and Claims

The paper claims that learning coordinate systems provides a powerful mechanism for enhancing operator learning. By bridging the gap between classical adaptive mesh techniques and modern operator learning, ACT allows networks to better track evolving structures, reduce operator complexity, and dynamically focus on critical features.

The authors maintain a modest stance regarding the scope of these improvements, noting that ACT is most beneficial when spatial misalignment, transport, mixing, or sharp moving interfaces are major sources of difficulty. For problems where dynamics are well-represented in fixed Eulerian coordinates, the improvement may be less significant. Furthermore, while empirical evidence strongly supports the claim that adaptive transformations simplify the effective operator seen by subsequent updates, a formal theoretical characterization of when and why this occurs remains an open direction for future work. The framework is currently limited to structured grids, with extensions to complex geometries and irregular meshes left for future research.