AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training

会议: NeurIPS 2026
arXiv: 2605.15793
论文: https://arxiv.org/abs/2605.15793
领域: 科学计算
关键词: 神经算子预训练, 偏微分方程基础模型, 自适应算子变换, 多流残差网络, Sinkhorn-Knopp投影

一句话总结

针对异质偏微分方程（PDE）解算子结构差异大导致的预训练拟合瓶颈，AOT-POT受经典数值变换启发引入输入自适应的多流算子变换块，仅增加3%参数即可将解算子投影为更易联合建模的统一空间，在12项基准上相对L2误差平均降低40.9%。

研究背景与动机

在大规模科学机器学习领域，基于异质PDE数据集预训练神经算子以构建通用的代理模型已成为核心研究方向。传统神经算子（如FNO、DeepONet、DPOT）通过自回归方式递推时间演化算子，相较经典数值求解器可提速数个数量级。然而，不同物理系统（如不可压缩流体、可压缩Navier-Stokes、浅水波方程与反应扩散系统）底层所对应的动力学微分算子结构截然不同，时空演化特性与频谱衰减规律差异极大。强行用单一固定的骨干网络直接联合拟合如此庞大且结构高度冲突的解算子族，必然会导致严重的逼近误差与负迁移现象。

面对这种异质解算子的拟合难题，现有工作（如DPOT扩维或MoE-POT稀疏专家路由）主要遵循大语言模型的扩容范式，通过增加网络深度、加宽通道或堆叠稀疏专家子网络来硬性提升模型容量，却未改变目标解算子本身的复杂形态。然而数值计算经典理论表明，通过适定的算子变换（如傅里叶基下的Laplacian对角化或线性代数中的预条件子加速）可将复杂的原算子转换为结构更良态、计算更简化的等价形式。更关键的是，不同PDE所需的最优变换在跨方程时往往互不兼容甚至相互冲突。

因此，单纯扩大容量无法解决算子本身的内在异质性，而静态预定义的变换又无法适配多样化方程。核心 idea：将经典数值算子变换思想内化为神经网络架构设计，构建输入自适应的多流演化与双随机混合模块，动态将异质PDE解算子投影至平滑统一的潜在表示空间。

方法详解

整体框架

