mlsanityData to 7 Oct 2026

Grokking of Diffusion Models: Case Study on Modular Addition

arXiv 2604.17673 ↗First posted 20 Apr 2026cs.LGcited by 0

On the map: Diffusion models → Diffusion foundations and sampling

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.

Key facts. No citing papers in the map yet (data through 7 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.

Builds on

47 of 54 references.

Closest ancestors · 15

Same direction, newest first

  1. 2025Selective Underfitting in Diffusion Models1 Oct
  2. 2025On the Edge of Memorization in Diffusion Models25 Aug
  3. 2025Memorization to Generalization: Emergence of Diffusion Models from Associative Memory27 May
  4. 2025Why Diffusion Models Don't Memorize: The Role of Implicit Dynamical Regularization in Training23 May
  5. 2024One-Step is Enough: Sparse Autoencoders for Text-to-Image Diffusion Models28 Oct
  6. 2024Going beyond Compositions, DDPMs Can Produce Zero-Shot Interpolations29 May
  7. 2023Compositional Abilities Emerge Multiplicatively: Exploring Diffusion Models on a Synthetic Task13 Oct
  8. 2023Generalization in diffusion models arises from geometry-adaptive harmonic representations4 Oct
  9. 2023Understanding the Latent Space of Diffusion Models through the Lens of Riemannian Geometry24 Jul
  10. 2022Scalable Diffusion Models with Transformers19 Dec
  11. 2022Diffusion Models already have a Semantic Latent Space20 Oct
  12. 2022gDDIM: Generalized denoising diffusion implicit models11 Jun
3 more
  1. 2021High-Resolution Image Synthesis with Latent Diffusion Models20 Dec
  2. 2020Score-Based Generative Modeling through Stochastic Differential Equations26 Nov
  3. 2020Denoising Diffusion Probabilistic Models19 Jun

Other references · 32

  1. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  2. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
  3. 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,539
  4. 2022Photorealistic Text-to-Image Diffusion Models with Deep Language Understandingcited by 6,322
  5. 2022Flow Matching for Generative Modelingcited by 5,426
  6. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  7. 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,550
  8. 2025Wan: Open and Advanced Large-Scale Video Generative Modelscited by 2,536
  9. 2020GLU Variants Improve Transformercited by 1,724
  10. 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,359
  11. 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,130
  12. 2022In-context Learning and Induction Headscited by 930
20 more
  1. 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
  2. 2022Transformers learn in-context by gradient descentcited by 714
  3. 2024Genie: Generative Interactive Environmentscited by 706
  4. 2023Towards Automated Circuit Discovery for Mechanistic Interpretabilitycited by 700
  5. 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
  6. 2018Deep Learning for Classical Japanese Literaturecited by 415
  7. 2017EMNIST: an extension of MNIST to handwritten letterscited by 212
  8. 2025Video models are zero-shot learners and reasonerscited by 208
  9. 2023The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networkscited by 153
  10. 2025HunyuanVideo 1.5 Technical Reportcited by 148
  11. 2022Omnigrok: Grokking Beyond Algorithmic Datacited by 142
  12. 2023Explaining grokking through circuit efficiencycited by 105
  13. 2024An analytic theory of creativity in convolutional diffusion modelscited by 103
  14. 2023Grokking as the Transition from Lazy to Rich Training Dynamicscited by 102
  15. 2024Understanding Hallucinations in Diffusion Models through Mode Interpolationcited by 87
  16. 2023A Survey on Audio Diffusion Models: Text To Speech Synthesis and Enhancement in Generative AIcited by 68
  17. 2023Unifying Grokking and Double Descentcited by 49
  18. 2024Emergence of Hidden Capabilities: Exploring Learning Dynamics in Concept Spacecited by 33
  19. 2025Towards Understanding Text Hallucination of Diffusion Models via Local Generation Biascited by 11
  20. 2025Sparse Autoencoder as a Zero-Shot Classifier for Concept Erasing in Text-to-Image Diffusion Modelscited by 8

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Understanding diffusion models requires rethinking (again) generalization
  2. 2024Towards a Mechanistic Explanation of Diffusion Model Generalization
  3. 2023On the Generalization Properties of Diffusion Models
  4. 2025Diffusion Models: A Mathematical Introduction
  5. 2025Generalization of Diffusion Models Arises with a Balanced Representation Space
  6. 2025From Navigation to Refinement: Revealing the Two-Stage Nature of Flow-based Diffusion Models through Oracle Velocity
  7. 2025The Principles of Diffusion Models
  8. 2024Tracing the Roots: Leveraging Temporal Dynamics in Diffusion Trajectories for Origin Attribution