METRO: Metric-Enhanced Token Routing Operator

2026-10-04

AI Summary

This work addresses the latent slice redundancy entanglement and isotropy limitations induced by linear scoring in neural operators. To overcome these issues, we propose a geometry-aware routing mechanism that introduces a learnable Mahalanobis metric in place of conventional linear projections. This enables latent slices to adaptively construct oriented ellipsoidal receptive fields, facilitating efficient feature sharding over complex meshes. The proposed mechanism functions as a plug-and-play module that integrates seamlessly into Transformer and Mamba backbone architectures. Extensive evaluations on partial differential equation (PDE) benchmarks and industrial tasks demonstrate that our approach significantly enhances solving performance while improving generalization robustness across varying Reynolds numbers and complex geometric configurations.

Abstract

State-of-the-art neural operators scale to complex meshes via slice-and-process architectures, yet many rely on linear compatibility scores for latent tokenization. Under common feature normalization, such scores are equivalent to isotropic Euclidean clustering, while without normalization they induce unbounded linear decision regions. In both cases, they lack slice-specific anisotropic locality, which can lead to redundant and entangled latent slices. To address this, we propose Metric-Enhanced Token Routing Operator (METRO), a geometry-aware routing mechanism that replaces linear projection with a learnable Mahalanobis metric. By enabling each latent slice to learn a local anisotropic tensor, METRO shapes receptive fields into exponentially localized, oriented ellipsoids that naturally align with flow features like boundary layers and wakes. As a drop-in replacement, METRO yields consistent improvements across both Transformer and Mamba backbones. Empirically, our method achieves substantial performance gains on irregular domains, outperforming baselines on both standard PDE benchmarks and complex industrial design tasks. Finally, METRO exhibits enhanced robustness in out-of-distribution regimes across varying Reynolds numbers and geometric configurations.

Problem

• neural operators
• latent tokenization
• anisotropic locality
• token routing
• slice-and-process

Innovation

• Neural Operator
• Mahalanobis Metric
• Token Routing
• Anisotropic Locality
• Partial Differential Equations