把指数级计算拆成四条一维路径：可分离 PI-DeepONet 如何破解高维 PDE
PINN 最难的地方，有时不是方程，而是梯度
我一直认为，物理信息机器学习最有吸引力的地方，是它让神经网络不再只追逐数据拟合，而是必须遵守偏微分方程、初始条件和边界条件。但当问题从一维扩展到三维空间加时间，PINN 类方法很快会撞上一个现实：训练损失里的导数计算，比网络本身更昂贵。
传统 PINN 针对一个具体的初始条件和边界条件，学习一个解函数；Physics-Informed DeepONet（PI-DeepONet）则更进一步，希望学习"输入函数到输出函数"的解算子。给它一组初始条件、载荷或源项，它应该直接输出整段时空解。这样训练完成后，就可以快速响应不同输入，适合数字孪生、实时控制和参数扫描。
问题在于，算子学习要在许多空间---时间位置上同时计算 PDE 残差。如果每个坐标轴各取 `n` 个点，一个 `d` 维时空域就有 `n^d` 个网格点。三维空间加时间就是四维；当 `n=50` 时，组合点数已经达到 625 万。更麻烦的是，反向自动微分还要构造规模随输出维度膨胀的雅可比矩阵。
我的理解是：**PI-DeepONet 并不是"网络太大"，而是把一个可以分解的多维问题，按笛卡尔积的方式完整展开了。**只要把这种结构重新利用起来，很多看似指数级的计算就有机会被拆成若干条一维路径。
一、先把 DeepONet 说清楚：它学习的不是一张图，而是一条映射
普通监督学习通常输入一个向量，输出一个向量；DeepONet 面对的是函数。它由两个部分组成：
-
branch 网络读取输入函数在若干传感器点上的值，例如初始温度、载荷时间序列或初始速度场；
-
trunk 网络读取查询坐标，例如 `(x,y,z,t)`，返回与该位置相关的基函数特征。
最后，两者做内积，得到这个输入函数在这个坐标处的解值。可以把它写成：
```text
输入函数 u(·) + 查询坐标 ζ
↓
DeepONet(u, ζ) → 解 s(ζ)
```
当查询坐标遍布整个时空域，DeepONet 就能一次输出一张完整的解场。它的优势是训练后推理快、可以覆盖一族初始条件和边界条件；但它仍然依赖大量标注解数据。对于复杂 PDE，生成这些输入---输出对本身就可能非常昂贵。
PI-DeepONet 的思路是减少对标注数据的依赖：训练时不必知道每个点的真实解，只要计算 PDE 残差、初始条件损失和边界条件损失即可。总损失可以理解为：
```text
总损失 = PDE 残差 + 初始条件误差 + 边界条件误差
```
网络如果真的学到了正确的解，代入控制方程后，残差就应该接近零。
二、瓶颈在哪里：从 n^d 个点到巨大的雅可比矩阵
以四维问题 `(x,y,z,t)` 为例，每个方向取 `n` 个点，普通 PI-DeepONet 的 trunk 网络需要在 `n^4` 个坐标上前向传播。假设有 `N` 个输入函数，网络还要为每个输入重复处理这些坐标。
在计算 PDE 导数时，常用反向模式自动微分。它适合"输出少、输入多"的场景，但这里的输出解场和坐标输入规模都很大，导致需要处理一个接近 `N·n^4 × N·n^4` 的雅可比结构。即使不真的把矩阵全部显式存下来，中间激活和反向图也会迅速吃满 GPU 内存。
这就是所谓的维度灾难：维度从 2 增加到 4，计算量不是增加两倍，而是从 `n^2` 变成 `n^4`。一旦还要加入材料参数、几何参数或随机载荷，训练成本会进一步上升。
我觉得这里最重要的观察是：这些 `n^4` 个点并不是完全独立的，它们本来就是四条一维坐标轴的笛卡尔积。既然解场可以用变量分离或低秩张量近似表示，就没必要让每条坐标组合都经过一次完整的 trunk 网络。
三、核心改造：让每个坐标轴拥有自己的小网络
可分离 PI-DeepONet（Sep-PI-DeepONet）的第一步，是把一个接收完整 `(x,y,z,t)` 的 trunk MLP，改成四个一维子网络：一个只看 `x`，一个只看 `y`，一个只看 `z`，一个只看 `t`。
每个子网络把一维坐标映射成低维特征。之后，通过外积把四条特征组合成四维解场，再沿低秩维度求和。用直觉化的表达：
```text
T(x,y,z,t) ≈ Σᵣ Txᵣ(x) · Tyᵣ(y) · Tzᵣ(z) · Ttᵣ(t)
```
这并不是简单地假设解完全可分离，而是用多个 rank 分量叠加，近似一个更复杂的张量。rank 越大，表达能力越强；rank 较小时，计算和内存更省。
原来需要 `n^4` 次完整 trunk 前向传播，现在只需要 `4n` 次一维前向传播，再用张量运算拼成完整场。输入函数的 branch 输出也不再被重复复制，而是和 trunk 的所有轴特征通过 einsum 一次性组合。
这一步改变了计算的增长方式：随着每个方向采样点数 `n` 增加，代价近似线性增长，而不是指数增长。更重要的是，它保留了 DeepONet 的函数到函数映射能力，不需要把问题退化成一个固定网格上的普通回归器。
四、第二个关键：为什么要换成前向自动微分
仅仅拆分 trunk 网络还不够。PI-DeepONet 的核心损失是 PDE 残差，而残差需要一阶、二阶甚至混合偏导数。可分离结构下，网络输入的四条一维坐标和最终输出场不再具有相同形状，原先高效的反向自动微分策略无法直接套用。
这里改用前向模式自动微分。可以把两种模式想象成两种问法：
-
反向模式：从一个输出反推它依赖了哪些输入；
-
前向模式：给定一个输入方向，追踪它如何影响所有输出。
当输出维度远大于输入维度时，前向模式的内存占用更友好。对于可分离结构，坐标轴上只有 `4n` 个一维节点，雅可比计算也从巨大的全场规模降到与 `n` 同阶的紧凑形式。
最终，训练流程变成：四个一维 trunk 网络生成可分离特征，外积构成解场，前向 AD 计算导数，PDE 残差和边界/初始条件共同更新 branch 与 trunk 参数。
五、算例一：Burgers 方程验证非线性能力
第一个测试是带黏性的 Burgers 方程。它包含对流非线性项 `u·∂u/∂x` 和扩散项 `ν·∂²u/∂x²`，常被用作检验物理信息网络是否能处理非线性输运的基准。
这里的目标不是求一条固定初始条件的解，而是学习从随机初始速度场到完整时空速度场的算子。训练只使用初始条件、周期边界条件和 PDE 残差，不依赖数值解标签。
在相近参数量下，普通 PI-DeepONet 的平均相对 L2 误差约为 `5.2×10^-2`，每次迭代约 136.6 ms；可分离版本误差约为 `6.2×10^-2`，每次迭代约 3.64 ms。精度基本同一量级，但单步训练计算下降了约两个数量级。
固定训练 83.69 秒比较时，普通模型只完成约 600 次迭代，相对 L2 误差约 `3.82×10^-1`；可分离模型完成约 21,500 次迭代，误差降到约 `8.98×10^-2`。这说明同样的墙钟时间里，效率提升不仅意味着"每一步更快"，还意味着优化器能走更多步。
可分离结构的 rank 和隐藏维度之间存在明显权衡。较大的隐藏维度配合较小的 rank 往往更有效；如果两个都盲目增大，模型可能出现接近平凡的解或训练不稳定。因此，低秩并不是越低越好，而是要在表达能力和张量组合成本之间寻找合适位置。
六、算例二：Biot 固结把两个物理场同时拉进来
第二个算例来自 Biot 固结理论，描述流体饱和多孔介质在载荷作用下的耦合响应。网络不仅要预测固体位移，还要预测孔隙压力。
输入是一条施加在排水边界上的载荷时间函数，输出是整根一维柱体随空间和时间变化的位移场与压力场。这个算例的意义在于，它不再是单一标量场，而是两个相互耦合的物理量。
普通 PI-DeepONet 的平均相对 L2 误差约为 `7.66×10^-2`，训练耗时约 25,415 秒；参数量相近的可分离模型，误差约 `7.83×10^-2`，训练只需约 552 秒，时间减少约 97.83%。更大一些的可分离模型，误差可降到约 `6.17×10^-2`，训练耗时仍只有约 997 秒。
这个结果让我看到，可分离结构并不是只适用于"一个方程、一个输出"。只要多个物理场共享相同坐标结构，并且可以通过最后的输出通道拆分，位移、压力、温度或浓度都可以一起预测。
七、算例三：四维参数化热方程是最有说服力的一关
最关键的测试是二维板上的参数化热方程。空间有 `x,y` 两个方向，时间是 `t`，热扩散率 `α` 又作为参数维度加入，整个问题相当于四维输入域。
网络输入初始温度和热扩散率，输出整块板随时间演化的温度场。训练过程中不使用数值解标签，而是把边界温度、初始温度和热方程残差写入损失。
可分离模型使用四个 trunk 子网络、约 57.7 万个参数，训练约 2.5 小时，最终相对 L2 误差约 `8.56×10^-2`。相同离散密度和训练轮数下，普通 PI-DeepONet 的单次迭代约需 10,416.7 ms，估计完整训练需要 289.35 小时。
这不是简单的"快一点"，而是从可运行变成不可运行。普通版本的计算成本已经让四维测试无法实际完成，而可分离版本可以在数小时内收敛，并对未见过的初始温度和扩散率组合给出合理温度场。
我认为，这个算例真正说明了可分离设计的价值：当维度增加时，传统 PI-DeepONet 不是精度慢慢变差，而是训练成本突然爆炸；可分离结构则把爆炸趋势压平，让高维参数化 PDE 仍然可以被训练。
八、额外测试：随机源项与卷积分支
为了展示结构的可扩展性，还可以把 branch 网络替换为 CNN，用于二维 Poisson 方程的随机源场。源项是一个空间相关的高斯随机场，网络输出满足齐次 Dirichlet 边界的解场。
在这个设置下，卷积分支负责读取二维源项，分离 trunk 负责处理空间坐标和扩散系数。模型约有 618.8 万个参数，平均相对 L2 误差约 `4.73×10^-2`，说明可分离思想并不局限于最基础的全连接 branch/trunk 组合。
这让我想到一个更大的设计空间：branch 可以是 CNN、RNN、Transformer 或图网络，trunk 仍然保留按坐标分解的结构。只要输入函数和查询坐标可以分别编码，算子输出就有机会通过低秩组合构建。
九、它真正解决的不是"高维"，而是高维里的重复计算
我不认为 Sep-PI-DeepONet magically 消除了维度本身。它做的是把高维网格中大量重复的结构提取出来：同一个 `x` 坐标会和许多 `y,z,t` 组合，同一个 `t` 坐标也会在不同空间位置反复出现。普通网络把这些组合逐个处理，可分离网络把每条轴的特征先算一次，再用外积复用。
这种思路与变量分离、低秩张量分解、Proper Generalized Decomposition 有相似精神，但这里的因子不是手工求出来的，而是由神经网络学习。它可以在保持非线性表达能力的同时，把坐标结构暴露给自动微分和硬件张量运算。
从复杂度角度看，三个变化叠加在一起：
-
分离 trunk 把 `n^d` 次前向传播降到 `d·n` 级别；
-
外积和求和代替重复拼接，减少输入函数和坐标的复制；
-
前向 AD 避免构造与全场输出同规模的反向雅可比。
这也是为什么论文中的时间改进可以达到两个数量级，而不是一个小的工程优化。
十、边界在哪里：可分离不是万能钥匙
我会特别注意这套方法的适用条件。第一，问题最好存在较规则的坐标结构，或者至少可以找到合理的因子化坐标。复杂不规则几何、移动边界、裂纹拓扑变化可能很难直接放进笛卡尔积网格。
第二，参数组合必须物理可行。把几何、材料和边界参数笛卡尔积采样，可能产生大量不相容组合；训练时需要筛选有效样本，或者在损失计算中跳过无效区域，否则内存和优化都会受到影响。
第三，外积虽然降低了前向传播成本，却仍可能产生较大的中间张量。加入更多参数维度、更多输出场或扩展到 branch 端的分离结构后，显存仍然会成为瓶颈。
第四，低秩近似有表达上限。如果解场包含尖锐界面、强局部奇异性或多尺度涡结构，较小 rank 可能无法表达细节；rank 增大又会带来新的计算和超参数问题。
第五，训练只使用物理残差并不意味着物理一定正确。方程、边界条件、参数范围和损失权重如果设定不当，网络仍可能收敛到不符合真实系统的近似解。
十一、结论：让 PINN 先学会"分解"，再谈高维扩展
我从这套方法中得到的最大启发是：解决 PINN 的维度灾难，不一定要继续堆更大的网络或更强的 GPU，先问一句"这个高维对象能不能按坐标和物理结构分解"往往更重要。
Sep-PI-DeepONet 把一个四维 trunk MLP 改成多个一维子网络，用低秩外积重新拼出解场，再用前向自动微分计算 PDE 导数。这样，模型仍然是 physics-informed 的，仍然能学习函数到函数的解算子，却把大量重复的坐标计算和雅可比开销压缩掉。
在 Burgers 方程上，它用接近的误差换来约两个数量级的单步加速；在 Biot 固结上，训练时间从数小时级降到十分钟级；在四维参数化热方程上，它把约 289 小时的估算训练压到约 2.5 小时，让原本难以完成的实验变得可执行。
我不会把这些数字理解成所有 PDE 都能直接获得百倍加速。真正的价值在于一种可迁移的设计原则：让网络的计算图贴合问题的张量结构，让自动微分模式贴合输入输出维度，让低秩表示承担高维组合中的冗余。
如果未来能进一步处理不规则几何、移动边界、多输入多输出和分支网络分解，这类方法有机会把 physics-informed neural operator 从漂亮的概念，推进到工程级数字孪生和高维科学计算中。
对我来说，最值得记住的一句话是：高维 PDE 的困难不只来自维度本身，也来自我们是否把每个坐标组合都当成了全新的问题。可分离 PI-DeepONet 做的，正是把这些重复结构重新找回来。
参考资料：Computer Methods in Applied Mechanics and Engineering 434 (2025) 117586