METRO: Metric-Enhanced Token Routing Operator

Nodens Koren; Thomas Hofmann; Georgios Kissas

cs.LGOct 4, 2026

METRO: Metric-Enhanced Token Routing Operator

Authors: Nodens Koren, Thomas Hofmann, Georgios Kissas

Organizations: Department of Computer Science ETH Zürich Zürich, Switzerland · Swiss Data Science Center ETH Zürich Zürich, Switzerland

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.

Figures & tables

Panel of Figure 1 — Figure 1 : Overview of the METRO Architecture. ( a ) The processing pipeline: Input features are mapped to centers, transformed by the metric, sliced, and processed. ( b ) Visual demonstration: Comparing standard spherical neighborhoods (Iso) to flow-aligned elliptical neighborhoods (Aniso), showing how the model explicitly attends to local physical structures.

Figure 2 : M-Transolver consistently outperforms baselines across complex geometries (Elasticity, Plasticity, Airfoil Pipe), achieving an average 23.4% relative error reduction .

Model	OOD Reynolds		OOD Angles
Model	$C_{L}$ $\downarrow$	$\rho_{L}$ $\uparrow$	$C_{L}$ $\downarrow$	$\rho_{L}$ $\uparrow$
Simple MLP	0.6205	0.9578	0.4128	0.9572
GraphSAGE ( Hamilton et al., 2017 )	0.4333	0.9707	0.2538	0.9894
PointNet ( Qi et al., 2017 )	0.3836	0.9806	0.4425	0.9784
Graph U-Net ( Gao and Ji, 2019 )	0.4664	0.9645	0.3756	0.9816
MeshGraphNet ( Pfaff et al., 2021 )	1.7718	0.7631	0.6525	0.8927

Table 2 : Generalization performance on Out-of-Distribution (OOD) tasks for the AirfRANS dataset. We evaluate extrapolation capabilities on unseen Reynolds numbers and Angles of Attack.

\downarrow

indicates lower is better;

\uparrow

indicates higher is better (Spearman’s rank correlation).

Model Variant	AirfRANS				AirCraft (Relative $L_{2}$ Error $\downarrow$ )
Model Variant	Volume $\downarrow$	Surf $\downarrow$	$C_{L}$ $\downarrow$	$\rho_{L}$ $\uparrow$	$C_{p}$	$\rho$	$U$	$V$	$W$	$P$
M-Transolver (Full)	0.0014	0.0061	0.0528	0.9991	0.0692	0.0353	0.0060	0.0919	0.0899	0.0389
w/o Anisotropy	0.0023	0.0066	0.0752	0.9988	0.0702	0.0358	0.0071	0.0987	0.0984	0.0402
w/o Quadratic Term	0.0018	0.0090	0.0780	0.9989	0.0718	0.0363	0.0073	0.1024	0.0984	0.0411
w/o All Components	0.0049	0.0105	0.1159	0.9974	0.0791	0.0386	0.0093	0.1125	0.0997	0.0448
Transolver (Ref.)	0.0037	0.0142	0.1030	0.9978	0.0748	0.0379	0.0091	0.1107	0.1016	0.0439

Table 3 : Unified Ablation Study on Geometric Slicing Components. We evaluate the contribution of the Quadratic term (Exponential Locality) and Anisotropy (Directionality) across both the AirfRANS (2D) and AirCraft (3D) datasets. The “w/o All” variant removes both the quadratic distance term and anisotropic precision structure within our metric-routing implementation, yielding a linear isotropic routing variant; the Transolver reference is reported separately.

Panel of Figure 3 — Figure 3 : Learned Attention Topology on Darcy Flow. Visualization of the attention weights for 64 latent slices. (a) The baseline Transolver suffers from high redundancy, with most slices collapsing into identical spherical shapes in the domain center. (b) M-Transolver learns a diverse set of anisotropic shapes (ellipses, boundary layers, curves) that efficiently tile the domain with minimal overlap, demonstrating superior spatial disentanglement.

Panel of Figure 4 — Figure 4 : Attention Robustness in Chaotic Flow (Navier-Stokes). Comparison of learned slice attention at a late timestep ( $T=20$ ). (a) The baseline Transolver slices become highly entangled, attending to the entire domain and losing local specificity. (b) M-Transolver maintains clear separation, with slices focusing on distinct, localized flow features despite the chaotic global dynamics.

