神经算子: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变换?

  1. 计算效率:FFT复杂度为 ,远优于直接积分的
  2. 全局感受野:一次FFT即可捕捉全分辨率的空间长程依赖
  3. 参数效率:谱权重 的参数量与空间分辨率无关

谱截断:高频成分通常对应噪声或不稳定模式,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 将神经算子与图神经网络结合:

其中 为节点特征, 为边特征。

关键设计:

  1. 图卷积层:用消息传递(Message Passing)替代 Fourier 变换
  2. 不变性/等变性:分子任务需保持旋转不变性,通过 SphereNet 等架构实现
  3. 多尺度聚合:类似图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, forces

6. 物理信息神经算子

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_physics

7. 与其他方法的比较

7.1 vs 有限元方法(FEM)

维度有限元方法神经算子
计算精度可达机器精度近似解(相对误差 ~
计算效率分钟~小时(逐次求解)毫秒级(算子推理)
泛化能力单次问题跨参数/跨分辨率泛化
可解释性强(数学理论完备)较弱(黑盒性质)
适用场景高精度需求、复杂几何多查询、设计优化、实时预测

7.2 vs 纯数据驱动方法

维度数据驱动(逐点)神经算子
训练数据需要大量输入-输出对可利用物理定律(减少数据依赖)
推理粒度逐点预测全函数输出
物理一致性无保证可通过 PINO 等方法保证
跨分辨率优(FNO 等架构天然支持)

7.3 神经算子家族概览

方法核心思想发表特点
FNOFourier域卷积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

  1. 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

  2. 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

  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