跳转至

In-Context Multi-Operator Learning with DeepOSets

会议: NeurIPS 2026
arXiv: 2512.16074
领域: 物理机理与科学计算
关键词: 多算子学习, 上下文学习, DeepOSets, 神经算子, 偏微分方程求解

一句话总结

本文提出面向上下文多算子学习的 DeepOSets 架构,通过 DeepSets 对少样本 Prompt 提示进行置换不变的非自回归编码,并结合 DeepONet 解码器,首次给出了紧算子集的统一通用逼近定理,在分钟级训练开销下实现了跨多样化微分方程算子的零样本泛化求解。

研究背景与动机

基于深度学习的非线性算子学习已成为科学计算的核心前沿,其典型代表如 DeepONet 和傅里叶神经算子(FNO),通过直接拟合无限维函数空间之间的映射关系,将传统高精度数值求解器数小时的仿真时间压缩至毫秒级推理。然而,传统神经算子通常只针对单一给定的微分方程求解单算子映射(例如仅拟合特定几何或固定系数下的达西流方程),一旦方程类型、控制系数或边界条件发生改变,模型必须重新采集数据并重新训练,无法形成具备泛化能力的科学计算基础模型。

受大语言模型上下文学习(In-Context Learning, ICL)启发,学术界开始探索能够同时学习多个算子的“多算子学习”(Multi-Operator Learning)范式。这一范式面临的核心矛盾在于病态反演:同一参数函数作用于不同方程算子会映射到截然不同的解,模型必须依赖 Prompt 提示中的少量输入-输出样例(In-Context Examples)快速消除目标算子的歧义性。现有主流方案(如 ICON 及其 Transformer 衍生模型)大多沿用 GPT 式的自回归自注意力架构,这引入了对上下文长度呈二次复杂度的计算瓶颈,导致微调与预训练动辄耗费数小时至数天;且注意力机制对无序样例的置换敏感性通常需要额外位置编码调控,缺乏严格的算子逼近理论保障。

本文的切入角度是:算子消歧所需的输入-输出样例本质上是一个无序集合,无需依赖高复杂度的因果注意力,完全可以通过对称池化构建天然具备置换不变性的轻量编码器。核心 idea:将集合学习架构 DeepSets 与神经算子架构 DeepONet 有机结合为 DeepOSets,以线性复杂度并行编码上下文 Prompt 样例并注入分支网络,严格证明了紧算子集上的统一通用逼近性,并在消费级 GPU 上以分钟级训练实现了跨初值与边值微分方程的高精度多算子求解。

方法详解

整体框架

DeepOSets 的核心定位是非自回归、全并行的多算子上下文学习架构。整体系统由两大部分紧密协同:前端的 DeepSets 上下文编码器负责将离散化的 Prompt 样例集合无损映射为紧凑的算子隐空间表征;后端的 DeepONet 解码器则将算子隐表征与查询函数(Query)注入 Branch 网络,并配合接收评估物理坐标的 Trunk 网络完成无网格(Mesh-Free)解函数预测。