AOT-POT将自适应算子变换机制无缝融入算子Transformer主干中。系统输入为历史 \(T\) 步离散观测场 \(\bm{u}^{<T} \in \mathbb{R}^{H \times W \times T \times C}\)，首先通过逐帧补丁化卷积与坐标位置编码映射为空间patch特征，并结合复指数加权MLP聚合时间维度动态。随后，特征被升维扩展至 \(n\) 条并行的特征流（隐空间完备基分量），交由 \(N\) 层堆叠的AOT Block进行逐层算子变换与多头傅里叶注意力交互。在每个AOT块内部，模型依次执行基于全局池化特征的输入自适应流聚合、傅里叶核积分演化、自适应输出重分布，以及经由Sinkhorn-Knopp投影保证双随机特性的流间混合，最终通过门控读出层折叠流维度并输出下一时刻预测场 \(\bm{u}^{t}\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["历史场序列输入<br/>u &lt; T, 补丁化与时间聚合"] --> B["并行流基底扩展<br/>升维至 n 条并行隐特征流"]
    B --> C["AOT Block: 自适应聚合<br/>输入依赖凸组合 a_l 提取演化基"]
    C --> D["Fourier Mixer: 频域核积分<br/>多头复数频域变换处理空间交互"]
    D --> E["AOT Block: 自适应重分布<br/>输入依赖权重 d_l 散射回各流通道"]
    E --> F["AOT Block: 双随机矩阵混合<br/>Sinkhorn 投影算子核 T_l 跨流正交混合"]
    F --> G["门控读出与场预测<br/>自适应折叠并行流输出 u_t"]

关键设计

1. 多流自适应聚合与重分布：实现输入条件化的基底变换 标准Transformer的残差连接遵循固定的加法路径 \(\bm{x}_{l+1} = \bm{x}_{l} + \mathcal{F}(\bm{x}_{l})\)，单一特征流无法针对不同类型的偏微分方程动态调整基底。AOT-POT在层内维持 \(n\) 个并行流 \(\bm{x}_{l} \in \mathbb{R}^{n \times C}\) 作为潜在解空间的过完备基底，并在子层 \(\mathcal{F}\) 前后引入输入依赖的聚合权重 \(\bm{a}_{l} \in \mathbb{R}^{1 \times n}\) 与重分布权重 \(\bm{d}_{l} \in \mathbb{R}^{1 \times n}\)：

\[\bm{x}_{l+1} = \bm{T}_{l}\,\bm{x}_{l} + \bm{d}_{l}^{\top}\,\mathcal{F}(\bm{a}_{l}\,\bm{x}_{l}, \mathcal{W}_{l})\]

其中权重向量通过对隐状态执行空间全局平均池化（GAP）及RMSNorm后，经过带门控标量的线性变换与激活函数动态生成。聚合权重满足概率单纯形约束 \(\bm{a}_{l} \geq 0\) 且 \(\bm{1}_{n}^{\top}\bm{a}_{l} = 1\)，构成输入依赖的基底凸组合提取；重分布权重经由 \(2\sigma(\tilde{\bm{d}}_{l})\) 调节各基分量接收更新的强度，使单一骨干能根据输入的流体系数或方程类型自主决定当前子层应处理的几何基底。

2. Sinkhorn投影双随机算子核：保证深度变换的数值稳定性 若直接允许流间转换核 \(\bm{T}_{l} \in \mathbb{R}^{n \times n}\) 任意参数化，随着网络层数加深，复合残差变换 \(\prod \bm{T}_{L-i}\) 会导致信号幅值指数级放大或快速衰减，引发深层算子网络常见的梯度崩溃。为此，AOT-POT将变换核强约束于Birkhoff多面体（双随机矩阵集 \(\mathcal{M}\)）：

\[\mathcal{M} = \left\{\bm{T}_{l} \in \mathbb{R}^{n \times n} \;\middle|\; \bm{T}_{l}\bm{1}_{n}=\bm{1}_{n}, \; \bm{1}_{n}^{\top}\bm{T}_{l}=\bm{1}_{n}^{\top}, \; \bm{T}_{l} \geq 0\right\}\]

模型使用动态生成的矩阵 \(\tilde{\bm{T}}_{l}\) 并通过Sinkhorn-Knopp算法进行行、列交替归一化投影得到 \(\bm{T}_{l}\)。该双随机性质赋予了算子变换两大核心优势：一是严格保持保积性质（谱范数有界于1），彻底杜绝跨层梯度爆炸；二是双随机矩阵相乘在代数上关于合成封闭，确保了多层连续变换在极深架构下的长期演化稳定。

3. 多头傅里叶频域混合与门控折叠：捕捉长程连续算子特征 在被自适应聚合的单一特征切片进入子层后，模型采用多头傅里叶注意力层实现空间全局卷积核积分计算。利用平移不变性将连续积分核算子转化为频域可学习权重矩阵点乘，各注意力头利用复数域两层MLP直接对频谱各模态赋予演化变换。在完成 \(N\) 层AOT处理后，末端的 \(n\) 条并行流并不采取破坏物理空间平滑性的粗暴拼接，而是通过轻量门控向量 \(\bm{g} = \text{softmax}(\bm{w})\)（初值置零以实现均等初始化）对各基底进行自适应凸组合读出，既保留了各变换基的解耦物理特征，又以极低开销无缝接入去噪预测目标。

损失函数 / 训练策略

训练采用自回归去噪目标函数，为抵御长时间积分过程中的复合误差累积，在前序历史帧上叠加高斯扰动 \(\bm{\varepsilon} \sim \mathcal{N}(0, \epsilon \|\bm{u}^{<t}\| \bm{I})\)。模型在混合数据集 \(\mathcal{D}\) 上通过均方相对误差联合训练：

\[\mathcal{L} = \sum_{1 \leq t \leq T} \left\| \mathcal{G}_{w}(\bm{u}^{<t} + \bm{\varepsilon}) - \bm{u}^{t} \right\|_{2}^{2}\]

优化器采用AdamW，学习率采用余弦退火策略，在8张NVIDIA H200 GPU上进行1000个epoch的大规模多PDE联合预训练。

实验关键数据

主实验

在涵盖FNO、PDEBench、PDEArena与CFDBench共12个基准数据集上，对比了专有小模型、主流零样本预训练模型（DPOT、MoE-POT、MPP）以及下游全微调设置。评测指标为相对 \(L_{2}\) 误差（L2RE，越低越好）：

模型体系	参数量	FNO (\(\nu=10^{-4}\))	CNS (1, 0.1)	CNS (0.1, 0.01)	PDEBench DR	PDEBench SWE	PDEArena NS	CFDBench	平均降幅
FNO (单任务)	0.5M	0.0834	0.0980	0.170	0.1200	0.00440	0.0912	0.00761	-
DPOT-T	7M	0.0606	0.0173	0.0220	0.0321	0.00560	0.1250	0.00952	基线
MoE-POT-T	17M	0.0588	0.0160	0.0215	0.0310	0.00527	0.1340	0.01030	~3.2%
AOT-POT-T (本文)	7M	0.0438	0.0110	0.0130	0.0241	0.00433	0.0985	0.00776	-28.7%
DPOT-S	30M	0.0442	0.0153	0.0188	0.0379	0.00657	0.0991	0.00696	基线
MoE-POT-S	90M	0.0417	0.0143	0.0181	0.0354	0.00629	0.1060	0.00748	~5.8%
AOT-POT-S (本文)	31M	0.0333	0.00767	0.00965	0.0146	0.00266	0.0784	0.00493	-40.9%
DPOT-M	122M	0.0285	0.0116	0.0129	0.0292	0.00290	0.0812	0.00752	基线
AOT-POT-M (本文)	124M	0.0257	0.00906	0.00954	0.0186	0.00200	0.0775	0.00543	-24.5%

消融实验

在AOT-POT-S配置下针对AOT块中的三大核心映射机制开展消融对比（统一在12个预训练数据集上从头预训练）：

实验变体配置	聚合映射 \(\bm{a}\)	重分布映射 \(\bm{d}\)	混合变换核 \(\bm{T}\)	12数据集平均L2RE	性能退化幅度	说明
AOT-POT-S (Full)	动态单纯形	动态Sigmoid	Sinkhorn双随机	0.04005	0.0%	完整模型性能最优
w/o \(\bm{a}\)	均等池化	动态Sigmoid	Sinkhorn双随机	0.04223	+5.4%	削弱了输入依赖的基底选择能力
w/o \(\bm{d}\)	动态单纯形	均等散射	Sinkhorn双随机	0.04172	+4.2%	限制了不同流接收梯度的灵活性
w/o \(\bm{T}\)	动态单纯形	动态Sigmoid	单位阵恒等	0.04554	+13.7%	丢失跨流算子空间旋转投影能力

关键发现

• 参数极低开销与跨级压制：AOT-POT仅在DPOT基础上增加了3%的模型参数量，但AOT-POT-S（31M）在12个数据集中有11个全面超越了参数量高出4倍的DPOT-M（122M），表明优化算子表征空间比堆砌模型深度与宽度更高效。
• 消融敏感度与变换核核心地位：去除变换核 \(\bm{T}\) 使得平均相对误差骤降13.7%，证明基于Birkhoff多面体的双随机基底旋转变换是重构异质算子学习难度最关键的数学枢纽。
• 长时程自回归稳定性：在500步长轨迹递推评估中，基线DPOT-FT-S误差增长了24.1倍，而AOT-POT-FT-S仅增长9.8倍（误差增速降低约2.5倍），双随机核的严格范数保持机制有效阻断了递归误差的发散。
• 算子可解释性：将各层学得的变换核矩阵 \(\bm{T}_{l}\) 展平后进行t-SNE可视化与最近邻分类，模型对预训练物理方程类型的无监督聚类准确率高达99.7%，并在未见过的波方程系统上达到100%判别率，证实网络自主习得了对应于具体物理算子的正交变换特征。

亮点与洞察

• 经典数值变换的神经内化：将传统计算数学中利用基变换对角化微分算子的先验，转化为可微的双随机特征流旋转操作，摆脱了暴力盲目增大模型容量的单一思维定式。
• 构造性双随机正则约束：借助Sinkhorn-Knopp将变换矩阵限定在Birkhoff流形，在赋予过完备基变换极大自由度的同时，利用范数有界和代数封闭特性彻底锁死深层演化中的数值发散隐患。
• 强大的零样本算子迁移性：自适应算子变换赋予了主干网络极高的泛化先验，微调时在未见过的全新PDE类型（如Wave-Layer）上仅需16个样本即可击败基线使用64个样本的微调表现。

局限与展望

• 逼近论理论严格证明缺失：论文通过详尽的动机实验和经验测试验证了自适应算子变换的有效性，但尚未从无穷维算子逼近论的角度严格证明Birkhoff投影流变换对误差上界的压缩速率。
• 显式动力学语义对应模糊：尽管分类探测器证明 \(\bm{T}_{l}\) 编码了高辨识度的PDE类型特征，但变换核权重的特定模式与真实物理参数（如粘性系数、耗散项阶数）之间的精确解析对应关系尚待破译。
• 复杂不规则网格扩展：当前实现深度依赖于基于规则网格的Patchify和傅里叶积分层，未来拓展至非欧几里得几何与非结构化非均匀网格时，需要探索流形上的等变双随机变换算子。

相关工作与启发

• vs DPOT / POSEIDON (通用预训练神经算子)：DPOT与POSEIDON试图以统一固定的通道结构拟合多源PDE数据，在遇到异质性强烈的算子时容易互相干扰；AOT-POT通过动态多流变换解耦了不同方程的几何投影，在同等甚至更小参数规模下大幅刷新预测精度。
• vs MoE-POT (稀疏专家架构)：MoE-POT通过硬性路由将不同数据分配给独立的专家网络，参数量随专家数激增且未改变算子逼近难度；AOT-POT采用软性多流重分布与双随机旋转，以3%极低参数增量实现了超越稀疏专家的表示容量。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 巧妙将经典数值计算算子变换与双随机流形结合，构思精巧。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖12个预训练数据集、下游跨维度迁移、500步长程递推及详细消融与可解释性分析。
• 写作质量: ⭐⭐⭐⭐⭐ 理论动机连贯清晰，数学推导严谨，实验叙述详实。
• 价值: ⭐⭐⭐⭐⭐ 为科学基础模型（AI for Science / PDE Foundation Models）开辟了模型扩容之外的算子自适应变换全新维度。

相关论文

• [CVPR 2026] NESTOR: A Nested MOE-based Neural Operator for Large-Scale PDE Pre-Training
• [NeurIPS 2026] Conservative Neural Posterior Estimation via Distributionally Robust Training
• [NeurIPS 2026] Curriculum Multiple Shooting for Robust Training of Neural and Universal Differential Equations
• [NeurIPS 2026] MC2: Monte Carlo Correction for Fast Elliptic PDE Solving
• [NeurIPS 2026] Cross-Fitting for Neural Posterior Estimation