跳转至

Reformulating Neural Operators in d+1 Dimensions for Embedding Evolution

会议: NeurIPS 2026
arXiv: 2505.11766
领域: 物理与科学计算 / Physics & Scientific Computing
关键词: Neural Operators, Operator Learning, Embedding Evolution, Partial Differential Equations, Fourier Analysis

一句话总结

本文通过引入辅助嵌入坐标将神经算子重构到 \(d+1\) 维乘积域,提出基底多样化辅助演化机制与薛定谔化核神经算子(SKNO),以结构化算子演化替代蛮力扩增嵌入宽度,在十余项偏微分方程基准上取得最优相对 \(L_2\) 误差。

研究背景与动机

神经算子(Neural Operators, NOs)作为学习无穷维函数空间映射的新兴范式,在偏微分方程(PDE)代理求解中取得了显著成功。已有研究大多聚焦于物理域 \(D_x \subset \mathbb{R}^d\) 上的积分核参数化创新,例如图神经网络、傅里叶谱卷积以及物理注意力机制等。然而,在将输入场升维(lifting)到高维嵌入空间之后,这些嵌入特征在网络深层中究竟如何演化,却长期缺乏原则性的算子级设计。

当特征表达能力不足时,现有主流模型(如 FNO、Transolver)几乎完全依赖蛮力扩大隐层宽度(通道数)或堆叠多头注意力来补偿。这种暴力升宽的做法伴随着昂贵的计算代价:通道密集混合的计算量随通道数二次方增长,而多头注意力虽然通过分块对角化缓解了部分计算量,却削弱了跨头之间的直接耦合;更关键的是,嵌入向量各通道之间的演化机制本质上仍停留在离散矩阵乘法,缺乏连续函数空间算子的结构化先验。

针对这一瓶颈,本文换了一个全新视角:既然物理坐标上的演化可以通过核积分算子连续建模,为何不将隐层嵌入本身也视为连续函数的一个辅助坐标?通过引入辅助嵌入维度 \(p \in D_p \subset \mathbb{R}\),将整个神经算子流水线建立在 \(d+1\) 维乘积域 \(D_x \times D_p\) 上。核心 idea:引入辅助函数维度将神经算子重构为 \(d+1\) 维连续算子,并在辅助坐标上采用原始域与傅里叶对偶域协同混合的基底多样化演化机制(SKNO),以算子级演化替代通道蛮力堆砌。

方法详解

整体框架

在传统神经算子中,输入场 \(a(x) \in \mathbb{R}^{d_a}\) 被逐点映射到离散向量嵌入 \(v(x) \in \mathbb{R}^{d_v}\),随后仅在物理坐标 \(x\) 上执行积分核卷积。本文提出的 \(d+1\) 维神经算子框架将输入标量化升维到乘积域 \(D_x \times D_p\) 上的标量函数 \(v_0(x, p) \in \mathbb{R}\)。整个演化流程由升维算子 \(\mathcal{P}\)、迭代的 \(d+1\) 维线性核模块 \(\mathcal{L}\) 与非线性映射 \(\sigma\)、以及最终的降维恢复算子 \(\mathcal{Q}\) 构成:

\[\mathcal{G}[a] \approx (\mathcal{Q} \circ \sigma_{L-1} \circ \mathcal{L}_{L-1} \circ \cdots \circ \sigma_0 \circ \mathcal{L}_0 \circ \mathcal{P})[a]\]

在具体实现中,薛定谔化核神经算子(Schrödingerised Kernel Neural Operator, SKNO)在物理坐标 \(x\) 上串联 \(L-1\) 个全局谱卷积传播器与 1 个局部微分传播器,同时在辅助坐标 \(p\) 上引入原始域与频域的双重线性混合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入场 a(x)"] --> P["分离升维算子 P<br/>映射到 (d+1) 维标量场 v0(x, p)"]
    P --> Global["全局谱传播器 (L-1 层)<br/>x 维截断谱卷积 + 双残差连接"]
    Global --> Aux["辅助维度演化模块<br/>原始坐标 Bl + 傅里叶对偶坐标 Fp^-1 Al Fp 协同混合"]
    Aux --> Local["局部微分传播器 (第 L 层)<br/>有限差分局部导数积分核"]
    Local --> Q["积分恢复算子 Q<br/>沿辅助坐标 Dp 正交加权求积"]
    Q --> Out["输出场 u(x)"]