Figure 5 : Learned Anisotropy ( $\mathbf{P}_{s}$ ). (L) The Full M-Transolver captures complex directional variance. (R) Without the quadratic term, the model collapses into isotropic horizontal banding.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Model	ShapeNet-Car
Model	Volume $\downarrow$	Surf $\downarrow$	$C_{D}$ $\downarrow$	$\rho_{D}$ $\uparrow$
GraphSAGE ( Hamilton et al., 2017 )	0.0461	0.1050	0.0270	0.9695
PointNet ( Qi et al., 2017 )	0.0494	0.1104	0.0298	0.9583
Graph U-Net ( Gao and Ji, 2019 )	0.0471	0.1102	0.0226	0.9725
MeshGraphNet ( Pfaff et al., 2021 )	0.0354	0.0781	0.0168	0.9840
GNO ( Li et al., 2020 )	0.0383	0.0815	0.0172	0.9834

Appendix

Table 4 : Quantitative results on the ShapeNet-Car benchmark (3D automotive aerodynamics). We compare against geometric baselines and state-of-the-art neural operators. We report relative

L_{2}

error (

\downarrow

) for Volume and Surface fields, and the error/correlation for the Drag Coefficient (

C_{D}

\rho_{D}

). The best baseline results are italicized , and our results are bolded .

Operator		Point Cloud	Regular Grid		Structured Mesh
Operator		Elasticity	Navier-Stokes	Darcy	Plasticity	Airfoil	Pipe
Classic	U-Net ( Ronneberger et al., 2015 )	0.0235	0.1982	0.0080	0.0051	0.0079	0.0065
	ResNet ( He et al., 2016 )	0.0262	0.2753	0.0587	0.0233	0.0391	0.0120
	Swin ( Liu et al., 2021 )	0.0283	0.2248	0.0397	0.0170	0.0270	0.0109
	DeepONet ( Lu et al., 2021 )	0.0965	0.2972	0.0588	0.0135	0.0385	0.0097
Frequency	WMT ( Gupta et al., 2021 )	0.0359	0.1541	0.0082	0.0076	0.0075	0.0077

Appendix

Table 5 : Comparative performance (Relative L2 Error) on standard PDE benchmarks. Datasets are grouped by geometry type.

Panel of Figure 6 — Figure 6 : Impact of Metric Learning on Attention. (a) The standard Transolver uses isotropic Euclidean distances, treating all directions equally. This often results in attention leaking into irrelevant regions (e.g., across boundary layers). (b) M-Transolver adapts the metric tensor $\mathbf{P}$ to the local physics. The attention mechanism naturally elongates along flow streamlines and compresses across normal gradients, resulting in cleaner, physics-aligned feature aggregation as described in Algorithm 1 .

Geometry	Benchmarks	#Dim	#Mesh	#Input	#Output	#Dataset
Point Cloud	Elasticity	2D	972	Structure	Inner Stress	(1000, 200)
Structured Mesh	Plasticity	2D+Time	3,131	External Force	Mesh Displacement	(900, 80)
	Airfoil	2D	11,271	Structure	Mach Number	(1000, 200)
	Pipe	2D	16,641	Structure	Fluid Velocity	(1000, 200)
Regular Grid	Navier-Stokes	2D+Time	4,096	Past Velocity	Future Velocity	(1000, 200)
Regular Grid	Darcy	2D	7,225	Porous Medium	Fluid Pressure	(1000, 200)

Appendix

Table 6 : Summary of experiment benchmarks. The first six datasets are from FNO ( Li et al., 2021 ) and geo-FNO ( Li et al., 2023c ) , AirfRANS is from ( Bonnet et al., 2022 ) , and AirCraft is a high-fidelity industrial dataset. #Mesh records the size of discretized meshes. #Dataset is organized as the number of samples in training and test sets.

Figure 7 : Latent versus physical domain partitioning in vehicle aerodynamics. While standard physical partition of unity methods must allocate independent, redundant spatial clusters to capture spatially disconnected structures—such as the separated flow around bilateral side mirrors—METRO’s latent Mahalanobis routing embeds these semantically identical regions into the same latent neighborhood. When projected back to the physical domain via the preimage $\phi^{-1}$ , a single latent slice naturally encompasses both mirrors. This demonstrates METRO’s ability to achieve efficient global feature localization without being constrained by physical contiguity.

	Training Configuration						Model Configuration