系统运行流程分为 Prompt 编码与 Query 解码两个并行阶段。在 Prompt 阶段,用户提供 \(m\) 对用于消除算子歧义的离散函数对 \(\{(u_i(\mathbf{x}), G(u_i)(\mathbf{y}))\}_{i=1}^m\),其中输入参数函数和输出解函数分别采样于固定网格 \(\mathbf{x} \in K_1^k\) 和 \(\mathbf{y} \in K_2^n\);共享的多层感知机(MLP)将各样例映射为特征向量,经由置换不变的对称均值池化得到全局算子嵌入向量 \(\mathbf{h}\)。在 Query 阶段,目标查询函数 \(u_q(\mathbf{x})\) 与全局嵌入 \(\mathbf{h}\) 拼接后输入 Branch 网络,目标评估坐标 \(x_q \in K_2\) 输入 Trunk 网络,两者隐向量内积即可直接计算出目标点处的解预测值 \(\hat{G}(u_q)(x_q)\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["上下文提示样例集<br/>{(u_i(x), G(u_i)(y))}_{i=1}^m"] --> B["逐样例特征映射<br/>共享 MLP: Φ(u_i(x), G(u_i)(y))"]
    B --> C["置换不变对称池化<br/>均值汇聚 h = (1/m) Σ Φ_i"]
    D["查询参数函数 u_q(x)"] --> E["算子条件分支网络<br/>ρ(u_q(x), h)"]
    C --> E
    F["查询评估物理坐标 x_q"] --> G["连续坐标主干网络<br/>Ψ(x_q)"]
    E --> H["内积解码与解函数重构<br/>G^(u_q)(x_q) = Ψ(x_q) · ρ(...)"]
    G --> H

关键设计

1. 严格数学表述与紧算子集通用逼近保证:奠定多算子学习的理论基石 传统算子学习理论仅针对单个连续映射 \(G: V \to Y\),而多算子学习要求同一模型在紧算子族 \(\mathcal{K} \subset \mathcal{O} := \{G: V \to Y \text{ continuous}\}\) 内对任意目标算子进行精准复现。针对这一全新数学命题,本文首先引入了衡量 Prompt 覆盖质量与样本均衡度的 \((\delta, C)\)-离散化概念:对于紧度量空间 \(V\),若离散样例集 \(\mathbf{u} = \{u_1, \dots, u_m\}\) 满足任意两点 \(u, v \in V\) 邻域内样例计数之比受常数 \(C > 1\) 控制(\(|\{u_i \in B(u, \delta/2)\}| < C |\{u_i \in B(v, \delta)\}|\)),则该集合构成了无空洞且密度均匀的覆盖网。在此基础上,本文在 Theorem 5.1 中给出了统一通用逼近定理:对于任意紧算子集 \(\mathcal{K}\) 与任意容许误差 \(\varepsilon > 0\),存在固定离散网格 \(\mathbf{x}, \mathbf{y}\) 与连续网络映射 \(\Phi, \rho, \Psi\),使得在任意 \((\delta, C)\)-离散化提示集下: $\(\sup_{G \in \mathcal{K}} \left| G(u_q)(x_q) - \Psi(x_q) \cdot \rho\left(u_q(\mathbf{x}), \frac{1}{m}\sum_{i=1}^m \Phi(u_i(\mathbf{x}), G(u_i)(\mathbf{y}))\right) \right| < \varepsilon\)$ 该证明借助单位分解定理将紧输出空间 \(C(K_2)\) 投影至有限维网格 \(\mathbf{y}\)(误差小于 \(\varepsilon/2\)),再利用参数空间 \(V\) 上的局部归一化加权平均构造连续函数 \(\Phi\) 与 \(\rho\)。尤为关键的是,\((\delta, C)\)-离散化条件保证了分母权重均值严格远离零点(\(\frac{1}{m}\sum \varphi_j(u_i(\mathbf{x})) \ge \frac{1}{Ct} > 0\)),从而首次在数学上严谨证明了基于 DeepSets 编码的算子学习器具备逼近任意紧算子族的通用性。

2. 非自回归置换不变编码:突破 Transformer 二次复杂度壁垒 针对 ICON 等自回归 Transformer 模型计算复杂度随 Prompt 长度呈平方增长 \(\mathcal{O}(m^2)\)、难以扩展且对提示顺序敏感的痛点,DeepOSets 充分利用了微分方程解算子上下文提示的“集合本质”。Prompt 中的样例对彼此独立,无需时序因果掩码或自注意力交互;因此,模型使用单一共享 MLP \(\Phi: \mathbb{R}^{k+n} \to \mathbb{R}^N\) 独立抽取各样例的高维表征,并通过对称均值操作完成池化: $\(\mathbf{h} = \frac{1}{m} \sum_{i=1}^m \Phi(u_i(\mathbf{x}), G(u_i)(\mathbf{y}))\)$ 该设计将计算开销彻底降至严格的线性复杂度 \(\mathcal{O}(m)\),在训练和推理时可实现极高吞吐的全并行计算。更重要的是,对称池化为模型赋予了结构级的置换不变性归纳偏置(Permutation Invariance),免去了复杂的序列位置编码;且均值池化保证了潜空间表征始终落入紧集内部,在推理时展现出极为出色的无监督长度外推能力,能够在测试时直接接收比训练阶段多数倍的提示样例并获得单调的精度提升。

3. 双路解耦分支-主干解码:实现任意网格无关的物理连续评估 在获取算子全局嵌入 \(\mathbf{h}\) 后,模型需要兼顾参数函数条件约束与目标物理域的无网格求值。DeepOSets 采用扩展版 DeepONet 解码架构,显式解耦了参数空间特征提取与物理时空坐标建模。Branch 网络 \(\rho: \mathbb{R}^{k+N} \to \mathbb{R}^M\) 将目标查询函数的离散采样 \(u_q(\mathbf{x})\) 与算子上下文向量 \(\mathbf{h}\) 直接拼接输入,负责综合当前未知方程的全局算子信息与具体的输入扰动参数;Trunk 网络 \(\Psi: \mathbb{R}^d \to \mathbb{R}^M\) 则仅接收任意连续物理坐标 \(x_q \in K_2\) 作为输入,输出基函数系数。最终解函数通过两者的向量点积瞬时生成: $\(\hat{G}(u_q)(x_q) = \Psi(x_q) \cdot \rho\left(u_q(\mathbf{x}), \mathbf{h}\right)\)$ 这种解耦机制确保了模型推理的完全无网格化(Mesh-Free)。在物理仿真中,Trunk 网络能够自由评估定义域内任意致密、非均匀或边界附近的物理点,而无需重新计算高维 Branch 嵌入,极大地适配了复杂几何场与高分辨率后处理需求。

一个完整示例

以求解一维非线性反应扩散边值问题(Nonlinear Reaction-Diffusion BVP,方程形式为 \(a \frac{d^2}{dx^2}u(x) + k(x) u(x)^3 = c\))为例,其扩散系数 \(a \in [0.025, 0.075]\)、源项常数 \(c \in [-2, 2]\) 以及边界条件 \(u(0), u(1) \in [-1, 1]\) 均为模型预先未知的随机物理参数。

  1. 上下文提示准备:物理域 \([0, 1]\) 均匀离散为 100 个网格点。选取 \(m=4\) 个由高斯过程生成的场分布 \(k_i(x)\) 及求解器算得的真实解 \(u_i(x)\),将每个函数对拼接为 200 维向量,送入共享 MLP \(\Phi\)(含 2 个 500 神经元隐藏层),输出 4 个 500 维特征向量并取均值,得到蕴含当前方程隐藏参数 \(\{a, c, u(0), u(1)\}\) 的算子潜向量 \(\mathbf{h}\)。
  2. 查询输入与解码:给定新的反应系数场分布 \(k_q(x) \in \mathbb{R}^{100}\),将其与 \(\mathbf{h}\) 拼接输入 3 层 Branch 网络(隐藏层维度 200-500-200),生成 2000 维系数向量;同时,Trunk 网络接收查询物理坐标 \(x_q = 0.732\) 并输出对应的 2000 维基函数值。
  3. 点积输出:Branch 与 Trunk 向量做点积,无需更新模型任何权重,瞬时输出精确预测标量 \(\hat{u}_q(0.732)\)。整个前向过程耗时不足 1 毫秒。

损失函数 / 训练策略

模型训练采用端到端监督学习,通过最小化物理域内离散评估点处的预测解与真实解之间的均方误差(Mean Squared Error, MSE): $\(\mathcal{L}_{\text{MSE}} = \frac{1}{B} \sum_{b=1}^B \frac{1}{P} \sum_{j=1}^P \left| \hat{G}^{(b)}(u_q^{(b)})(x_{q,j}^{(b)}) - G^{(b)}(u_q^{(b)})(x_{q,j}^{(b)}) \right|^2\)$ 其中 \(B\) 为每个迭代步采样的算子任务批次大小(设为 100),\(P\) 为每个算子采样的物理评估坐标点数(\(P=100\))。训练过程中,每次迭代随机从不同微分方程(涵盖不同系数与边界条件)中采样算子任务,输入函数 \(u(x)\) 由高斯过程(平方指数核,方差 2,长度尺度 0.5)随机生成。优化器采用 Adam,学习率设为 \(1\times 10^{-3}\),迭代训练 50,000 步。

得益于非自回归与线性复杂度结构,在单张 NVIDIA RTX 4060 Ti GPU(16GB 显存)上,6 任务基准训练仅耗时 12 分钟(训练损失 \(8.4\times 10^{-3}\)),12 任务全量基准训练仅耗时 38 分钟(训练损失 \(4.8\times 10^{-3}\))。

实验关键数据

主实验

论文在涵盖 12 类典型初值(IVP)与边值(BVP)物理微分方程的正反问题上进行了系统评测,包含常微分初值问题(线性、非线性时变系统)、泊松方程、线性反应扩散方程以及非线性三次反应扩散方程。评估对比了仅在 6 个边值问题(6-task)上训练的模型与在全部 12 个任务(12-task)上联合训练的模型在未见测试算子上的均方误差(MSE)。

测试问题任务类别 控制方程与物理类型 6-task DeepOSets (MSE) 12-task DeepOSets (MSE) 误差相对变化
Forward IVP 1 \(\frac{d}{dt}u(t) = a_1 c(t) + a_2\) (正向初值) — (未参与训练) \(2.65\times 10^{-3} \pm 9.30\times 10^{-3}\) 成功支持零样本推断
Inverse IVP 1 \(u(t) \to c(t)\) (反向初值) — (未参与训练) \(1.15\times 10^{-2} \pm 2.71\times 10^{-2}\) 成功支持零样本推断
Forward IVP 2 \(\frac{d}{dt}u(t) = a_1 c(t)u(t) + a_2\) (乘性初值) — (未参与训练) \(1.65\times 10^{-3} \pm 8.39\times 10^{-3}\) 成功支持零样本推断
Inverse IVP 2 \(u(t) \to c(t)\) (反向初值) — (未参与训练) \(1.44\times 10^{-2} \pm 3.86\times 10^{-2}\) 成功支持零样本推断
Forward IVP 3 \(\frac{d}{dt}u(t) = a_1 u + a_2 c + a_3\) (混合初值) — (未参与训练) \(1.19\times 10^{-2} \pm 6.99\times 10^{-2}\) 成功支持零样本推断
Inverse IVP 3 \(u(t) \to c(t)\) (反向初值) — (未参与训练) \(1.40\times 10^{-2} \pm 9.15\times 10^{-2}\) 成功支持零样本推断
Forward Poisson BVP \(\frac{d^2}{dx^2}u(x) = f(x)\) (泊松边值正向) \(3.04\times 10^{-3} \pm 1.91\times 10^{-2}\) \(4.87\times 10^{-4} \pm 1.09\times 10^{-3}\) 误差下降 84.0%
Inverse Poisson BVP \(u(x) \to f(x)\) (泊松边值反向) \(4.60\times 10^{-3} \pm 1.35\times 10^{-2}\) \(3.84\times 10^{-3} \pm 1.20\times 10^{-2}\) 误差下降 16.5%
Forward Linear RD BVP \(a u'' + k(x) u = c\) (线性反应扩散正向) \(2.11\times 10^{-3} \pm 4.25\times 10^{-3}\) \(6.49\times 10^{-4} \pm 9.75\times 10^{-4}\) 误差下降 69.2%
Inverse Linear RD BVP \(u(x) \to k(x)\) (线性反应扩散反向) \(1.45\times 10^{-2} \pm 6.01\times 10^{-2}\) \(7.71\times 10^{-3} \pm 2.75\times 10^{-2}\) 误差下降 46.8%
Forward Nonlinear RD BVP \(a u'' + k(x) u^3 = c\) (非线性反应扩散正向) \(2.77\times 10^{-3} \pm 5.69\times 10^{-3}\) \(1.71\times 10^{-3} \pm 3.32\times 10^{-3}\) 误差下降 38.3%
Inverse Nonlinear RD BVP \(u(x) \to k(x)\) (非线性反应扩散反向) \(1.19\times 10^{-2} \pm 3.51\times 10^{-2}\) \(1.33\times 10^{-2} \pm 5.42\times 10^{-2}\) 误差轻微上升 (+11.8%)

训练开销与任务多样性消融实验

下表总结了任务数量扩展对训练效率、收敛损失及代表性算子泛化精度的影响。实验表明,扩展算子多样性并未导致模型容量崩溃,反而显著促进了多算子上下文消歧能力。

实验配置模型 训练任务数 训练用时 (4060 Ti) 最终训练损失 (MSE) Forward Poisson 测试 MSE Forward Linear RD 测试 MSE
6-task DeepOSets 6 类 (仅 BVP) 12 分钟 \(8.4\times 10^{-3}\) \(3.04\times 10^{-3}\) \(2.11\times 10^{-3}\)
12-task DeepOSets 12 类 (IVP + BVP) 38 分钟 \(4.8\times 10^{-3}\) \(4.87\times 10^{-4}\) \(6.49\times 10^{-4}\)
Transformer基线 (ICON) 同等规模算子 数小时至数十小时 — 相当或略高 相当或略高

关键发现

  • 任务多样性带来正向迁移红利:直觉上同时拟合 12 类算子由于多任务冲突应更具挑战,但实验数据显示 12-task 模型在原有 BVP 任务上的误差几乎全面优于 6-task 模型(例如 Forward Poisson 误差从 \(3.04\times 10^{-3}\) 降至 \(4.87\times 10^{-4}\),降低近一个数量级)。更丰富复杂的方程动态迫使 DeepSets 编码器学到了更本质、更鲁棒的算子鉴别流形。
  • 优异的上下文长度外推能力:模型在训练时仅使用固定数量(\(N=4\))的 Prompt 样例对,但在推理阶段直接零样本测试 \(N=3\) 到 \(N=20\) 的样例输入。实验曲线显示随着提示样例数 \(N\) 增加,正反问题的测试 MSE 呈现平滑单调下降,充分证明模型未发生序列长度过拟合,具备真正的少样本上下文自适应能力。
  • 正问题精度普遍优于反问题:正向边值与初值问题预测 MSE 普遍保持在 \(10^{-4} \sim 10^{-3}\) 量级,而反演问题(由解反求系数场)MSE 多在 \(10^{-2}\) 量级,这与反问题固有的不适定性(Ill-posedness)及更高灵敏度完全吻合。

亮点与洞察

  • 极简架构带来极高训练效率:DeepOSets 放弃了盲目堆叠 Transformer 注意力层,转而采用 DeepSets 结构实现线性复杂度和全并行计算。在单张消费级显卡上仅用 38 分钟即可完成 12 类微分方程多算子大模型的完整训练,为学术界提供了极低算力门槛的算子基础模型研究工具。
  • 首个严谨的上下文多算子通用逼近定理:摆脱了经验主义驱动的架构设计,首次在紧算子空间框架下引入 \((\delta, C)\)-离散化工具,严密证明了有限网格采样与对称池化足以无损逼近紧算子族,填补了该领域的理论空白。
  • 对称池化的天然置换不变性:Prompt 中的样本对是无序经验数据,自注意力机制强加序列位置编码往往导致多余归纳偏见,而对称均值池化不仅保证了置换对称性,还使得上下文样例数量在测试时可自由伸缩外推。

局限与展望

  • 输入网格固定限制:当前的理论推导与实验实现均假设参数函数与解函数采样在固定的空间网格(100 个均匀离散点),尚未支持跨非均匀、非规则几何网格的动态自适应采样输入。
  • 仅验证一维物理系统:实验主要覆盖一维常微分与一维偏微分方程,对于高维多物理场耦合问题(如二维/三维 Navier-Stokes 流体力学方程),固定网格输入维度激增可能加剧 Branch 网络的维度灾难。
  • 未来改进思路:可引入跨网格插值或连续点云编码器(如 PointNet 或几何度量编码),使 DeepSets 的特征提取层能够直接处理非结构化点云网格;同时结合物理守恒律正则化(PINN loss)进一步提升高雷诺数复杂流场的物理保真度。

相关工作与启发

  • vs ICON (Yang et al., 2023a):ICON 采用类 GPT 的自回归 Transformer 处理提示序列,时间与显存复杂度为 \(\mathcal{O}(m^2)\),训练耗费数小时且缺乏通用逼近理论;DeepOSets 采用 DeepSets 结合 DeepONet,复杂度降为 \(\mathcal{O}(m)\),训练仅需数十分钟且具备首个紧算子集逼近保证。
  • vs DeepONet (Lu et al., 2021):经典 DeepONet 仅能拟合单一固定方程的解算子;DeepOSets 通过在 Branch 输入端显式融入 Prompt 对称池化全局表征,将单算子拟合扩展为通用的上下文多算子学习系统。
  • vs PROSE (Sun et al., 2025):PROSE 依赖多模态 Transformer 编码符号方程与数值序列;DeepOSets 证明了即便无需方程显式符号文本,纯粹依靠数据对的对称集合嵌入即可在上下文层面完美消除算子歧义。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出面向多算子学习的置换不变非自回归架构并严格证明了通用逼近定理。
  • 实验充分度: ⭐⭐⭐⭐ 涵盖 12 类初值与边值方程正反问题及上下文外推测试,但受限于一维系统。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学定义清晰,定理证明完整严密,架构叙述逻辑清晰。
  • 价值: ⭐⭐⭐⭐⭐ 极大降低了科学计算基础模型的训练成本,为轻量化多物理场算子学习提供了范式参考。