Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2604.17673 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866908978396528640 |
|---|---|
| author | Kim, Joon Hyeok Park, Yong-Hyun Østby, Mattis Dalsætra Gu, Jiatao |
| author_facet | Kim, Joon Hyeok Park, Yong-Hyun Østby, Mattis Dalsætra Gu, Jiatao |
| contents | 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 the task into an arithmetic computation phase followed by a visual denoising phase, separated by a critical timestep threshold. Our work provides the mechanistic decomposition of algorithmic learning in diffusion models, revealing how these models bridge continuous pixel-space generation and discrete symbolic reasoning. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2604_17673 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Grokking of Diffusion Models: Case Study on Modular Addition Kim, Joon Hyeok Park, Yong-Hyun Østby, Mattis Dalsætra Gu, Jiatao Machine Learning 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 the task into an arithmetic computation phase followed by a visual denoising phase, separated by a critical timestep threshold. Our work provides the mechanistic decomposition of algorithmic learning in diffusion models, revealing how these models bridge continuous pixel-space generation and discrete symbolic reasoning. |
| title | Grokking of Diffusion Models: Case Study on Modular Addition |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2604.17673 |