METRO: Metric-Enhanced Token Routing Operator¶

会议: NeurIPS 2026
arXiv: 2610.05100
论文: arXiv
领域: 其他
关键词: 神经算子, 偏微分方程求解, Token路由, 马氏距离, 各向异性

一句话总结¶

针对传统切片处理神经算子在线性点积路由下退化为各向同性欧氏聚类或无界超平面的几何局限，METRO 引入可学习的马氏距离度量张量，使隐空间切片自适应拉伸为对齐物理流场（边界层与尾迹）的局部化旋转椭球，显著提升复杂几何与非均匀域上的流体及固力预测精度。

研究背景与动机¶

在航空航天、汽车外流场、能源与叶轮机械等工业设计领域中，偏微分方程（PDE）的大规模数值求解是气动外形优化与不确定性量化的基础支柱。传统基于网格的数值求解器计算代价高昂，而近年来发展起来的神经算子虽然具备快速代理仿真的潜力，但在工业复杂几何场景下面临多重核心挑战：物理边界不规则导致无法简单插值到规则笛卡尔网格，工业级离散往往包含海量自由度，且物理现象普遍包含跨越全域的远场尾迹交互以及高度各向异性的局部流场（例如厚度极薄但沿壁面极度拉伸的黏性边界层，以及剧烈震荡的激波面）。

为了实现非结构化网格上的线性可扩展性，以 Transolver 为代表的“切片处理”（Slice-and-Process）范式应运而生。该范式先将全域的大规模网格点软路由聚合到少量潜在切片 Token（\(G \ll N\)），在紧凑隐空间通过 Transformer 或状态空间模型（SSM）进行全局特征交互，再反向投影解码至原始物理分辨率。然而，现有方法在切片路由打分机制上存在本质的几何缺陷：原始 Transolver 采用线性点积计算网格点与原型之间的兼容度。在标准特征归一化（如 LayerNorm 或 \(\ell_2\) 范数约束）的假设下，点积最大化在代数上严格等价于各向同性的欧氏距离最近邻原型分配，导致其注意力和聚类倾向于形成刚性的球形感受野；若缺少归一化，点积打分则退化为无界超平面划分。这两种情况均缺乏切片特异的各向异性局部化能力，极易导致多个切片在空间上严重重叠纠缠，并在物理特征剧烈变化的区域（如近壁剪切层）将无关的远场背景信息混入切片中。

针对这一几何表征错配问题，本文的核心切入点在于重构网格到切片的分配几何学。不同于在整个系统上施加全局特征重加权，本文让每个潜在切片自主学习独立的局部度量张量。核心 idea：提出度量增强型 Token 路由算子 METRO，用带可学习正定精度矩阵的马氏距离对数高斯打分替代传统线性点积，借助编码器映射在物理坐标系诱导出由拉回伪度量支配的定向椭球感受野，以紧凑且解耦的各向异性切片精确捕获边界层与定向尾迹。

方法详解¶

METRO 的核心设计是将算子学习中的切片降维视为隐空间的度量引导软聚类问题。通过在特征空间为每个切片引入可学习的中心向量与精度矩阵，模型在隐空间构建起高斯局部核，使得注意力的等高线自然退化为指数局部化的紧致椭球，进而通过非线性编码器的几何拉回在物理域自适应弯曲与定向。

整体框架¶

整个网络在空间上维持与标准 Slice-and-Process 一致的线性复杂度管线，依次由特征几何投影、各向异性度量分配、切片加权聚合、隐空间骨干交互以及点级反投影重构五大步骤组成：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入网格特征与坐标 X"] --> B["特征与几何投影<br/>X_geo 与 V"]
    B --> C["各向异性马氏距离打分与归一化对数权重"]
    C --> D["切片质量守恒聚合与隐空间全局交互"]
    D --> E["非线性反投影与物理拉回重构"]
    E --> F["输出物理场预测 Y_out"]

