Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment
作者: Haozhe Jia, Pengyu Yin, Wenshuo Chen, Shaofeng Liang, Lei Wang, Bowen Tian, Xiucheng Wang, Nanqian Jia, Yutao Yue
分类: cs.LG, cs.CV
发布日期: 2026-05-20
🔗 代码/项目: GITHUB
💡 一句话要点
提出REPA-P框架,通过表征对齐打破科学扩散模型中的捷径学习。
🎯 匹配领域: 支柱四:生成式动作 (Generative Motion)
关键词: 物理信息扩散模型 表征对齐 捷径学习 偏微分方程 科学计算
📋 核心要点
- 现有物理信息扩散模型仅约束最终输出,中间表征易受边界条件变化影响,导致捷径学习。
- REPA-P通过一阶原理残差对齐中间特征与物理状态,避免捷径学习,提升泛化能力。
- 实验表明,REPA-P在多个PDE任务中显著提升收敛速度、降低物理残差并增强分布外鲁棒性。
📝 摘要(中文)
本文提出REPA-P,一种无需教师模型且架构无关的框架,旨在通过一阶原理残差将中间特征与物理状态对齐,从而解决物理信息扩散模型中存在的捷径学习问题。该问题源于现有方法通常仅在最终输出上强制执行偏微分方程(PDE)约束,而中间表征不受约束,容易在边界条件变化时产生捷径学习。REPA-P将轻量级的$1{ imes}1$投影头连接到选定的层,将隐藏激活解码为物理量,并在训练期间应用PDE残差损失。这些投影头在推理时被丢弃,引入零开销。在包括达西流、拓扑优化、静电势和湍流通道流在内的四个PDE任务中,REPA-P加速收敛高达2倍,降低物理残差高达66.4%,并提高分布外鲁棒性高达49.3%,在U-Net和Diffusion Transformer骨干网络上均表现出一致的增益。消融实验表明,监督少量中间层可以捕获大部分好处,并补充输出级别的物理损失。
🔬 方法详解
问题定义:现有物理信息扩散模型主要通过在最终输出上施加偏微分方程(PDE)约束来融入物理信息。然而,中间层的表征缺乏约束,模型容易学习到与特定训练数据相关的捷径,导致在新的边界条件下泛化能力下降。这种捷径学习阻碍了模型真正理解和模拟物理现象的能力。
核心思路:REPA-P的核心思路是通过对中间层的表征进行监督,使其与物理状态对齐,从而迫使模型学习更具物理意义的表征,避免捷径学习。具体来说,REPA-P利用PDE残差作为监督信号,鼓励中间层的表征满足物理定律。这种方法类似于在模型的“思考”过程中进行监督,使其遵循物理原理。
技术框架:REPA-P框架包含以下几个主要步骤:1) 在扩散模型的选定层添加轻量级的$1{ imes}1$投影头。2) 使用这些投影头将隐藏激活解码为物理量(例如,速度、压力等)。3) 计算这些物理量对应的PDE残差。4) 将PDE残差作为损失函数的一部分,用于训练扩散模型。在推理阶段,这些投影头被移除,因此不会增加推理开销。
关键创新:REPA-P的关键创新在于它提供了一种无需教师模型,且架构无关的方法来对齐扩散模型的中间表征与物理状态。与以往主要关注最终输出的物理信息方法不同,REPA-P通过监督中间层,迫使模型学习更具物理意义的表征。此外,REPA-P在推理阶段无需额外的计算开销,使其易于集成到现有的扩散模型中。
关键设计:REPA-P的关键设计包括:1) 轻量级的$1{ imes}1$投影头,用于将隐藏激活解码为物理量。2) PDE残差损失函数,用于衡量中间层表征与物理定律的偏差。3) 选择合适的中间层进行监督,作者发现监督少量中间层即可获得显著的性能提升。4) 投影头在推理时被丢弃,保证零推理开销。
🖼️ 关键图片
📊 实验亮点
实验结果表明,REPA-P在四个PDE任务中均取得了显著的性能提升。具体而言,REPA-P加速收敛高达2倍,降低物理残差高达66.4%,并提高分布外鲁棒性高达49.3%。这些结果表明,REPA-P能够有效打破科学扩散模型中的捷径学习,提升模型的泛化能力和物理一致性。
🎯 应用场景
REPA-P框架可广泛应用于涉及物理建模和仿真的领域,如流体力学、电磁学、热传导等。它能够提升科学扩散模型在复杂物理环境下的预测精度和泛化能力,加速科学发现和工程设计,例如优化飞行器设计、预测气候变化、改进医疗成像等。
📄 摘要(原文)
Physics-informed diffusion models typically enforce PDE constraints only on final outputs, leaving intermediate representations unconstrained and prone to shortcut learning under shifted boundary conditions. We introduce REPA-P, a teacher-free, architecture-agnostic framework that aligns intermediate features with physical states using first-principles residuals. REPA-P attaches lightweight $1{\times}1$ projection heads to selected layers, decodes hidden activations into physical quantities, and applies PDE residual losses during training. These heads are discarded at inference, introducing zero overhead. Across four PDE tasks, including Darcy flow, topology optimization, electrostatic potential, and turbulent channel flow, REPA-P accelerates convergence by up to $2{\times}$, reduces physics residuals by up to $66.4\%$, and improves out-of-distribution robustness by up to $49.3\%$, with consistent gains on both U-Net and Diffusion Transformer backbones. Ablations show that supervising a small set of intermediate layers captures most benefits and complements output-level physics losses. Code is available at https://github.com/Hxxxz0/REPA-P.