Benchmarks	Loss	Epochs	Init LR	Optimizer	Batch	Scheduler	Layers $L$	Heads	Channels $D$	Slices $G$
Elasticity	Relative L2	500	$10^{-3}$	AdamW	1	CosineAnnealing	8	8	128	64
Plasticity	Relative L2	500	$10^{-3}$	AdamW	4	OneCycle	8	8	128	64
Airfoil	Relative L2	500	$10^{-3}$	AdamW	4	OneCycle	8	8	128	64
Pipe	Relative L2	500	$10^{-3}$	AdamW	4	OneCycle	8	8	128	64
Navier-Stokes	Relative L2	500	$10^{-3}$	AdamW	2	OneCycle	8	8	256	32

Appendix

Table 7 : Training and model configurations for M-Transolver and M-LaMO. Our settings strictly follow the baselines (Wu et al., 2024) to ensure fair comparison.

L_{rL2}

denotes the relative L2 loss,

L_{v}

and

L_{s}

represent volume and surface field losses, and

L_{g}

is the spatial gradient regularization term. Note that AirCraft adopts the same training configuration as AirfRANS.

	OOD Reynolds		OOD Angles
Dataset	Range	Samples	Range	Samples
Training Set	$[3\times 10^{6},5\times 10^{6}]$	500	$[-2.5^{\circ},12.5^{\circ}]$	800
Test Set	$[2\times 10^{6},3\times 10^{6}]\cup[5\times 10^{6},6\times 10^{6}]$	500	$[-5^{\circ},-2.5^{\circ}]\cup[12.5^{\circ},15^{\circ}]$	200

Appendix

Table 8 : Settings for OOD generalization experiments on AirfRANS. The training and test sets contain disjoint ranges for Reynolds numbers and angles of attack.

Precision Structure	Airfoil	Darcy
Diagonal (Standard)	0.0045	0.0045
Diagonal + Rank-1 Update	0.0047	0.0045
Diagonal + Rank-2 Update	0.0048	0.0046

Appendix

Table 9 : Ablation study on the precision matrix structure. We compare our standard Diagonal formulation against adding low-rank non-diagonal updates on the Airfoil and Darcy datasets. The results show that increasing rank complexity yields no performance gain (and even degradation on Airfoil), suggesting that axis-aligned scaling is sufficient and easier to optimize for these standard benchmarks.

Ablations		Transolver L2RE		M-Transolver L2RE
Number of Slices		Elasticity	Darcy	Elasticity	Darcy
Number of Slices	1	0.0148	0.0386	0.0221	0.0379
	8	0.0071	0.0096	0.0052	0.0054
	16	0.0067	0.0067	0.0048	0.0048
	32	0.0067	0.0063	0.0049	0.0047
	64	0.0064	0.0059	0.0044	0.0045

Appendix

Table 10 : Ablation study on the scalability with respect to the number of slices

G

. We compare the performance (Relative

L_{2}

) of the baseline Transolver and our M-Transolver on the Elasticity (Point Cloud) and Darcy (Regular Grid) datasets.

Figure 8 : Data Efficiency Analysis. Convergence of Relative

L_{2}

Error as the number of training samples increases. M-Transolver (red) exhibits superior data efficiency, achieving significantly lower error rates in the low-data regime (200-400 samples) across all three benchmarks.

Figure 9 : Depth Efficiency Analysis. Performance comparison across varying network depths (8 to 40 layers). M-Transolver demonstrates high stability and robustness, maintaining low error rates even with shallower networks, whereas the baseline often requires greater depth to match performance.

Tier (by Eccentricity)	Mean Ecc.	Mean Alignment	Count
Top 10% Most Aniso.	0.9330	0.8000	636
Top 25% Most Aniso.	0.8952	0.7567	1591
Top 50% Most Aniso.	0.8454	0.7265	3182
Bottom 10% (Least)	0.4347	0.6613	636

Appendix

Table 11 : Geometric analysis of attention slices in the Navier-Stokes dataset. M-Transolver learns highly anisotropic receptive fields that are strongly aligned with the fluid flow velocity.