具体计算流程如下： 1. 特征投影：输入的非规则网格节点特征与坐标 \(\mathbf{X} \in \mathbb{R}^{N \times D_{\text{in}}}\) 分别经过线性层投影为几何路由特征 \(\mathbf{X}_{\text{geo}} \in \mathbb{R}^{N \times D}\) 与值特征 \(\mathbf{V} \in \mathbb{R}^{N \times C}\)。 2. 度量打分与权重分配：对 \(G\) 个潜在切片，利用当前切片特定的中心 \(\boldsymbol{\mu}_g\) 与对角正定精度矩阵 \(\mathbf{P}_g\)，计算各节点到切片中心的加权距离与体积修正，得到归一化的行分配概率矩阵 \(\mathbf{A}^{\text{row}}\)。 3. 聚合与隐空间处理：利用按切片总质量列归一化的权重矩阵 \(\mathbf{A}^{\text{slice}}\)，将点特征 \(\mathbf{V}\) 压缩成紧凑切片 Token \(\mathbf{T}_{\text{slice}} \in \mathbb{R}^{G \times C}\)；随后在切片级别使用 Transformer 块或 Mamba 等长序列骨干网络进行高效全局物理混合，输出 \(\mathbf{T}_{\text{out}}\)。 4. 反向投影解切片：利用初始的行分配权重 \(\mathbf{A}^{\text{row}}\) 将切片表征插值广播回全部 \(N\) 个网格点，并经由输出线性变换得到最终预测物理场。

关键设计¶

1. 各向异性马氏距离打分与归一化对数权重：消除球形聚类假设与方差失衡

针对点积打分在特征归一化下导致各向同性球形聚类、无法匹配流体大展弦比结构的缺陷，METRO 将每个潜在切片 \(g \in \{1,\dots,G\}\) 参数化为几何中心 \(\boldsymbol{\mu}_g \in \mathbb{R}^D\) 和正定对角精度矩阵 \(\mathbf{P}_g = \text{diag}(p_{g,1},\dots,p_{g,D})\)。为了确保严格的正定性，对角元素由未约束变量通过非线性映射参数化为 \(p_{g,i} = \text{softplus}(\hat{p}_{g,i}) + \varepsilon\)（其中 \(\varepsilon > 0\) 为极小常数）。对于节点 \(n\) 的几何特征 \(\mathbf{x}_n \in \mathbb{R}^D\)，切片打分函数 \(S_{n,g}\) 定义为带有高斯对数行列式惩罚的二次形式：

\[S_{n,g} = -\frac{1}{2} \|\mathbf{x}_n - \boldsymbol{\mu}_g\|_{\mathbf{P}_g}^2 + \frac{1}{2} \log \det(\mathbf{P}_g) = -\frac{1}{2} \sum_{k=1}^D p_{g,k} (x_{n,k} - \mu_{g,k})^2 + \frac{1}{2} \sum_{k=1}^D \log p_{g,k}\]

该设计的精妙之处在于两个数学特性的有机结合：首先，二次型项 \(- \frac{1}{2} \|\mathbf{x}_n - \boldsymbol{\mu}_g\|_{\mathbf{P}_g}^2\) 在 Softmax 激活后表现为指数衰减核 \(\exp(-\frac{1}{2} \|\mathbf{x}_n - \boldsymbol{\mu}_g\|_{\mathbf{P}_g}^2)\)，使得远离几何中心的无关特征被指数级压制，形成了严格局部化的椭球等值面；其次，\(\frac{1}{2} \log \det(\mathbf{P}_g)\) 项扮演了连续高斯分布的体积归一化常数，在行间 Softmax 竞争中，它能够精确平衡方差较大（精度较小）的切片与高精细切片的 logits 幅度，避免宽广切片在路由中以大体积虚假主导权重分配。

2. 切片质量守恒聚合与隐空间全局交互：维持计算效率与物理守恒

