CATO: Charted Attention for Neural PDE Operators

2026-05-09

AI Summary

This work addresses the high computational cost and difficulty in capturing intrinsic geometric structures of existing Transformer-based neural operators when solving partial differential equations on complex geometries. The authors propose CATO, a geometry-adaptive, derivative-aware neural operator that learns continuous latent coordinate charts to map input meshes into an implicit atlas space. Within this space, chart-conditioned axial attention efficiently captures long-range dependencies. CATO further incorporates a derivative-aware physical loss—encompassing solution values, gradients, and fluxes—and a mesh-consistent gradient regularization. This approach is the first to jointly leverage learnable, geometry-adaptive atlas spaces and derivative-aware supervision, with theoretical guarantees on controllable approximation error under chart perturbations. Experiments demonstrate that CATO significantly outperforms state-of-the-art methods across all benchmark datasets, achieving an average accuracy improvement of 26.76% while reducing parameter count by 81.98%.

Abstract

Neural operators have emerged as powerful data-driven solvers for PDEs, offering substantial acceleration over classical numerical methods. However, existing transformer-based operators still face critical challenges when modeling PDEs on complex geometries: directly processing over massive mesh points is computationally expensive, while operating in raw discretization coordinates may obscure the intrinsic geometry where physical interactions are more naturally expressed. To address these limitations, we introduce the Charted Axial Transformer Operator (CATO), a geometry-adaptive and derivative-aware neural operator for PDEs on general geometries. Instead of applying attention directly in the physical coordinate system, CATO learns a continuous latent chart that maps mesh coordinates into a learned chart space, where chart-conditioned axial attention efficiently captures long-range dependencies with reduced computational cost. In addition, CATO introduces a derivative-aware physics loss for steady-state PDEs that jointly supervises solution values, mesh-consistent gradients, and an auxiliary flux-like field, improving physical fidelity and reducing oversmoothing. We further provide a theoretical approximation result showing that, under a favorable chart, charted axial attention can represent low-rank axial solution operators with controlled error, and that small chart perturbations induce bounded approximation degradation. CATO achieves the best performance across all evaluated datasets, yielding an average improvement of approximately 26.76\% over the strongest competing baselines while reducing the number of parameters by 81.98\%. These results highlight the effectiveness of learning geometry-adaptive charts and derivative-aware physical supervision for accurate and efficient PDE operator learning.

Problem

• neural operators
• PDEs on complex geometries
• computational efficiency
• geometry representation
• physical fidelity

Innovation

• charted attention
• neural PDE operators
• geometry-adaptive
• derivative-aware loss
• axial transformer