PDE-JEPA: 面向参数化 PDE 的隐动力学建模预测表示学习

TL;DR PDE-JEPA 用掩码潜在预测预训练表征，再经几何投影对齐物理场演化，结合物理结构化潜在预测器分解参数无关演化与参数响应，在 9 个 PDE 基准上 ID/OOD 平均提升 33.4%/51.4%。

问题

问题背景

参数化偏微分方程（PDE）的快速求解是科学计算与工程仿真中的关键需求，深度学习代理模型近年来受到广泛关注。当前主流方法（如 FNO、DeepONet 及其变体）多采用重建式预训练，以恢复观测物理场为主要目标。

现有方法局限

重建目标强制模型最小化原始场与解码场的像素级误差，但这并不等同于捕获系统演化的底层规则：

• 模型可能将大量容量用于重建高频/噪声细节，反而忽略低维流形上的动力学结构；
• 对观测噪声和数值误差敏感，学到的潜在表示难以直接支持多步 rollout 或参数外推；
• 不同参数条件下的物理场差异较大，重建式特征可能混淆参数效应与空间模式，导致泛化受限。

为什么这个问题难/重要

PDE 轨迹同时包含状态快照与演化规律双重信息。预测式表示学习（如 JEPA）能保留丰富物理信息，但论文发现仅靠预测不足以保证准确场演化——潜在空间需要显式对齐动力学几何。例如，在 unseen governing parameters 上外推时，系统响应与参数之间存在强非线性耦合，普通潜在预测器容易发散。该问题直接关系代理模型在工程设计优化、实时控制等场景中的可用性，工程价值高。

行业类比

类似视频预测中的 JEPA 用潜在预测替代像素重建，但若潜在空间未对齐时间演化，预测的表示无法解码为物理上合理的未来帧。

核心洞察

• 预测性表示学习揭示了重建目标的局限，但仅靠预训练无法保证准确的场演化，必须显式对齐潜在空间几何。与以往基于自编码器重建物理场的参数化 PDE 表示学习不同，PDE-JEPA 首先通过 masked-latent prediction 预训练编码器，证明该方式能捕获丰富的物理规律；然而实验显示，这种表示虽然信息量大，其潜在轨迹几何与真实物理场演化并不一致，直接用于预测会导致偏差。因此作者引入 geometry projector 对齐潜在轨迹与演化几何，这种“先预测预训练、后几何适配”的路线区别于端到端监督范式，为 PDE 表示学习提供了新的工程切入点。
• 物理结构化分解显著提升对未见参数的泛化能力。PDE-JEPA 将潜在动力学分解为参数无关演化分量和参数相关响应分量，使模型分离通用物理规律与参数特定效应。与常见的黑盒潜在预测器相比，这种结构归纳偏置在参数外推任务中带来平均 51.4% 的提升，表明在潜在空间注入物理结构比单纯增加模型容量更有效。对工程实践而言，这一思路可迁移到其他涉及分布偏移的动力学建模场景，通过分解机制降低对训练参数覆盖的依赖。

方法

输入：给定参数化 PDE 的初始条件场 u0 与物理参数 μ（或一段观测轨迹），目标预测后续时刻的物理场演化。

关键模块分三阶段训练：

1. 预测式表示预训练：编码器将输入场映射到潜在表示，并采用 masked-latent prediction 任务，即随机遮挡部分潜在变量，训练模型从可见潜在预测被遮挡的潜在，以此捕获 PDE 动力学的不变规律。此阶段无需重建原始场。
2. 物理对齐几何投影：引入一个几何投影器，将预训练 latent 轨迹的几何结构与真实物理场演化轨迹的几何结构对齐。通过残差几何投影和物理轨迹对齐损失，同时保持表示的信息量，使 latent 空间与物理动力学状态空间一致。
3. 物理结构化潜在预测器：在几何对齐的 latent 空间上，构建一个分解式预测器，将潜在动力学分解为参数无关的演化分量和参数相关的响应分量。演化分量负责通用时间推进，响应分量根据参数 μ 调节动力学的具体行为。通过连续时间传播，实现多步 rollout 预测。

输出：解码器将预测的潜在变量映射回物理场，得到未来时刻的场预测。训练时解码器与预测器联合优化，推理时可直接从初始条件与参数进行自回归预测。

与同类方法差异：不同于依赖重构损失的表示学习（如 autoencoder），PDE-JEPA 采用联合嵌入预测学习，并通过几何对齐与物理分解显式注入动力学先验，从而在参数外推上显著提升。

实验

实验设计

