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

TL;DR

REPA-P is a novel framework that enhances physics-informed diffusion models by aligning intermediate representations with physical laws, leading to faster convergence and better robustness across various PDE tasks.

Contribution

It introduces a teacher-free, architecture-agnostic method that supervises intermediate features with physical residuals, improving physics-informed learning without additional inference overhead.

Findings

Accelerates convergence by up to 2x across tasks.

Reduces physics residuals by up to 66.4%.

Improves out-of-distribution robustness by up to 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).