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

arXiv: https://arxiv.org/abs/2605.20780 · HTML: https://mlsanity.com/paper/2605.20780/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 22 of its 31 references that are in the map.
- First posted on arXiv 20 May 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 22 of 22)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification](https://mlsanity.com/paper/1801.06879/) (2018)
- [Fast Sampling of Diffusion Models via Operator Learning](https://mlsanity.com/paper/2211.13449/) (2022)
- [DCTdiff: Intriguing Properties of Image Generative Modeling in the DCT Space](https://mlsanity.com/paper/2412.15032/) (2024)
- [Physics-Informed Representation Alignment for Sparse Radio-Map Reconstruction](https://mlsanity.com/paper/2501.19160/) (2025)
- [ANT: Adaptive Neural Temporal-Aware Text-to-Motion Model](https://mlsanity.com/paper/2506.02452/) (2025)
- [POLARIS: Projection-Orthogonal Least Squares for Robust and Adaptive Inversion in Diffusion Models](https://mlsanity.com/paper/2512.00369/) (2025)
- [LUMA: Low-Dimension Unified Motion Alignment with Dual-Path Anchoring for Text-to-Motion Diffusion Model](https://mlsanity.com/paper/2509.25304/) (2025)
- [RadioFlow: Efficient Radio Map Construction Framework with Flow Matching](https://mlsanity.com/paper/2510.09314/) (2025)

## Abstract (arXiv)

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