Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment

Abstract

Physics-informed diffusion models typically impose PDE constraints only on the final output, leaving intermediate features unconstrained. This can enable shortcut solutions that fit training statistics yet generalize poorly under shifted boundary conditions. We introduce \textbf{REPA-P}, a \emph{teacher-free} physics-informed representation alignment framework that uses first-principles residuals as supervision. REPA-P attaches lightweight projection heads to a few early/mid layers of a diffusion backbone, decodes hidden activations into physical states, and applies PDE and boundary-condition residual losses to these intermediate predictions during training. The heads are discarded at inference, preserving the original architecture and sampling cost. Across three 2D scientific field benchmarks (Darcy flow, topology optimization, and Electrostatic Charge Potential), REPA-P accelerates convergence, reduces physics residuals by up to 80\%, and improves out-of-distribution robustness to boundary-condition shifts while maintaining generation quality with zero inference overhead. Ablations show that supervising only a small set of intermediate layers captures most gains and complements output-level physics losses.