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

2026-05-20

AI Summary

This work addresses the susceptibility of physics-informed diffusion models to shortcut learning when boundary conditions shift, a failure mode often caused by insufficient physical constraints on intermediate representations. To mitigate this, the authors propose the REPA-P framework, which employs a lightweight 1×1 convolutional projection head to decode intermediate features into physical quantities and enforces alignment with first-principles physics through PDE residual losses during training. This approach achieves teacher-free, architecture-agnostic physical consistency in intermediate layers without incurring any inference overhead. Compatible with diverse backbones such as U-Net and Diffusion Transformer, REPA-P demonstrates significant improvements across four PDE tasks, accelerating convergence by up to 2×, reducing physical residuals by 66.4%, and enhancing out-of-distribution robustness by 49.3%.

Abstract

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).

Problem

• shortcut learning
• physics-informed diffusion
• intermediate representations
• PDE constraints
• out-of-distribution robustness

Innovation

• representation alignment
• physics-informed diffusion
• shortcut learning
• PDE residuals
• intermediate supervision