关键设计

1. 分离升维与积分恢复:将隐层特征赋予连续坐标意义 现有方法将每个空间点独立映射为高维隐向量,通道间缺乏几何拓扑约束。SKNO 引入辅助维度 \(D_p \subset \mathbb{R}\),采用分离升维算子 \(\mathcal{P}\): $\(v_0(x,p) = w^\top(p) a(x)\)$ 其中 \(w(p) \in \mathbb{R}^{d_a}\) 为可学习的参数权重剖面。在经过 \(L\) 层算子演化得到 \(v_L(x,p)\) 后,恢复算子 \(\mathcal{Q}\) 沿辅助坐标执行连续数值积分求积: $\(u_{\text{pred}}(x) = \int_{D_p} \chi(p) v_L(x,p) \, dp\)$ 实际网络中,\(\chi(p) \in \mathbb{R}^{d_u}\) 通过轻量 MLP 与末层线性积分投影实现。该设计使特征通道获得连续坐标解释,并将最后的预测重构建模为对辅助空间的投影度量,为跨层特征解耦与模式能量聚集奠定了基础。

2. 辅助维度基底多样化演化:双坐标视角规避通道冗余 为了打破传统 MLP 通道混合在单一基底下单纯增加宽度导致的参数低效,SKNO 在每个空间位置 \(x_i\) 处,对沿辅助轴离散化的剖面 \(v_l(x_i, \cdot) \in \mathbb{R}^{N_p}\) 引入基底多样化演化模块: $\(v_l(x_i, \cdot) \mapsto \left( \mathcal{F}_p^{-1} \widetilde{A}_l \mathcal{F}_p + B_l \right) v_l(x_i, \cdot)\)$ 其中 \(\mathcal{F}_p\) 为沿辅助维度的离散傅里叶变换矩阵,\(\widetilde{A}_l\) 在频域对偶坐标中进行模式加权与混合,\(B_l\) 则在原始坐标空间中混合信号。命题分析表明,单纯复制同一基底的通道混合(\(B+B\))生成的算子空间仍退化为原有子空间;而引入傅里叶对偶共轭分支(\(A+B\))能够严格扩充有限参数预算下可达的线性算子流形族,且傅里叶变换的正交酉性保证了算子范数不会非受控膨胀。