在获取到切片打分后，模型通过跨切片 Softmax 建立逐网格点到隐切片的配分关系：\(A_{n,g}^{\text{row}} = \text{softmax}_g(S_{n,g})\)，满足对每个点 \(n\) 恒有 \(\sum_{g=1}^G A_{n,g}^{\text{row}} = 1\)。为了保证在非均匀网格密度下积分的守恒性与无偏性，不能直接用行归一化权重做均值聚合。METRO 显式计算每个切片在离散网格点集上所汇聚的总统计质量 \(m_g = \sum_{n=1}^N A_{n,g}^{\text{row}}\)，并构建列归一化的切片聚合矩阵 \(A_{n,g}^{\text{slice}} = A_{n,g}^{\text{row}} / m_g\)。潜在切片 Token 的生成严格遵循质量加权平均：

\[\mathbf{T}_{\text{slice}} = (\mathbf{A}^{\text{slice}})^\top \mathbf{V} \in \mathbb{R}^{G \times C}\]

在连续测度极限下，该离散和收敛于对输入场的连续加权积分 \(\int_\Omega w_g(\xi) f(\xi) d\nu(\xi) / \int_\Omega w_g(\xi) d\nu(\xi)\)。紧接着，这组小规模切片被输入至任意现代序列骨干网络（例如由多头自注意力构成的 TransformerBlock，或由选择性状态空间构成的 Mamba 模块）进行充分的长程物理混合与非线性演化。由于切片数量 \(G \ll N\) 且对角精度矩阵的度量计算只需逐元素操作，整体路由算子的计算复杂度为严格的 \(\mathcal{O}(N \cdot G \cdot D)\)，与标准线性注意力保持完全一致的计算开销，却具备显著更高的空间利用率。

3. 非线性反投影与物理拉回重构：从隐空间轴对称到物理空间旋转椭球

虽然在隐特征空间中 \(\mathbf{P}_g\) 被约束为对角矩阵，其在特征轴上呈现轴对齐的椭球形态，但这并不限制模型在物理空间中捕捉任意方向倾斜的涡流与壁面结构。METRO 的理论基础证明了，通过非线性特征编码器 \(\phi: \Omega \to \mathbb{R}^D\) 的几何作用，隐空间的马氏度量诱导了物理网格域上的局部拉回伪度量（Pullback Pseudometric）。对于切片中心对应的物理位置 \(\mathbf{u}_0\)，二次距离函数在物理坐标下的局部 Hessian 矩阵严格满足：

\[\nabla^2 q(\mathbf{u}_0) = 2 J_\phi(\mathbf{u}_0)^\top \mathbf{P}_g J_\phi(\mathbf{u}_0)\]

其中 \(J_\phi(\mathbf{u}_0)\) 是特征编码器关于物理坐标的雅可比矩阵（Jacobian）。只要雅可比矩阵具有满列秩，\(J_\phi^\top \mathbf{P}_g J_\phi\) 便在物理网格切空间构成正定局部黎曼度量。即便 \(\mathbf{P}_g\) 纯为对角阵，编码器的非线性耦合也使得诱导出的感受野能够沿流动流线自由旋转与拉伸。最后，经过隐空间处理更新后的切片 Token \(\mathbf{T}_{\text{out}}\)，通过与行分配矩阵相乘 \(\mathbf{Y} = \mathbf{A}^{\text{row}} \mathbf{T}_{\text{out}}\) 广播还原到物理空间所有节点，并由输出线性层映射出最终预测场。

损失函数 / 训练策略¶

