Grokking of Diffusion Models: Case Study on Modular Addition

• URL: http://arxiv.org/abs/2604.17673v1
• Date: Mon, 20 Apr 2026 00:02:00 GMT
• Title: Grokking of Diffusion Models: Case Study on Modular Addition
• Abstract summary: We show that diffusion models trained with flow-matching objectives exhibit grokking--delayed generalization after overfitting--on modular addition.<n>Our work provides the mechanistic decomposition of algorithmic learning in diffusion models.
• Score: 20.821116485381154
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Diffusion Model for Manifold Data: Score Decomposition, Curvature, and Statistical Complexity [93.52679474816728]
  Diffusion models have become a leading framework in generative modeling.<n>This paper investigates how diffusion models learn structured data, focusing on two key aspects: statistical complexity and influence of data geometric properties.<n>We also highlight the interplay of manifold curvature with the structures in the score function.
  arXiv  Detail & Related papers  (2026-03-21T04:40:19Z)
• Diffusion models under low-noise regime [3.729242965449096]
  We show that diffusion models are effective denoisers when the corruption level is small.<n>We quantify how training set size, data geometry, and model objective choice shape denoising trajectories.<n>This work starts to address gaps in our understanding of generative model reliability in practical applications.
  arXiv  Detail & Related papers  (2025-06-09T15:07:16Z)
• Simple and Critical Iterative Denoising: A Recasting of Discrete Diffusion in Graph Generation [0.0]
  dependencies between intermediate noisy states lead to error accumulation and propagation during the reverse denoising process.<n>We propose a novel framework called Simple Iterative Denoising, which simplifies discrete diffusion and circumvents the issue.<n>Our empirical evaluations demonstrate that the proposed method significantly outperforms existing discrete diffusion baselines in graph generation tasks.
  arXiv  Detail & Related papers  (2025-03-27T15:08:58Z)
• One-for-More: Continual Diffusion Model for Anomaly Detection [63.50488826645681]
  Anomaly detection methods utilize diffusion models to generate or reconstruct normal samples when given arbitrary anomaly images.<n>Our study found that the diffusion model suffers from severe faithfulness hallucination'' and catastrophic forgetting''<n>We propose a continual diffusion model that uses gradient projection to achieve stable continual learning.
  arXiv  Detail & Related papers  (2025-02-27T07:47:27Z)
• Continuous Diffusion Model for Language Modeling [64.7425225935854]
  Existing continuous diffusion models for discrete data underperform compared to discrete methods.<n>We propose a continuous diffusion model for language modeling that incorporates the geometry of the underlying categorical distribution.<n>Our method outperforms existing discrete diffusion models and approaches the performance of autoregressive models.
  arXiv  Detail & Related papers  (2025-02-17T08:54:29Z)
• Understanding Representation Dynamics of Diffusion Models via Low-Dimensional Modeling [29.612011138019255]
  We study the emergence of unimodal representation dynamics in diffusion models.<n>The unimodality arises from an interplay between denoising strength and class confidence across noise scales.<n>In classification tasks, the presence of unimodal dynamics reliably reflects the generalization of the diffusion model.
  arXiv  Detail & Related papers  (2025-02-09T01:58:28Z)
• Distillation of Discrete Diffusion through Dimensional Correlations [21.078500510691747]
  "Mixture" models are capable of treating dimensional correlations while remaining scalable.<n>Loss functions enable the mixture models to distill such many-step conventional models into just a few steps by learning the dimensional correlations.<n>Results show the effectiveness of the proposed method in distilling pretrained discrete diffusion models across image and language domains.
  arXiv  Detail & Related papers  (2024-10-11T10:53:03Z)
• How Diffusion Models Learn to Factorize and Compose [14.161975556325796]
  Diffusion models are capable of generating photo-realistic images that combine elements which likely do not appear together in the training set. We investigate whether and when diffusion models learn semantically meaningful and factorized representations of composable features.
  arXiv  Detail & Related papers  (2024-08-23T17:59:03Z)
• An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations [21.411327264448058]
  We propose an expectation-maximization (EM) approach to train diffusion models from corrupted observations. Our method alternates between reconstructing clean images from corrupted data using a known diffusion model (E-step) and refining diffusion model weights based on these reconstructions (M-step) This iterative process leads the learned diffusion model to gradually converge to the true clean data distribution.
  arXiv  Detail & Related papers  (2024-07-01T07:00:17Z)
• The Emergence of Reproducibility and Generalizability in Diffusion Models [10.188731323681575]
  Given the same starting noise input and a deterministic sampler, different diffusion models often yield remarkably similar outputs. We show that diffusion models are learning distinct distributions affected by the training data size. This valuable property generalizes to many variants of diffusion models, including those for conditional use, solving inverse problems, and model fine-tuning.
  arXiv  Detail & Related papers  (2023-10-08T19:02:46Z)
• Reflected Diffusion Models [93.26107023470979]
  We present Reflected Diffusion Models, which reverse a reflected differential equation evolving on the support of the data. Our approach learns the score function through a generalized score matching loss and extends key components of standard diffusion models.
  arXiv  Detail & Related papers  (2023-04-10T17:54:38Z)
• Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC [102.64648158034568]
  diffusion models have quickly become the prevailing approach to generative modeling in many domains. We propose an energy-based parameterization of diffusion models which enables the use of new compositional operators. We find these samplers lead to notable improvements in compositional generation across a wide set of problems.
  arXiv  Detail & Related papers  (2023-02-22T18:48:46Z)
• Diffusion Models in Vision: A Survey [73.10116197883303]
  A diffusion model is a deep generative model that is based on two stages, a forward diffusion stage and a reverse diffusion stage.<n> Diffusion models are widely appreciated for the quality and diversity of the generated samples, despite their known computational burdens.
  arXiv  Detail & Related papers  (2022-09-10T22:00:30Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.