Grokking of Diffusion Models: Case Study on Modular Addition
Despite their empirical success, how diffusion models generalize remains poorly understood from a mechanistic perspective. We demonstrate that diffusion models trained with flow-matching objectives exhibit grokking--delayed generalization after overfitting--on modular addition, enabling controlled analysis of their internal computations. We study this phenomenon across two levels of data regime. In a single-image regime, mechanistic dissection reveals that the model implements modular addition by composing periodic representations of individual operands. In a diverse-image regime with high intraclass variability, we find that the model leverages its iterative sampling process to partition...
| Authors | Joon Hyeok Kim, Yong-Hyun Park, Mattis Dalsætra Østby et al. |
|---|---|
| Primary category | cs.LG |
| Published | 20 April 2026 |
| Last updated at source | 20 April 2026 |
Silytics shows public metadata from arXiv and links to the original. Information may be incomplete or out of date. Last refreshed 17 October 2026. Licenses and terms of the original apply to the work itself. Not financial, legal or professional advice.