模型训练遵循端到端优化策略。对于不同几何尺度的物理任务，采用与基准统一的损失监督： 1. 标准 PDE 与场预测：直接采用相对 \(L_2\) 范数损失 \(\mathcal{L}_{rL2} = \|\hat{\mathbf{u}} - \mathbf{u}\|_2 / \|\mathbf{u}\|_2\)。对于 Darcy Flow 任务额外加入空间梯度正则化项 \(\mathcal{L}_{\text{total}} = \mathcal{L}_{rL2} + 0.1 \mathcal{L}_g\)。 2. 工业外流场任务（AirfRANS 与 AirCraft）：联合优化体流场损失与壁面场损失 \(\mathcal{L}_{\text{total}} = \mathcal{L}_v + \mathcal{L}_s\)；在 ShapeNet-Car 汽车阻力任务中采用 \(\mathcal{L}_{\text{total}} = \mathcal{L}_v + 0.5 \mathcal{L}_s\)。 3. 优化器与学习率调度：采用 Adam 或 AdamW 优化器，基础学习率设为 \(10^{-3}\)，配合 OneCycleLR 或余弦退火策略（CosineAnnealingLR），在 200 到 500 个 epoch 内完成充分收敛。

实验关键数据¶

主实验¶

评估涵盖工业级大尺度外流场基准（AirfRANS 机翼外流、AirCraft 3D 全机气动）与标准 PDE 基准。在 AirfRANS 与 AirCraft 上的对比结果如下（相对 \(L_2\) 误差越低越好，Spearman 秩相关系数 \(\rho_L\) 越高越好）：

数据集 / 物理指标	之前 SOTA (Transolver / LaMO)	本文 M-Transolver	本文 M-LaMO	相对提升幅度
AirfRANS: 体流场误差 (Volume \(\downarrow\))	0.0020 (LaMO)	0.0014	0.0010	+50.0%
AirfRANS: 壁面场误差 (Surf \(\downarrow\))	0.0066 (LaMO)	0.0061	0.0053	+19.7%
AirfRANS: 升力系数误差 (\(C_L \downarrow\))	0.0977 (LaMO)	0.0528	0.0495	+49.3%
AirfRANS: 升力秩相关性 (\(\rho_L \uparrow\))	0.9983 (LaMO)	0.9991	0.9994	显著更强
AirCraft: 压力系数 (\(C_p \downarrow\))	0.0748 (Transolver)	0.0692	0.0734	+7.5%
AirCraft: 流体密度 (\(\rho \downarrow\))	0.0379 (Transolver)	0.0353	0.0358	+6.9%
AirCraft: 主流向速度 (\(U \downarrow\))	0.0091 (Transolver)	0.0060	0.0075	+34.1%
AirCraft: 垂向速度 (\(V \downarrow\))	0.1107 (Transolver)	0.0919	0.1029	+17.0%
AirCraft: 展向速度 (\(W \downarrow\))	0.1016 (Transolver)	0.0899	0.0953	+11.5%
AirCraft: 静压 (\(P \downarrow\))	0.0438 (LaMO)	0.0389	0.0405	+11.2%

在汽车气动基准 ShapeNet-Car 上，M-Transolver 同样将体流场误差从 LaMO 的 0.0212 降至 0.0173（提升 18.4%），并将阻力系数 \(C_D\) 相对误差由 0.0119 降至 0.0103（提升 13.4%），秩相关性 \(\rho_D\) 达 0.9938。在标准 PDE 基准（Elasticity、Plasticity、Airfoil、Pipe、Navier-Stokes、Darcy）中，M-Transolver 在非均匀域上取得了平均 23.4% 的相对误差降低，其中超弹性应力预测误差降至 0.0044，管流速度预测降至 0.0023。

消融实验¶

在统一实验配置下，针对二次马氏距离项（指数局部性 Metric Locality）与各向异性精度结构（方向自适应 Anisotropy）在 AirfRANS（2D 流动）与 AirCraft（3D 复杂几何）上的独立贡献进行消融：

