发表机构
ETH Zürich; Swiss Data Science Center(苏黎世联邦理工学院; 瑞士数据科学中心)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
METRO提出用可学习的马氏距离度量替代线性投影,实现几何感知的令牌路由,在Transformer和Mamba骨干上均提升性能,并在PDE基准和工业设计任务中优于基线,且对分布外场景更鲁棒。
AI 中文摘要
最先进的神经算子通过切片和处理架构扩展到复杂网格,但许多算子依赖线性兼容性分数进行潜在令牌化。在常见的特征归一化下,此类分数等价于各向同性欧几里得聚类,而在无归一化时,它们会引发无界的线性决策区域。在这两种情况下,它们都缺乏切片特定的各向异性局部性,可能导致冗余且纠缠的潜在切片。为解决此问题,我们提出了度量增强的令牌路由算子(METRO),这是一种几何感知的路由机制,用可学习的马氏距离度量取代线性投影。通过使每个潜在切片学习局部各向异性张量,METRO将感受野塑造成指数局部化的定向椭球体,这些椭球体自然与边界层和尾流等流动特征对齐。作为即插即用的替代组件,METRO在Transformer和Mamba骨干网络上均带来一致的改进。实验上,我们的方法在不规则域上取得了显著的性能提升,在标准PDE基准和复杂工业设计任务上均优于基线。最后,METRO在不同雷诺数和几何配置的分布外场景中展现出增强的鲁棒性。
英文摘要
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.