神经算子:FNO、DeepONet与图神经算子
神经算子(Neural Operators)是一类旨在学习无限维函数空间之间映射的深度学习架构。与传统神经网络学习有限维向量映射不同,神经算子直接学习算子(函数到函数的映射),这使其特别适合求解偏微分方程(Partial Differential Equations, PDEs)、反问题等科学计算任务。1
1. 引言:什么是神经算子?
1.1 从函数逼近到算子逼近
传统深度学习关注的是有限维映射:
例如图像分类网络中,输入是固定维度的像素向量,输出是类别概率分布。
然而,在科学计算中,我们往往需要学习无限维映射——即算子:
典型场景:
- 输入:PDE的初始条件、边界条件或参数(如外力场、介质系数)
- 输出:对应的解函数
1.2 为什么要学习算子?
考虑一个参数化PDE族:
其中 表示参数(如初始条件), 表示解。传统方法(如有限元法)每次只能求解一个特定参数下的方程。而算子学习方法一次训练后,可以泛化到任意参数输入:
| 方法 | 每次推理时间 | 泛化能力 |
|---|---|---|
| 有限元方法 | 分钟~小时 | 单次求解 |
| 神经算子 | 毫秒级 | 算子级泛化 |
这使得神经算子在多尺度模拟、实时预测、设计优化等场景具有显著优势。1
2. 算子学习理论基础
2.1 Banach空间与算子理论
Banach空间是完备的赋范向量空间,是研究无限维函数空间的数学框架。
设 和 为两个 Banach 空间,算子 是从 到 的映射。
有界线性算子:若存在常数 使得 对所有 成立,则 为有界线性算子。线性算子的范数定义为:
\|\mathcal{G}\| = \sup_{a \neq 0} \frac{\|\mathcal{G}(a)\|_{\mathcal{U}}}{\|a\|_{\mathcal{A}}紧算子(Compact Operator):若 将有界集映射为相对紧集(在 中列紧),则称 为紧算子。紧算子具有特殊的谱性质——其谱是可数集,零点外仅有有限或可数个孤立特征值。
2.2 函数空间的逼近
神经算子的核心问题是:神经网络能否逼近任意连续算子?
Stone-Weierstrass定理(实数版本):闭区间上的任意连续函数可用多项式一致逼近。
算子逼近的Universal Approximation Theorem2:设 为紧 Hausdorff 空间(如有界区域), 为连续函数空间。则单隐藏层神经网络在适当的激活函数下,可以在 的任意紧子集上一致逼近任意连续非线性算子。
具体而言,对于非线性算子 ,存在神经网络 使得:
这一定理为神经算子的存在性提供了理论基础。
2.3 Urysohn算子
Urysohn算子是一类重要的非线性积分算子,具有以下形式:
其中 称为核函数(kernel)。当 仅依赖前两个变量时,退化为线性积分算子:
线性积分算子与Fourier分析密切相关—— 的Fourier变换正好对应频域中的乘法运算。这启发了FNO的设计。
3. Fourier神经算子(FNO)
Fourier神经算子(Fourier Neural Operator, FNO)由 Li 等人于 ICLR 2021 提出,是第一个能学习分辨率无关解算子的深度学习架构。1
3.1 架构设计
FNO的核心思想是将积分算子限制为Fourier空间中的卷积,从而实现高效的参数化。
整体架构包含以下组件:
输入 a(x) ∈ ℝ^{H×W×c_in}
│
▼
┌─────────────────┐
│ 升维层 P │ 将 c_in 通道升至 c_hidden
└────────┬────────┘
│
▼ (重复 L 层)
┌─────────────────┐
│ Fourier层 │ │
│ ℱ⁻¹(R·ℱ) │ │ 频域线性变换 + 激活
│ + 局部激活 │ │
└────────┬────────┘
│
▼
┌─────────────────┐
│ 降维层 Q │ 将 c_hidden 降至 c_out
└────────┬────────┘
│
▼
输出 u(x) ∈ ℝ^{H×W×c_out}
3.2 FFT在频域的线性变换
Fourier层的核心运算是:
其中:
- :二维Fourier变换
- :Fourier谱权重,是学习得到的复数矩阵
- :局部线性变换(可选残差连接)
- :激活函数(如GELU)
为什么用Fourier变换?
- 计算效率:FFT复杂度为 ,远优于直接积分的
- 全局感受野:一次FFT即可捕捉全分辨率的空间长程依赖
- 参数效率:谱权重 的参数量与空间分辨率无关
谱截断:高频成分通常对应噪声或不稳定模式,FNO通过截断高频(保留前 个模式)实现降噪和正则化:
3.3 参数效率分析
设输入网格大小为 ,通道数为 。
| 组件 | 参数量 |
|---|---|
| 传统积分算子 | |
| FNO Fourier层 | ,其中 |
| 局部卷积层 | , 为卷积核大小 |
FNO的参数量与输入分辨率无关,这使其能实现分辨率无关(resolution-invariant)的推理——在低分辨率数据上训练,可在高分辨率上测试。
3.4 代码示例(简化版PyTorch)
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List
class SpectralConv2d(nn.Module):
"""Fourier空间的频谱卷积层"""
def __init__(
self,
in_channels: int,
out_channels: int,
modes: int = 16 # 截断的高频模式数
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.modes = modes
# 复数权重: (out_ch, in_ch, modes, modes)
self.weight_real = nn.Parameter(
torch.randn(out_channels, in_channels, modes, modes)
)
self.weight_imag = nn.Parameter(
torch.randn(out_channels, in_channels, modes, modes)
)
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.weight_real)
nn.init.xavier_uniform_(self.weight_imag)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
x: (batch, channels, height, width)
"""
batch, _, h, w = x.shape
# FFT: (batch, ch, h, w) -> (batch, ch, h, w) 复数
x_ft = torch.fft.rfft2(x)
# 取前 modes x modes 的频率
x_ft_trunc = x_ft[:, :, :self.modes, :self.modes]
# 频域乘法: (out_ch, in_ch, modes, modes) @ (batch, in_ch, modes, modes)
weight = torch.complex(self.weight_real, self.weight_imag)
out_ft = torch.einsum('oiwh,biwh->bowh', weight, x_ft_trunc)
# 补零至原尺寸
out_ft_full = F.pad(out_ft, (0, w - self.modes, 0, h - self.modes))
# 逆FFT
return torch.fft.irfft2(out_ft_full, s=(h, w))
class FNO2d(nn.Module):
"""二维Fourier神经算子"""
def __init__(
self,
in_channels: int,
out_channels: int,
hidden_channels: int = 64,
num_layers: int = 4,
modes: int = 16
):
super().__init__()
# 升维: in_ch -> hidden_ch
self.lift = nn.Sequential(
nn.Conv2d(in_channels, hidden_channels, kernel_size=1),
nn.GELU()
)
# Fourier层
self.fourier_layers = nn.ModuleList([
SpectralConv2d(hidden_channels, hidden_channels, modes)
for _ in range(num_layers)
])
self.norms = nn.ModuleList([
nn.LayerNorm([hidden_channels, 0, 0]) # 适配任意分辨率
for _ in range(num_layers)
])
# 降维: hidden_ch -> out_ch
self.project = nn.Conv2d(hidden_channels, out_channels, kernel_size=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.lift(x)
for fourier, norm in zip(self.fourier_layers, self.norms):
# 残差连接
x = x + fourier(x)
# LayerNorm (在channel维度)
x = x.permute(0, 2, 3, 1) # (B, H, W, C)
x = norm(x)
x = x.permute(0, 3, 1, 2) # (B, C, H, W)
x = F.gelu(x)
return self.project(x)
# 使用示例
if __name__ == "__main__":
model = FNO2d(
in_channels=1,
out_channels=1,
hidden_channels=32,
num_layers=4,
modes=12
)
# 输入: 1个样本,1通道,64x64分辨率
x = torch.randn(1, 1, 64, 64)
with torch.no_grad():
y = model(x)
print(f"输入形状: {x.shape}")
print(f"输出形状: {y.shape}")
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")4. DeepONet
DeepONet(Deep Operator Network)由 Lu 等人于 2019 年提出,2021 年发表于 Nature Machine Intelligence。2
4.1 分支网+主干网结构
DeepONet 的核心思想是将算子表示为两个神经网络的积结构:
其中:
- Branch Net(分支网) :编码输入函数 到 维向量
- Trunk Net(主干网) :编码位置坐标 到 维向量
- :感知机数(number of “ensors”),控制网络容量
输入函数 a(s) 位置坐标 y
│ │
▼ ▼
┌──────────────┐ ┌──────────────┐
│ Branch Net │ │ Trunk Net │
│ β: ℱ → ℝ^p │ │ τ: Ω → ℝ^p │
└──────┬───────┘ └──────┬───────┘
│ β₁(a), β₂(a)... │ τ₁(y), τ₂(y)...
│ │
└─────────┬─────────────────┘
│
▼
┌──────────────┐
│ 线性组合 │
│ Σ βₖ(a)·τₖ(y)│
└──────────────┘
│
▼
输出 u(y) ≈ ℒ(a)(y)
4.2 算子逼近理论保证
DeepONet 的理论基础是算子泛化定理。2
定理(DeepONet Universal Approximation):设 为紧集, 为连续算子。则对任意 ,存在整数 和神经网络 ,使得:
逼近误差界:假设分支网和主干网分别用 个神经元近似,则期望误差为:
4.3 训练策略
采样策略:DeepONet 需要同时采样输入函数和输出位置点。
混合采样(Mixed Sampling):
class DeepONet(nn.Module):
def __init__(
self,
branch_layers: List[int], # e.g., [100, 128, 128, 128, 256]
trunk_layers: List[int], # e.g., [2, 128, 128, 128, 256]
p: int = 256 # 感知机数
):
super().__init__()
self.p = p
# 分支网: 编码输入函数(离散采样点)
self.branch = nn.Sequential(
nn.Linear(branch_layers[0], branch_layers[1]),
nn.ReLU(),
# ... 更多层
nn.Linear(branch_layers[-2], p)
)
# 主干网: 编码位置坐标
self.trunk = nn.Sequential(
nn.Linear(trunk_layers[0], trunk_layers[1]),
nn.ReLU(),
# ... 更多层
nn.Linear(trunk_layers[-2], p)
)
def forward(self, a: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
"""
a: (batch, num_sensors) 输入函数在传感器的值
y: (batch, num_points, dim) 位置坐标
"""
# 分支输出: (batch, p)
b = self.branch(a)
# 主干输出: (batch, num_points, p)
t = self.trunk(y)
# 逐点乘积求和
return torch.einsum('bp,bp->b', b, t).unsqueeze(-1)
def train_deeponet(
model: DeepONet,
train_data: list, # [(a_i, u_i)] 输入函数和对应解
epochs: int = 1000,
lr: float = 1e-3
):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
for epoch in range(epochs):
total_loss = 0.0
# 随机采样: 输入函数 + 位置点
for a, u in train_data:
# a: (num_sensors,) u: (num_points,)
a_t = torch.tensor(a, dtype=torch.float32).unsqueeze(0)
# 随机采样 m 个位置
m = 100
indices = np.random.choice(len(u), m, replace=False)
y = torch.tensor(indices.reshape(1, m, 1), dtype=torch.float32)
u_true = torch.tensor(u[indices], dtype=torch.float32)
# 前向传播
u_pred = model(a_t, y) # (1, m, 1)
loss = criterion(u_pred.squeeze(), u_true)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 100 == 0:
print(f"Epoch {epoch}, Loss: {total_loss:.6f}")5. GraphONet:图结构数据上的算子学习
5.1 图神经算子的动机
许多科学问题涉及图结构数据:
- 分子:原子为节点,化学键为边
- 交通网络:路口为节点,道路为边
- 电力系统:母线为节点,输电线路为边
传统 FNO 和 DeepONet 主要处理规则网格或无结构数据,难以直接应用于图。
5.2 GraphONet架构
GraphONet 将神经算子与图神经网络结合:
其中 为节点特征, 为边特征。
关键设计:
- 图卷积层:用消息传递(Message Passing)替代 Fourier 变换
- 不变性/等变性:分子任务需保持旋转不变性,通过 SphereNet 等架构实现
- 多尺度聚合:类似图Transformer的多头注意力机制
5.3 分子动力学应用
分子动力学(Molecular Dynamics, MD)模拟原子间的相互作用,传统方法计算量大。
ATOM(A Pretrained Neural Operator for Multitask Molecular Dynamics)等工作表明,图神经算子可以:
- 预测原子受力(能量关于位置的负梯度)
- 实现端到端的势能面学习
- 在未见过的分子构型上零样本泛化
# 分子势能预测的简化示例
class MolecularOperator(nn.Module):
def __init__(
self,
node_dim: int = 128,
edge_dim: int = 64,
num_layers: int = 6
):
super().__init__()
# 原子类型嵌入
self.atom_embedding = nn.Embedding(100, node_dim)
# 边嵌入(距离、角度等)
self.edge_encoder = nn.Sequential(
nn.Linear(edge_dim, node_dim),
nn.ReLU()
)
# 图消息传递层
self.graph_layers = nn.ModuleList([
GraphConvBlock(node_dim) for _ in range(num_layers)
])
# 输出: 能量 + 原子受力
self.energy_head = nn.Linear(node_dim, 1)
self.force_head = nn.Linear(node_dim, 3)
def forward(
self,
atom_types: torch.Tensor, # (N,) 原子序数
positions: torch.Tensor, # (N, 3) 坐标
edge_index: torch.Tensor, # (2, E) 边连接
edge_attrs: torch.Tensor # (E, edge_dim) 边属性
) -> tuple[torch.Tensor, torch.Tensor]:
# 初始化节点特征
x = self.atom_embedding(atom_types)
# 图卷积
for layer in self.graph_layers:
x = layer(x, edge_index, edge_attrs)
# 原子能量(每原子贡献)
atomic_energy = self.energy_head(x) # (N, 1)
# 总能量 = 求和
total_energy = atomic_energy.sum() # scalar
# 力 = -梯度(能量)
# 需要 autograd,计算每个原子的力
forces = -torch.autograd.grad(
total_energy,
positions,
create_graph=True
)[0] # (N, 3)
return total_energy, forces6. 物理信息神经算子
6.1 PINN思想在算子学习中的应用
物理信息神经算子(Physics-Informed Neural Operator, PINO)将 PINN 的物理约束思想引入算子学习。3
PINO 的目标函数结合两项:
数据损失(Data Loss):
物理损失(Physics Loss):
其中 为 PDE 残差,例如对 Navier-Stokes 方程:
6.2 约束编码方式
硬约束(Hard Constraint):通过设计网络结构满足物理约束
- 边界条件:将边界值直接编码进网络输出
- 不可压缩条件:使用满足 的速度场表示
软约束(Soft Constraint):通过损失函数惩罚违规
def physics_informed_loss(
model: FNO2d,
a: torch.Tensor, # 输入参数/边界条件
u_true: torch.Tensor, # 真实解(可选)
pde_residual: callable, # PDE残差函数
lambda_physics: float = 1.0
) -> torch.Tensor:
"""
计算 PINO 损失
"""
# 预测解
u_pred = model(a)
# 数据损失(如果提供了真值)
loss_data = F.mse_loss(u_pred, u_true) if u_true is not None else 0.0
# 物理损失:计算 PDE 残差
# 需要对输入坐标创建计算图
x = torch.linspace(0, 1, u_pred.shape[-1], requires_grad=True)
y = torch.linspace(0, 1, u_pred.shape[-2], requires_grad=True)
X, Y = torch.meshgrid(x, y, indexing='ij')
# 计算导数(自动微分)
u = u_pred.squeeze()
u_x = torch.autograd.grad(u, X, grad_outputs=torch.ones_like(u))[0]
u_xx = torch.autograd.grad(u_x, X, grad_outputs=torch.ones_like(u))[0]
u_yy = torch.autograd.grad(u_y, Y, grad_outputs=torch.ones_like(u))[0]
# 简化的扩散方程残差: u_t = α∇²u
residual = u_t - alpha * (u_xx + u_yy)
loss_physics = torch.mean(residual ** 2)
return loss_data + lambda_physics * loss_physics7. 与其他方法的比较
7.1 vs 有限元方法(FEM)
| 维度 | 有限元方法 | 神经算子 |
|---|---|---|
| 计算精度 | 可达机器精度 | 近似解(相对误差 ~) |
| 计算效率 | 分钟~小时(逐次求解) | 毫秒级(算子推理) |
| 泛化能力 | 单次问题 | 跨参数/跨分辨率泛化 |
| 可解释性 | 强(数学理论完备) | 较弱(黑盒性质) |
| 适用场景 | 高精度需求、复杂几何 | 多查询、设计优化、实时预测 |
7.2 vs 纯数据驱动方法
| 维度 | 数据驱动(逐点) | 神经算子 |
|---|---|---|
| 训练数据 | 需要大量输入-输出对 | 可利用物理定律(减少数据依赖) |
| 推理粒度 | 逐点预测 | 全函数输出 |
| 物理一致性 | 无保证 | 可通过 PINO 等方法保证 |
| 跨分辨率 | 差 | 优(FNO 等架构天然支持) |
7.3 神经算子家族概览
| 方法 | 核心思想 | 发表 | 特点 |
|---|---|---|---|
| FNO | Fourier域卷积 | ICLR 2021 | 适合周期边界问题,高效全局建模 |
| DeepONet | 分支-主干网积结构 | Nat Mach Intell 2021 | 通用架构,理论保证 |
| GraphONet | 图神经网络+算子 | NeurIPS 2021+ | 分子/图结构数据 |
| PINO | 数据+物理双监督 | ICLR 2022 | 减少数据需求,保证物理 |
| KnoIN | 核积分网络 | ICML 2023 | 可解释性强 |
8. 应用场景
8.1 PDE求解
神经算子最直接的应用是学习PDE的解算子:
- Navier-Stokes方程:预测流场演化
- 弹性力学方程:结构力学分析
- 热传导方程:传热问题
- Maxwell方程:电磁场模拟
优势:一次训练,多次推理;支持参数化PDE族(如不同雷诺数的流动)。
8.2 反问题
反问题旨在从观测数据推断系统参数:
其中 为待推断的参数(如介电系数、源项), 为观测数据。
神经算子可以作为代理模型,加速反问题的迭代求解:
- 每次迭代只需调用前向算子(毫秒级)
- 可与贝叶斯推断、MCMC等方法结合
8.3 分子模拟
分子模拟中的关键算子:
| 算子 | 输入 | 输出 |
|---|---|---|
| 势能面 | 原子坐标 | 势能 + 力 |
| 力场 | 分子构型 | 原子受力 |
| 动力学传播 | 时刻构型 | 时刻构型 |
图神经算子可以学习这些算子,实现:
- 分子动力学加速(MDNet等)
- 药物-靶标相互作用预测
- 新材料发现
9. 参考资料
本文档涵盖神经算子的核心概念、主流架构及最新进展。如有补充或修正,欢迎通过 Wiki 编辑提交。
Footnotes
-
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2021). Fourier Neural Operator for Parametric Partial Differential Equations. International Conference on Learning Representations (ICLR). arXiv:2010.08895 ↩ ↩2 ↩3
-
Lu, L., Jin, P., Pang, G., Zhang, Z., & Karniadakis, G. E. (2021). Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3), 218-227. doi:10.1038/s42256-021-00302-5 ↩ ↩2 ↩3
-
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Stuart, A., Bhattacharya, K., & Anandkumar, A. (2023). Physics-Informed Neural Operator for Learning Partial Differential Equations. ACM Transactions on Mathematical Software. arXiv:2111.03794 ↩