模型配置	AirfRANS Volume \(\downarrow\)	AirfRANS Surf \(\downarrow\)	AirfRANS \(C_L \downarrow\)	AirCraft \(U \downarrow\)	AirCraft \(V \downarrow\)	说明
M-Transolver (完整模型)	0.0014	0.0061	0.0528	0.0060	0.0919	各向异性 + 二次型局部衰减
w/o 各向异性 (退化为球形度量)	0.0023	0.0066	0.0752	0.0071	0.0987	保持二次衰减，但约束 \(\mathbf{P}_g = \alpha_g \mathbf{I}\)
w/o 二次型项 (退化为自适应线性)	0.0018	0.0090	0.0780	0.0073	0.1024	移除 \(\mathbf{x}^\top \mathbf{P} \mathbf{x}\)，仅保留带权线性打分
去除全部组件 (w/o All)	0.0049	0.0105	0.1159	0.0093	0.1125	移除二次项与各向异性，退化为线性各向同性变体
Transolver 基线参考	0.0037	0.0142	0.1030	0.0091	0.1107	原始点积线性切片机制

关键发现¶

几何机制的正交与互补性：消融数据表明，二次型距离项主要主导指数局部性，移除它会导致 AirfRANS 的壁面场误差从 0.0061 恶化至 0.0090（近壁高梯度特征被远场平滑干扰）；而各向异性度量主要负责几何流向解耦，移除它使体流场误差从 0.0014 骤增至 0.0023。两者共同作用产生的增益远超单项叠加。
极端分布外（OOD）外推鲁棒性：在 AirfRANS 未见雷诺数外推任务中，M-Transolver 的 \(C_L\) 相对误差为 0.1761，较原始 Transolver（0.2996）大幅降低 41.2%；在未见攻角（Angle of Attack）旋转外推中，误差由 0.1500 降至 0.0960。各向异性度量矩阵随攻角自适应旋转主轴，有效防止了刚性球形核被迫大半径包含无关背景噪声。
参数对齐下的极致 Token 压缩率：在严格参数匹配测试中（参数量均锁在约 2.83M），Transolver 需要 96 个切片（误差 0.0061，耗时约 13s/epoch），而 M-Transolver 仅需 16 个各向异性切片即可达到 0.0048 误差（耗时仅 8s/epoch），实现了表征紧凑性、高精度与计算提速的三重提升。

亮点与洞察¶

从隐空间各向异性到物理流向拉回：理论分析揭示了局部雅可比拉回机制 \(J_\phi^\top \mathbf{P}_g J_\phi\)，证明即使在隐特征空间使用对角正定度量，经由非线性几何编码器亦可在物理空间呈现出沿流线高度拉伸弯曲的旋转椭球感受野，避免了高阶张量参数优化的病态问题。
隐空间单位分解破除物理连续性局限：传统物理空间的单位分解要求空间邻域连续，面对对称分离涡（如汽车双侧后视镜脱落涡）必须分配冗余核；METRO 在隐语义空间进行划分，使得单个隐切片能够同时投影覆盖物理空间互不相连但动力学相似的离散区域。
骨干网络完全解耦：METRO 作为即插即用的几何路由算子，在 Transformer（M-Transolver）与状态空间模型（M-LaMO）上均实现了稳健的性能飞跃，展现了优异的模块通用性。

局限与展望¶

高阶非对角全张量优化的鞍点挑战：虽然理论上全精度矩阵具有更强的表征能力，但附录消融表明引入 Rank-1/Rank-2 低秩非对角更新在部分基准上反而因参数冗余和优化地形复杂而略微降点，目前仍依赖对角结构配合非线性编码器。
三维超大规模网格的内存压力：虽然时间复杂度保持对点数线性，但在数十万点级的全机精细网格下，多头注意力与路由 logits 的显存峰值仍较为可观，进一步拓展需结合分块局部聚合或稀疏化调度。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次指明线性切片路由的球形欧氏等价缺陷，系统构建了各向异性马氏路由与拉回度量理论。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D/3D、固力/流体、规则/非规则、工业外流场及 OOD 外推等共 9 项核心基准，消融详实。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，从隐空间代数到物理几何图解的映射清晰生动。
价值: ⭐⭐⭐⭐⭐ 为大规模三维几何物理代理模型提供了极具实用价值的各向异性降维基石。