A Mechanistic Study of Algorithmic Generalization in Diffusion Models on Modular Arithmetic

arXiv:2604.17673v2 Announce Type: replace Abstract: Despite their empirical success, how diffusion models generalize remains poorly understood from a mechanistic perspective. In this work, we provide a controlled mechanistic account of algorithmic learning in diffusion models. By utilizing the grokking phenomenon on an image-based modular addition task as a tractable testbed, we dissect the model's internal dynamics across two distinct visual regimes. In a single-image regime, mechanistic dissection reveals that the model implements modular addition by encoding operands as periodic representations and composing them at the self-attention--feedforward interface in a way that recovers the trigonometric addition identities. In a diverse-image regime with high intraclass variability, we find that the iterative sampling process partitions the task into an arithmetic computation phase and a subsequent visual refinement phase, with their transition precisely captured by an internal frequency-entropy signal. Together, our analysis clarifies the mechanics of algorithmic generalization of diffusion models in a controlled setting, while leaving open how broadly these mechanisms transfer beyond modular addition.