3. 全局-局部物理传播器协同:多尺度捕捉界面与涡旋结构 截断的傅里叶谱卷积虽然具备全局感受野和准线性计算效率,但在处理不连续介质交界面、高雷诺数涡旋脱落等局部剧烈变化时容易发生吉布斯震荡。SKNO 在物理坐标 \(D_x\) 上采用非对称两阶段设计:前 \(L-1\) 层采用标准截断谱卷积算子: $\(\mathrm{SpectralConvOp}(v_l) = \mathcal{F}_x^{-1} \left( \mathrm{trun\_diag}(\mathcal{F}_x(v_l)) \right)\)$ 在最后一层则挂载局部微分算子,利用 query 点局域邻域 \(D_{loc} \subset D_x\) 内的有限差分导数加权核积分补充高频局部特征: $\(\mathcal{K}_{L-1}[v_{L-1}](x,p) = \frac{\iint_{D_{loc} \times D_p} \kappa_{L-1}(x,y,p,p')(v_{L-1}(y,p') - v_{L-1}(x,p')) \, dy \, dp'}{\iint_{D_{loc} \times D_p} \kappa_{L-1}(x,y,p,p')(y-x,p') \, dy \, dp'}\)$ 实验证实,这种全局打底、末层局部锐化的组合显著降低了输出特征的纠缠熵,精细保留了边界层物理梯度。

损失函数 / 训练策略

模型直接采用算子学习标准的相对 \(L_2\) 范数损失进行端到端优化,不添加任何额外的物理偏置正则项:

\[\mathcal{L}(\theta) = \frac{1}{S} \sum_{s=1}^S \left[ \frac{\lVert \mathcal{G}_\theta[a_s] - u_s \rVert_2}{\lVert u_s \rVert_2} \right] = \frac{1}{S} \sum_{s=1}^S \left( \frac{\sum_j (\mathcal{G}_\theta[a_s](x_j) - u_s(x_j))^2}{\sum_j (u_s(x_j))^2} \right)^{1/2}\]

优化器采用 Adam,学习率初始化为 0.001 并配合余弦退火(Cosine Annealing)调度器。在与 Transformer 类基线(如 Transolver)对比时,不仅对比同等训练轮数(500 epochs),还严格对齐训练墙钟时间(Wall-clock time)以验证算力转化效率。

实验关键数据

主实验

在涵盖 1D、2D 与极端不稳定的 3D 复杂物理系统(单卡 RTX 4090)上,SKNO 全面超越传统 FNO、DeepONet 与最新注意力架构 Transolver:

基准任务 维度与物理机制 FNO (Rel. \(L_2\) / 耗时) Transolver (Rel. \(L_2\) / 耗时) SKNO (本文) 相对 FNO 误差降幅
1D Burgers 1D 非线性对流扩散 (\(\nu=0.1\)) 6.479e-4 / 6.86 min 6.277e-3 / 94.23 min 5.475e-4 / 15.51 min 15.5%
2D Darcy Flow 2D 间断系数二阶椭圆方程 6.201e-3 / 10.21 min 5.853e-3 / 115.48 min 5.555e-3 / 17.59 min 10.4%
2D Stress 2D 复合材料弹性力学应力场 6.221e-2 / 9.43 min 4.812e-2 / 117.39 min 4.724e-2 / 19.36 min 24.1%
2D Strain 2D 复合材料弹性力学应变场 5.753e-2 / 9.43 min 4.657e-2 / 117.39 min 4.635e-2 / 19.36 min 19.4%
2D Gray-Scott 2D 双变量反应扩散系统 2.425e-2 / 9.64 min 3.573e-2 / 234.96 min 1.298e-2 / 20.11 min 46.5%
2D Incomp. NS 2D 低粘性不可压 NS (\(\nu=10^{-5}\)) 1.280e-1 / 773.33 min 1.002e-1 / 1263.33 min 8.717e-2 / 708.93 min 31.9%
3D Comp. NS 3D 跨音速可压缩湍流 Navier-Stokes 2.631e-1 / 10.12 min 2.947e-1 / 361.41 min 2.376e-1 / 23.33 min 9.7%
3D Rayleigh-Taylor 3D 强扰动多相界面失稳流动 (\(A=0.25\)) 5.219e-2 / 5.12 min 5.723e-2 / 135.03 min 4.471e-2 / 10.31 min 14.3%

消融实验与容量效率分析

在 2D Darcy Flow 基准上的模块消融(表 3):

模型配置 Rel. \(L_2\) 误差 说明与影响分析
SKNO (Full Model) 5.555e-3 包含双重残差、双基底辅助演化与全局+局部传播器
w.o. Local Propagator 5.715e-3 移除末层局部微分算子,锐度捕捉能力下降
w.o. \(B_l\) along \(p\) dim 6.827e-3 移除原始辅助坐标混合分支,误差显著上升
w.o. Non-linear Residual 7.621e-3 移除外层非线性残差映射,非线性特征拟合退化
w.o. Linear Residual (\(\bar{\mathcal{W}}_l\)) 8.284e-3 移除线性算子残差,欧拉前向演化被破坏
w.o. \(\widetilde{A}\) along \(p\) dim 8.804e-3 移除频域对偶辅助演化分支,丢失对偶空间模式表达
w.o. Double Residuals 1.363e-2 同时移除线性和非线性残差,误差剧增 145%
w.o. Global Propagators 1.584e-1 移除全局谱卷积传播器,长程上下文丢失导致崩塌

在 2D Gray-Scott 上的容量效率对齐实验(表 10,6 随机种子均值 \(\pm\) 标准差):

模型变体 参数量 (# Params) 计算量 (FLOPs) Rel. \(L_2\) 误差 \(\pm\) Std
FNO (基准宽度) 1,337,450 9.40e8 2.169e-2 \(\pm\) 3.35e-3
FNO (\(2\times\) 宽度) 5,347,538 2.74e9 1.685e-2 \(\pm\) 3.73e-3
FNO (\(4\times\) 宽度) 21,385,634 8.93e9 1.839e-2 \(\pm\) 1.10e-3
2-Stack FNO (同基底并联) 2,685,962 2.33e9 2.532e-2 \(\pm\) 4.29e-3
SKNO (B+B, 同基底双分支消融) 2,690,762 2.52e9 2.447e-2 \(\pm\) 9.65e-4
SKNO (A+B, 频域+空域双基底) 2,141,844 2.11e9 1.256e-2 \(\pm\) 1.14e-3

关键发现

  • 对偶基底分支的关键性:SKNO(A+B) 仅用 2.14M 参数与 2.11 GFLOPs 便达到了 1.256e-2 的相对误差;而同参数量同结构的 SKNO(B+B) 误差高达 2.447e-2,甚至劣于原版 FNO。这从实证上严格证明了收益源自频域-空域双重视角的正交互补性,而非简单的容量扩增。
  • 超分辨率与跨分辨率稳健性:在 1D Burgers 零样本超分测试中,网格从 128 扩展至 8192 时,Transolver 误差从 6.277e-3 激增至 8.993e-3(方差 9.04e-7),而 SKNO 稳定在 5.48e-4 附近(方差仅 2.54e-13);在混合分辨率训练评估中,SKNO 展现出强劲的网格不变性。
  • 时序外推优越性:在 2D Navier-Stokes 自回归长程预测中,面对未见过的未来时间步(+11 至 +20 步),SKNO 平均误差比 FNO、Transolver 分别低 34.5% 与 58.7%,且无自回归误差累积爆发问题。

亮点与洞察

  • 将通道混合提升为辅助连续坐标算子:传统网络把通道视为孤立索引,SKNO 将其视为连续紧致函数域,利用连续核积分统一了特征维度与物理维度的建模范式。
  • 频域与原始域的双基底表征扩充:利用傅里叶变换的酉算子特性在辅助域构建对偶基底,以极轻的计算代价显著扩张了有限预算下可表示的线性算子类。
  • 正交字典稀疏性与纠缠熵抑制:理论推导与 SVD 能量分析证明,SKNO 演化后的特征在最后输出线性层展现出显著高于 FNO 的字典稀疏度与更低的纠缠熵,使得能量高度集中在前几阶奇异模式上。

局限与展望

  • 辅助维度物理可解释性较弱:辅助坐标 \(p\) 目前主要充当特征演化的隐式数学载体,其离散点数 \(N_p\) 与真实物理相空间坐标之间的先验映射关系仍不够直观。
  • 非规则几何网格扩展受限:由于核心传播器基于物理域与辅助域的双重快速傅里叶变换(FFT),模型高度依赖均匀笛卡尔网格,难以直接零样本迁移至非结构化网格复杂外形流动。
  • 未来方向:探索与几何形变网络结合以适配任意复杂外形,并量化研究高维核参数优化轨迹对局部极小值选择的影响机制。

相关工作与启发

  • vs FNO (Fourier Neural Operator): FNO 在物理域利用 FFT 谱卷积,但在通道维仅使用简单的全连接矩阵;SKNO 在辅助维度同样引入对偶谱演化,在保持准线性时间复杂度的同时大幅提升了表示容量与基底多样性。
  • vs Transolver / Physics-Attention: Transolver 依赖自注意力切片聚合捕获复杂几何,计算复杂度高且严重依赖多头与参考网格位置编码;SKNO 依托纯傅里叶算子化演化,在更低 FLOPs 和数倍训练速度下取得了更高精度与跨分辨率稳定性。
  • vs 量子模拟中的薛定谔化方法 (Schrödingerisation): 传统数值算法中的薛定谔化是通过引入虚拟维度将耗散/对流方程转化为高维酉演化;SKNO 借用了引入辅助相空间的几何直觉,将其转化为神经算子隐层特征演化的归纳偏置。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将通道特征混合形式化为辅助连续维度的算子演化,立意深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1D、2D、3D 极端湍流与实际天气数据,包含详尽的消融、分辨率跨尺度与时间外推验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学论证严谨,定理证明与经验设计相辅相成。
  • 价值: ⭐⭐⭐⭐⭐ 为解决神经算子宽度扩展瓶颈提供了极具启发性的低算力替代路线。