Probability-Signature Dynamics: Unpacking Modular Addition Learning Within Two-Layer Networks

2026-10-08

AI Summary

This study addresses how gradient-based training selects Fourier modes from data distributions to support generalization in modular addition tasks, and the unclear mechanisms underlying early loss reduction under label noise. To this end, this work proposes a probabilistic signature framework that translates gradient interactions into conditional statistics. By integrating the discrete Fourier transform with gradient dynamics modeling for two-layer networks, it analytically characterizes the cyclic shift operator. The proposed approach provides a unified explanation for structure emergence and noise effects, elucidating the mechanisms of Fourier sparsity and frequency matching. Furthermore, it successfully predicts the frequencies associated with operators such as XOR and reveals the decoupling mechanism of generalization under noisy conditions.

Abstract

Neural networks trained on modular addition tasks often develop Fourier-structured representations that support exact generalization. While prior work has identified these Fourier circuits, the mechanism by which gradient-based training selects them from the data distribution remains unclear. We address this question using probability signatures, which express leading gradient interactions through conditional statistics of the training distribution. For modular addition, these signatures are cyclic shift operators and are diagonalized by the discrete Fourier transform, yielding approximately decoupled Fourier-mode dynamics. This explains the emergence of Fourier sparsity, frequency matching, and phase alignment. The same framework resolves a puzzle under label noise: corrupted examples can show faster early loss decrease than clean examples, despite lacking a coherent generalization rule. We show that noise increases conditional label collisions, strengthening early shared-coordinate reinforcement. Finally, this method can be applied to other operators. Taking XOR as an example, we observed the predicted frequency in experiments.

Problem

• modular addition
• Fourier representations
• gradient-based training
• label noise
• probability signatures

Innovation

• Probability Signatures
• Modular Addition
• Fourier Representations
• Gradient Dynamics
• Label Noise