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 $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](https://github.com/Hxxxz0/REPA-P).

Many AI models can generate scientific simulations, such as fluid flow, electric fields, or engineering designs. However, these models are often only checked at the final output, so they may learn shortcuts that look correct on familiar examples but fail when the physical setting changes. Our work proposes REPA-P, a training method that encourages the model to follow physical laws not only at the final answer, but also during its internal reasoning process. During training, we add small helper modules that translate hidden features into physical quantities and check whether they obey known physical rules. These helpers are removed after training, so the model does not become slower to use. Across four scientific tasks, REPA-P helps models learn faster, produce results that better satisfy physics, and generalize better to new conditions.