# Grokking of Diffusion Models: Case Study on Modular Addition

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 47 of its 54 references that are in the map.
- First posted on arXiv 20 Apr 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 47)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [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)
- [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)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Genie: Generative Interactive Environments](https://mlsanity.com/paper/2402.15391/) (2024)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Deep Learning for Classical Japanese Literature](https://mlsanity.com/paper/1812.01718/) (2018)
- [Diffusion Models already have a Semantic Latent Space](https://mlsanity.com/paper/2210.10960/) (2022)
- [Video models are zero-shot learners and reasoners](https://mlsanity.com/paper/2509.20328/) (2025)
- [EMNIST: an extension of MNIST to handwritten letters](https://mlsanity.com/paper/1702.05373/) (2017)
- [Generalization in diffusion models arises from geometry-adaptive harmonic representations](https://mlsanity.com/paper/2310.02557/) (2023)
- [HunyuanVideo 1.5 Technical Report](https://mlsanity.com/paper/2511.18870/) (2025)
- [The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks](https://mlsanity.com/paper/2306.17844/) (2023)
- [Omnigrok: Grokking Beyond Algorithmic Data](https://mlsanity.com/paper/2210.01117/) (2022)
- [gDDIM: Generalized denoising diffusion implicit models](https://mlsanity.com/paper/2206.05564/) (2022)

## Abstract (arXiv)

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.