MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries

MoNo (Multiscale Optimal Transport Neural Operator) is a progressive multiscale neural operator for solving PDEs on general geometries. The paper reveals that existing learnable projection mechanisms for constructing latent spaces in Transformer-based neural operators cannot ensure stable and balanced assignments from observation points to latent tokens, causing some latent tokens to be over-assigned while others remain underutilized. This limitation further restricts the design of hierarchical architectures, as assignment imbalance is continuously inherited and amplified across latent spaces, eventually causing severe token collapse in deeper spaces.

To address these issues, the paper proposes MoNo, which efficiently solves PDEs on general geometries through stable latent-space construction. At its core is CoTAP (Cross-scale Optimal Transport Assignment and Projection), "a novel latent-space construction method that formulates cross-space assignment between adjacent spaces as an entropy-regularized optimal transport problem, thereby constructing balanced bidirectional projections and stable latent spaces. CoTAP also ensures stable information transfer across multiple latent spaces, further enabling multiscale architectures on general geometries, which in turn support more efficient learning of long-range physical interactions."

The paper's contributions are summarized as follows:

• "We propose MoNo, a progressive multiscale neural operator for solving PDEs on general geometries by constructing a stable hierarchy of latent spaces at multiple compression scales to efficiently learn physical interactions across scales."

• We propose CoTAP, which formulates cross-space assignment as an entropy-regularized optimal transport problem and constructs balanced bidirectional projections, thereby enabling the stable construction of latent spaces.

• Extensive experiments on seven PDE benchmarks show that MoNo outperforms SOTA methods in both prediction performance and computational efficiency.

The method works as follows: Given observation points and input observations, MoNo embeds spatial positions and observations into anchor tokens and physical state tokens. It then progressively builds stable latent spaces at multiple compressed scales to learn physical interactions. For each latent space, CoTAP constructs the cross-space assignment and projection between adjacent spaces by solving an entropy-regularized optimal transport problem with uniform marginal constraints. This produces balanced encoding and decoding projection matrices. CoTAP is solved with log-domain Sinkhorn iterations, implemented as a Triton-based fused solver for efficiency.

MoNo defines two variants: MoNo-light (Dg = 96, Df = 96, latent token numbers 512, 256, 128, 64) and MoNo (Dg = 192, Df = 192, latent token numbers 1024, 512, 256, 128), both with four latent spaces.

Extensive experiments on seven PDE benchmarks demonstrate that MoNo outperforms existing state-of-the-art neural operators in both prediction performance and computational efficiency. On standard benchmarks, MoNo achieves the best performance across all tasks. On the AirfRANS advanced benchmark, MoNo achieves the best performance in predicting the surrounding physical fields, surface physical fields, and lift coefficient. MoNo also shows strong performance in out-of-distribution evaluations and at multiple spatial resolutions. In terms of efficiency, MoNo-light contains only 1.73M parameters and achieves the lowest computational cost, while MoNo requires only 97.26 GFLOPs, which is 19.7% lower than that of Transolver++ when processing 65,536 spatial observation points.

Ablation studies verify the critical role of CoTAP and progressive multiscale modeling. Visualizations show that CoTAP constructs balanced assignment patterns with clear multiscale characteristics, mitigating token collapse in deeper latent spaces. The paper concludes that MoNo outperforms SOTA methods in both prediction performance and computational efficiency, highlighting the critical role of stable latent-space construction in developing Transformer-based neural operators with stronger representational capacity and higher computational efficiency.

Improvements to AI systems:

1. Stable Latent-Space Construction via Optimal Transport: Replace heuristic or learned clustering-based token assignment in Transformer-based neural operators with entropy-regularized optimal transport (CoTAP). This ensures balanced, uniform marginal constraints on token assignments, preventing over-assignment of some tokens and underutilization of others. The improved system can maintain token diversity and avoid collapse even in deep, multi-scale architectures.

2. Progressive Multiscale Hierarchy with Balanced Bidirectional Projections: Implement a hierarchy of latent spaces at multiple compression scales (e.g., 1024→512→256→128 tokens), where each adjacent pair is connected via CoTAP’s balanced encoding and decoding projections. This enables stable information transfer across scales, allowing the system to efficiently learn long-range physical interactions without losing fine-grained details.

3. Triton-Based Fused Sinkhorn Solver for Efficiency: Integrate a log-domain Sinkhorn iteration solver, fused and implemented in Triton, to compute optimal transport assignments in real-time during training and inference. This reduces computational overhead, enabling the system to handle large numbers of observation points (e.g., 65,536) with lower FLOPs (e.g., 19.7% less than Transolver++) while maintaining high accuracy.

4. General-Geometry PDE Solving without Mesh Dependency: Use the CoTAP-based latent spaces to handle arbitrary observation point distributions (non-uniform, irregular geometries). The improved system can directly operate on point clouds or scattered data, eliminating the need for structured meshes or predefined grids, thus broadening applicability to real-world engineering and scientific problems.

5. Out-of-Distribution and Multi-Resolution Generalization: Train the multiscale hierarchy with CoTAP to learn scale-invariant representations. The improved system can generalize to unseen spatial resolutions and out-of-distribution input conditions (e.g., different Reynolds numbers or geometries) without retraining, as demonstrated by MoNo’s strong performance on AirfRANS and multi-resolution benchmarks.

6. Token-Collapse Mitigation in Deep Networks: Use CoTAP’s balanced assignment to prevent latent token collapse in deeper spaces, which is a common failure mode in hierarchical Transformers. The improved system can scale to more latent spaces (e.g., 4+ levels) without degradation, enabling deeper and more expressive models for complex physical phenomena.

What the improved AI system can do:

• Solve partial differential equations (PDEs) on arbitrary, irregular geometries (e.g., airfoils, complex 3D shapes) with state-of-the-art accuracy, outperforming existing neural operators (e.g., Transolver, FNO, GNOT) on benchmarks like Darcy flow, Navier-Stokes, and AirfRANS.

• Predict multiple physical fields simultaneously (e.g., velocity, pressure, surface forces) and derived quantities (e.g., lift coefficient) with high fidelity.

• Operate at multiple spatial resolutions and handle out-of-distribution inputs (e.g., different flow conditions) robustly, making it suitable for real-time simulation and digital twin applications.

• Achieve high computational efficiency: with only 1.73M parameters (MoNo-light) or 97.26 GFLOPs (MoNo), it can process 65,536 points in a single forward pass, enabling deployment on resource-constrained hardware (e.g., edge devices, GPUs with limited memory).

• Scale to larger problem sizes and deeper hierarchies without performance degradation, thanks to stable latent-space construction, enabling future extensions to 3D turbulent flows, multi-physics coupling, and inverse design problems.