Panel of Figure 10 — Figure 10 : Showcases of the learned pullback metric alignment of M-Transolver. The visualizations demonstrate how the model learns anisotropic receptive fields that align with the physical flow structures.

	Point Cloud			Structured Mesh		Reg. Grid	Unstructured Mesh
Model	Elasticity	Plasticity	Airfoil	Pipe	Navier-Stokes	Darcy	ShapeNet-Car	AirfRANS
Best Baseline	$0.50{\scriptstyle\pm 0.021}$	$0.07{\scriptstyle\pm 0.008}$	$0.41{\scriptstyle\pm 0.013}$	$0.27{\scriptstyle\pm 0.026}$	$7.19{\scriptstyle\pm 0.072}$	$0.49{\scriptstyle\pm 0.025}$	$99.15{\scriptstyle\pm 0.01}$	$99.83{\scriptstyle\pm 0.010}$
	(LaMO)	(LaMO)	(LaMO)	(Transolver++)	(Transolver++)	(Transolver++)	(Transolver)	(LaMO)
Ours	$\mathbf{0.44}{\scriptstyle\pm 0.015}$	$\mathbf{0.03}{\scriptstyle\pm 0.006}$	$\mathbf{0.37}{\scriptstyle\pm 0.017}$	$\mathbf{0.23}{\scriptstyle\pm 0.028}$	$\mathbf{6.78}{\scriptstyle\pm 0.17}$	$\mathbf{0.45}{\scriptstyle\pm 0.01}$	$\mathbf{99.38}{\scriptstyle\pm 0.025}$	$\mathbf{99.94}{\scriptstyle\pm 0.010}$
	(M-Transolver)	(M-LaMO)	(M-LaMO)	(M-Transolver)	(M-Transolver)	(M-Transolver)	(M-Transolver)	(M-LaMO)

Appendix

Table 12 : Standard deviations of our model compared to the best previous baselines across all experiments. We report the mean

\pm

standard deviation over 4 independent runs. Spearman’s rank correlations for AirfRANS and ShapeNet-Car are presented as percentages (scaled by 100). All other values are relative errors (

\times 10^{-2}

). The best baseline model for each task is indicated in parentheses.

Explore similar work

MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries
Aug 10, 2026Zijiang Yang, Xiaomeng Wu, Dongmei FuNeural PDE Solvers PDE Operator Learning
School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging
Advectra: Asymmetric Latent Transport for Non-Stationary Physics
Oct 4, 2026Nodens Koren, Thomas Hofmann, Georgios KissasPDE Operator Learning Neural Operators
Department of Computer Science ETH Zürich Zürich, Switzerland · Swiss Data Science Center ETH Zürich Zürich, Switzerland
RETO: A Rotary-Enhanced Transformer Operator for High-Fidelity Prediction of Automotive Aerodynamics
Apr 30, 2026Bojun Zhang, Huiyu Yang, Yunpeng Wang +4Rotary Positional Embeddings PDE Operator Learning
Department of Mechanics and Aerospace Engineering, Southern University of Science and Technology, Shenzhen, 518055, China · Shenzhen Key Laboratory of Complex Aerospace Flows, Southern University of Science and Technology,Shenzhen 518055, China · Ningbo Key Laboratory of Advanced Manufacturing Simulation, Eastern Institute of Technology, Ningbo, 315200, China +1

Aug 10, 2026, cs.LGJ/K move · Enter open · S save

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

Zijiang Yang, Xiaomeng Wu, Dongmei Fu

School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging

Neural PDE Solvers PDE Operator Learning Neural Operators Optimal Transport

Open paper AskPDF arXiv

Transformer-based neural operators have achieved substantial progress in solving Partial Differential Equations (PDEs) by projecting spatial observations into compact latent tokens and learning physical interactions in latent spaces. However, we reveal that existing learnable projection mechanisms 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, we propose MoNo (Multiscale Optimal Transport Neural Operator), a progressive multiscale neural operator that 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. Extensive experiments demonstrate that MoNo outperforms existing state-of-the-art neural operators in both prediction performance and computational efficiency. Code is available at https://github.com/ZijiangY1116/MoNo.

Abstract

Explore similar work

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

Advectra: Asymmetric Latent Transport for Non-Stationary Physics

RETO: A Rotary-Enhanced Transformer Operator for High-Fidelity Prediction of Automotive Aerodynamics