论文在 9 个一维/二维 PDE benchmark 上评估, 涵盖 Advection、Burgers、Heat、Vorticity、Gray-Scott 等典型方程. 实验分 In-distribution (ID) 与 Out-of-distribution (OOD) 两种协议, OOD 重点考察未见 governing parameters 下的外推能力. 基线包括 reconstruction-based 表示学习与 parametric PDE solver 的 SOTA 方法. 指标为相对 L2 误差, 结果取平均改进.

关键发现

• PDE-JEPA 在 ID 场景下平均误差降低 33.4%, OOD 场景下平均降低 51.4%, 表明预测式预训练能有效捕获物理动态规律.
• 消融显示 Physics-Aligned Geometry (PAG) 模块显著提升 ID 性能; Physics-Structured Latent Predictor (PSP) 通过解耦 parameter-independent evolution 与 parameter-dependent response, 大幅增强 OOD 外推.
• 预测表征本身保留丰富物理信息, 但需几何对齐才能准确预测场演化.

与基线对比解读

与 reconstruction-based 预训练相比, PDE-JEPA 不依赖重建原始场, 而是通过 masked latent prediction 学习动态规律, 降低对高频噪声的敏感性. 但其直接预测潜在表征不足以保证演化精度, 因此引入几何投影器将潜在轨迹对齐到物理场演化几何, 使解码器更易恢复动态. PSP 的分解设计类似于将 PDE 解分为齐次部分与特解, 从而在分布外参数下保持结构一致性, 这解释了 OOD 提升更大的原因. 工程上, 该框架的预训练-对齐-预测三阶段可复用, 适合接入现有 PDE 求解流程.

行业影响

落地场景

PDE-JEPA 的潜空间预测能力，适合嵌入 CAE 仿真软件、数字孪生平台、气象海洋预报系统、新能源热管理 与 药物分子动力学筛选 等场景。例如，在车辆空气动力学优化中，PDE-JEPA 可快速预测不同几何或工况参数下的流场，替代部分高成本 CFD 求解；在电池热管理中，针对不同充放电倍率与材料参数，实时输出温度场分布，支撑在线控制。

商业价值

• 降本：减少高精度数值求解器调用次数，将小时级仿真缩短至秒级推理，降低计算集群或云仿真成本。
• 增收：以低延迟代理模型嵌入交互式设计工具，提升设计师迭代速度，加速产品上市周期。
• 体验提升：数字孪生场景中提供实时物理场反馈，支撑预测性维护与风险预警。

与现有工作流接口

• 作为 代理模型 接入现有求解器前端：输入参数与初始/边界条件，输出预测轨迹，供下游优化器调用。
• 通过 预训练编码器 提取物理状态嵌入，可与现有 Transformer / GNN 架构结合，无需重新设计特征工程。
• 输出的物理场可转换为 VTK / OpenFOAM 兼容网格，集成到 ParaView 等后处理流程。

具体落地 use case：

1. 自动驾驶空气动力学开发：将 PDE-JEPA 嵌入车辆外形优化管线，快速评估不同车身参数下的风阻与升力场，减少风洞试验与大涡模拟（LES）次数。
2. 医疗血流动力学评估：面向个性化血管模型，快速预测不同血压/黏度参数下的血流速度与壁面剪切应力，辅助支架设计与手术规划。

局限

• **训练流程复杂与计算开销**：方法包含 masked-latent prediction 预训练、geometry projector 对齐、physics-structured latent predictor 和解码器等多个阶段，整体训练管道比现有 reconstruction-based 方法（如 FNO、GNOT）或端到端 latent evolution 更长，需要调节的超参数更多。几何投影需要额外计算物理场演化几何，对网格分辨率和边界条件可能敏感，在 2D/3D 大规模场景下推理延迟与显存占用未在论文中给出详细对比，实际部署效率存疑。
• **外推能力边界未充分探索**：虽然报告对未见控制参数平均提升 51.4%，但实验中的参数外推可能局限于训练分布附近的插值或小范围外推。当参数大幅偏离训练域（例如雷诺数跨越多个数量级或方程非线性显著变化）时，参数无关演化与参数依赖响应的分解假设可能失效，导致潜在轨迹几何对齐误差累积。论文未系统分析外推失败案例、适用范围及参数敏感性，可能高估该方法的鲁棒性。
• **与同类预测式表示方法对比不足**：主要对比基线集中在 reconstruction-based 求解器和少量 latent evolution 方法，缺少与近期 latent JEPA 变体（如 MC-JEPA、V-JEPA 在物理场预测上的适配）的直接对照，无法完全剥离 geometry projector 与 physics-structured predictor 各自的独立贡献。消融实验虽然涉及 PAG，但未对 masked latent prediction 的 mask 策略、几何对齐损失权重、潜在维度等关键超参数进行充分敏感性分析，限制了方法可复现性与调参指导。