Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment
摘要
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 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 , reduces physics residuals by up to , and improves out-of-distribution robustness by up to , 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](https://github.com/Hxxxz0/REPA-P).
引用
@article{arxiv.2605.20780,
title = {Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment},
author = {Haozhe Jia and Pengyu Yin and Wenshuo Chen and Shaofeng Liang and Lei Wang and Bowen Tian and Xiucheng Wang and Nanqian Jia and Yutao Yue},
journal= {arXiv preprint arXiv:2605.20780},
year = {2026}
}