Metric aware token routing improves neural mesh processing accuracy
METRO: Metric-Enhanced Token Routing Operator
Machine Learning
Summary
Processing complex meshes in neural networks often uses simple scoring methods that miss important details, causing overlapping and unclear groupings. The authors introduced METRO, a new method that uses a learnable geometric metric to better organize data slices into shapes that fit natural flow features. This approach improves accuracy on standard physics problems and tough industrial designs, and it also performs well on new, unseen conditions. METRO can be easily swapped into existing neural network models to enhance their performance.
What this means in practice
- •For computational fluid dynamics engineers: Improve simulation accuracy of fluid flows on irregular meshes in industrial design workflows using enhanced token routing.
- •For machine learning engineers: Integrate METRO into Transformer and Mamba-based architectures to boost spatial feature learning and robustness in complex domain tasks.
Authors
Nodens Koren, Thomas Hofmann, Georgios Kissas
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.