模加法学习的机制与动力学：傅里叶特征、彩票机制与Grokking On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking
揭示两层神经网络学习模加法的傅里叶机制、特征涌现动力学与grokking三阶段过程
前置知识
模加法任务 (Modular Addition)
模加法任务是指学习映射 $(x, y) \mapsto (x + y) \bmod p$，其中 $x, y \in \mathbb{Z}_p$，$p$ 为素数。完整数据集 $\mathcal{D}_{\text{full}}$ 包含所有 $p^2$ 个可能的输入对及其对应的模和。该任务虽然形式简单，但能揭示神经网络学习中特征学习、泛化和相变等深层现象，因此成为深度学习理论研究的经典测试平台。
模加法是研究神经网络机制可解释性的标准任务，理解网络如何学会这个任务是本文的核心目标。
离散傅里叶变换 (DFT) 与傅里叶特征
离散傅里叶变换将时域信号分解为不同频率的正弦和余弦分量的叠加。在本文中，DFT 被应用于神经网络的权重参数 $\theta_m$ 和 $\xi_m$，将其投影到傅里叶基 $\mathbf{B}_p$ 上，得到傅里叶系数 $g_m = \mathbf{B}_p^\top \theta_m$。每个频率 $k$ 对应一个幅度 $\alpha_{km}$ 和相位 $\phi_{km}$，这种极坐标表示直接关联到相移余弦函数。傅里叶特征是理解网络如何将加法转化为几何旋转的关键工具。
论文的核心分析方法就是通过DFT在频域中追踪参数演化，理解傅里叶特征是读懂全文技术细节的前提。
Grokking 现象
Grokking 是指神经网络在训练数据上先完全过拟合（训练准确率100%但测试准确率接近0），经过大量额外训练步数后突然泛化（测试准确率突然跳升至接近100%）的现象。这一现象由 Power 等人在 2022 年首次系统描述，在模加法等简单算术任务中尤为显著。其背后的原因涉及权重衰减、特征稀疏化等机制的复杂交互。
本文将 grokking 刻画为三阶段过程，是论文的重要贡献之一，理解该现象的定义和背景是必要的。
相位对齐 (Phase Alignment)
在傅里叶特征的语境下，相位对齐是指神经网络第一层参数 $\theta_m$ 的相位 $\phi_m$ 和第二层参数 $\xi_m$ 的相位 $\psi_m$ 之间满足 $\psi_m = 2\phi_m \bmod 2\pi$ 的关系，即输出相位是输入相位的两倍。这种相位耦合关系在训练过程中动态涌现，是网络能够正确实现模加法算法的关键条件。相位差异 $D_m = 2\phi_m - \psi_m$ 的演化由梯度流驱动，最终收敛到零。
相位对齐是本文观察到的核心现象之一，也是推导多数理论结果的关键假设，理解它对把握论文技术主线至关重要。
彩票机制 (Lottery Ticket Mechanism)
本文中的彩票机制描述了在随机初始化下，每个神经元内多个频率分量之间的竞争过程。每个频率 $k$ 携带一张由初始幅度 $(\alpha_{km}(0), \beta_{km}(0))$ 和初始相位失配 $\tilde{D}_{km}(0)$ 决定的'彩票'。在训练过程中，幅度增长速率正比于 $\cos(D_{km})$，相位旋转速率正比于 $\sin(D_{km})$，形成'更好对齐→更快增长→更快对齐'的正反馈循环。最终，初始幅度最大且相位失配最小的频率赢得竞争，成为该神经元的主导频率。
彩票机制是解释傅里叶特征如何从随机初始化中涌现的核心理论机制，是论文的另一重要贡献。
过参数化与全多样化条件 (Full Diversification)
全多样化条件定义了网络在过参数化（神经元数 $M$ 足够大）时自然学到的理想结构，包含三个子条件：(i) 频率多样化——每个频率 $k \in [(p-1)/2]$ 都恰好有 $N$ 个神经元表示；(ii) 幅度同质性——所有神经元的 $\alpha_m \beta_m^2 = a$（常数）；(iii) 相位对称性——每个频率组内高阶相位矩 $\exp(i \cdot \iota \sum_{m \in N_k} \phi_m) = 0$ 对 $\iota \in \{2, 4\}$ 成立。这些条件共同保证网络能通过多数投票机制抵消单个神经元的噪声，逼近正确的指示函数。
全多样化条件是连接局部神经元行为与全局网络功能的桥梁，是理解论文机制解释的核心概念。
研究动机
现有研究虽然已发现神经网络在模加法任务上学习傅里叶特征表示，将输入嵌入圆上并利用几何旋转实现加法，但这些工作存在明显不足。具体而言，Nanda 等人（2023）和 Zhong 等人（2023）的逆向工程工作识别出单个神经元学习单频率特征并展现相位对齐，但未能充分解释这些局部特征如何组合成全局解决方案。Wang 和 Wang（2025）的均场近似方法忽略了有限神经元层面的动力学细节；Morwani 等人（2023）和 Tian（2024）分析的是非标准损失函数（如 $\ell_{2,3}$ 正则化或修改后的 $\ell_2$ 损失），与标准训练设置存在差距。此外，关于 grokking 现象的理论解释也缺乏一个从记忆化到泛化的完整端到端叙事。这些理论空白导致三个根本问题悬而未决：(Q1) 训练好的网络如何精确利用傅里叶特征实现模加法算法？(Q2) 这些特定傅里叶特征如何从随机初始化的梯度训练中可靠涌现？(Q3) 这些机制如何解释 grokking 的完整时间线？
本文的目标是本文的目标是提供一个完整的端到端逆向工程，解释两层神经网络如何从训练动力学到最终学到的模型来学习模加法。具体而言，作者旨在：(1) 在机制可解释性层面，精确定量刻画网络如何将局部傅里叶特征综合为全局解决方案，证明全多样化条件下的多数投票机制能逼近指示函数；(2) 在训练动力学层面，通过梯度流分析严格证明单频率模式的保持和相位对齐的涌现，解释'彩票机制'如何从随机初始化决定每个神经元最终学习的频率；(3) 在 grokking 层面，将其刻画为损失最小化与权重衰减之间竞争驱动的三阶段过程，覆盖从记忆化到延迟泛化的完整时间线。
与已有工作不同的是，本文的独特切入角度在于同时从机制解释（静态）和动力学（动态）两个互补视角进行分析，填补了现有工作的理论空白。与均场理论不同，本文直接在有限神经元层面工作，避免了均场近似带来的精度损失；与非标准损失函数的分析不同，本文在标准交叉熵损失下进行推导；与仅关注梯度下降收敛性的理论工作不同，本文追踪参数在傅里叶域中的演化轨迹，揭示了频率竞争和相位耦合的微观机制。此外，本文的关键技术洞察是利用二次激活的可分析性来建立理论，同时证明 ReLU 训练的模型在推理时对激活函数替换具有鲁棒性（保持 $|x|, x^2, x^4$ 等偶数阶分量时准确率达100%），从而建立了理论分析与实际训练之间的桥梁。这种从实验观察出发、以理论证明回应的双向方法论也是本文区别于纯理论工作的特色。
核心方法
本文的方法论遵循'实验观察→理论建模→严格证明'的路线。整体思路是：首先，通过大量实验（$p=23, M=512$, ReLU, AdamW）观察训练后网络权重的傅里叶结构，提炼出六个关键实验观察（Observation 1-6）；然后，利用这些观察构建理想化的机制模型，证明在全多样化条件下网络通过多数投票逼近模加法的指示函数；接着，在梯度流框架下分析训练动力学，证明单频率保持和相位对齐的涌现；最后，将彩票机制严格化，并扩展到 ReLU 激活和 grokking 场景。技术上，核心工具是离散傅里叶变换（DFT），将参数从时域投影到频域，使得每个频率分量的演化可以独立追踪。在理论分析中采用二次激活 $\sigma(x) = x^2$ 以获得解析解，同时通过 Observation 4 的激活替换鲁棒性实验保证理论结论对 ReLU 也适用。
本文的核心创新在于发现并形式化了'全多样化条件'（Definition 4.1），它由频率多样化和相位对称两部分组成，是连接单个神经元的局部傅里叶行为与全局网络功能的关键桥梁。与已有方法的本质区别在于：(1) 已有工作（如 Kumar 等人 2024）仅将相位均匀性作为构造性假设引入，本文通过实验验证这是过参数化下的稳定涌现性质，并推导出比严格均匀性更弱但足以支撑精确分析的条件；(2) 已有理论分析停留在均场层面或非标准损失，本文首次在有限神经元、标准交叉熵损失下给出完整证明；(3) 关于 grokking，本文提出三阶段刻画（记忆化→泛化I→泛化II），其中泛化I阶段由损失最小化和权重衰减共同驱动，泛化II阶段由权重衰减主导，这一细致分阶段比之前的二阶段（惰性/富特征）描述更精确。
方法步骤详情
方法分为三个层次展开。**层次一：机制解释**。给定训练好的模型参数 $\theta_m, \xi_m$，对每个神经元应用 DFT，发现每个神经元在频域中仅有一个主导频率 $\phi(m)$，且满足相位对齐 $\psi_m = 2\phi_m$。在全多样化条件下（频率覆盖全部 $(p-1)/2$ 个频率，相位在组内均匀分布），将所有神经元的贡献求和，推导出输出 logit 的形式为 $$f(x,y;\xi,\theta)[j] = \frac{aN}{2}\left[-1 + \frac{p}{2} \cdot \mathbf{1}(x+y \bmod p = j) + \frac{p}{4}\sum_{z \in \{x,y\}} \mathbf{1}(2z \bmod p = j)\right]$$ 其中信号项在 $j = (x+y) \bmod p$ 处峰值超过噪声项 $aNp/8$，经 softmax 后网络输出集中于正确答案。**层次二：训练动力学**。在梯度流 $\partial_t \Theta_t = -\nabla \ell(\Theta_t)$ 下，通过小初始化假设（$\kappa_{\text{init}}$ 足够小），证明神经元解耦——每个神经元的动力学仅依赖自身参数。然后在频域中追踪傅里叶系数的演化：(a) 利用傅里叶基的正交性证明单频率保持（Theorem 5.2）；(b) 推导相位差异 $D_m^\star = 2\phi_m^\star - \psi_m^\star$ 的 ODE，证明其单调收敛到零（Theorem 5.3）。**层次三：彩票机制**。将分析扩展到多频率初始化，利用 ODE 比较引理证明：在 $(p-1)/2$ 个频率的竞争中，初始相位失配 $\tilde{D}_{km}(0)$ 最小的频率 $k^\star$ 以 $\tilde{O}(\log p / (p\kappa_{\text{init}}))$ 的时间尺度主导所有其他频率（Corollary 6.1）。
技术新颖性
本文的技术新颖性体现在多个层面。首先，在理论工具上，作者创新性地将 ODE 比较引理（Smith 1995）应用于频率竞争分析，这在神经网络理论中是新颖的。其次，在数学建模上，定义了相位差异 $D_m = 2\phi_m - \psi_m$ 作为关键变量，发现它同时控制幅度增长速率（通过 $\cos D_m$）和相位旋转速度（通过 $\sin D_m$），形成了'对齐越好→增长越快→对齐更快'的正反馈闭环。第三，在实验发现上，Table 1 展示的激活替换鲁棒性极具洞察力：模型在 ReLU、$|x|$、$x^2$、$x^4$、$\log(1+e^{2x})$、$e^x$ 上保持100%准确率，但在 $x$ 和 $x^3$ 上准确率接近0%，这证明关键性质是激活函数的偶数阶分量而非具体形式，从而为用二次激活做理论分析提供了坚实的经验依据。第四，在 grokking 解释上，提出的'记忆化原则：从常见到稀有'（Common-to-Rare）揭示了训练早期模型先记忆对称数据对 $(i,j)$ 和 $(j,i)$，再记忆仅出现一次的稀有样本的分阶段策略，这比笼统的'先记忆后泛化'描述更细致。
实验结果
本文的核心发现可归纳为三个层面。**机制解释层面**：Proposition 4.2 证明，在全多样化条件下，网络输出 logit 的信号项 $\frac{ap}{4} \cdot \mathbf{1}(x+y \bmod p = j)$ 超过噪声项 $aNp/8$，取 $a \gtrsim (Np)^{-1} \log(p/\epsilon)$ 时 softmax 输出与真实标签的 $\ell_{1,\infty}$ 距离不超过 $\epsilon$。消融实验（§A.2）进一步表明全多样化参数化在相同神经元预算下具有最大的 logit 间隙，是最参数高效的方案。**训练动力学层面**：Theorem 5.2 证明在小初始化 ($\kappa_{\text{init}} = o(M^{-1/3})$) 下，非特征频率在初始阶段保持 $o(\kappa_{\text{init}})$ 的量级。Theorem 5.3 证明相位对齐时间 $t_\delta \asymp (p\kappa_{\text{init}})^{-1} \cdot \left[1 - (\sin D_m^\star(0)/\delta)^{-1/3} + \max\{\pi/2 - |D_m^\star(0) - \pi|, 0\}\right]$，且对齐后参数幅度仍保持 $O(\kappa_{\text{init}})$ 量级。Corollary 6.1 证明彩票胜出频率 $k^\star = \arg\min_k \tilde{D}_{km}(0)$ 以 $\tilde{O}(\log p/(p\kappa_{\text{init}}))$ 时间主导其他频率，成功概率至少 $1 - e^{-\Theta(p^{-c})}$。**Grokking 层面**：实验将 grokking 刻画为三阶段过程——记忆化阶段（步数 0-1000，训练损失快速降到零）、泛化I阶段（步数 1000-10000，权重衰减修剪非特征频率，测试损失急剧下降）、泛化II阶段（步数 10000+，权重衰减主导，测试准确率缓慢趋近完美）。频率稀疏度的逆参与比（IPR）在泛化I阶段快速上升，证实了权重衰减的频率域稀疏化效应。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 模加法学习（$p=23$, $M=512$, ReLU, 全数据集） | 训练/测试准确率 | 训练和测试均达到100%，损失 $1.194 \times 10^{-8}$ | 线性网络或宽度不足的网络 | 从 0.041 准确率提升至 1.000 |
| 激活函数替换鲁棒性 | 替换 ReLU 后的测试准确率 | $|x|$, $x^2$, $x^4$, $e^x$, $\log(1+e^{2x})$ 均达100%准确率 | 原始 ReLU 激活（准确率100%） | 证明模型依赖偶数阶分量而非具体激活形式 |
| Grokking 泛化（$p=23$, 75%训练划分，权重衰减 2.0） | 测试准确率 | 经过约 20000 步后测试准确率趋近100% | 无权重衰减的训练（测试准确率始终接近0%） | 揭示权重衰减驱动的三阶段泛化过程 |
| 全多样化条件消融（§A.2） | Logit 间隙（正确标签 vs 最高错误标签） | 全多样化参数化具有最大 logit 间隙 | 有限频率或非均匀相位的参数化 | 证明全多样化是最参数高效的方案 |
局限与改进
本文存在以下局限性。首先，理论分析主要基于二次激活 $\sigma(x) = x^2$，虽然通过 Observation 4 的鲁棒性实验建立了与 ReLU 的联系，但严格的 ReLU 理论仅给出了频率泄漏的渐近阶（Proposition 6.3：泄漏衰减为 $\Theta(r_k^{-2})$），未能完全复现二次激活下的精确结论。其次，全多样化条件（Definition 4.1）假设了理想的频率平衡和相位均匀性，虽然实验表明这些条件在 $M$ 足够大时近似成立，但对有限 $M$ 下条件违反程度的量化分析不足。第三，分析局限于两层网络和模加法这一特定任务，对更深网络（如 Transformer）和更复杂算术任务（如模乘法）的推广尚待验证。第四，grokking 的三阶段刻画虽然比之前更细致，但阶段之间的精确边界（如步数 1000 和 10000）依赖于特定超参数设置（$p=23$, 权重衰减 2.0），其普适性需要更多实验验证。最后，论文的 Theorem 5.2 和 5.3 都有'非正式'版本，正式版本的条件假设较多（如 $\log M / M \lesssim c^{-1/2}$, $\kappa_{\text{init}} = o(M^{-1/3})$），实际训练中这些条件是否严格满足值得商榷。
独立分析的弱点
本文有几个值得注意的弱点。第一，理论与实验之间存在'激活函数鸿沟'：理论在二次激活下建立，实验在 ReLU 下进行，虽然 Observation 4 提供了桥梁，但 ReLU 下的频率泄漏（Observation 7）使得单频率假设仅为近似成立，这削弱了理论结论的精确性。改进方向是发展 ReLU 激活下的严格动力学理论，可能需要处理 $|x|$ 项引入的非光滑性。第二，全多样化条件中的频率平衡假设（每个频率恰好 $N$ 个神经元）在实际随机初始化中仅近似成立，网络中可能出现某些频率过度代表或欠代表的情况，论文对此缺乏鲁棒性分析。改进方向是引入随机化的频率分布理论，分析频率覆盖不均匀时的性能退化。第三，grokking 分析中的记忆化'从常见到稀有'原则虽然直观，但缺乏严格的理论证明来解释为什么对称数据对会被优先记忆。第四，论文未讨论学习率、优化器选择等超参数对所发现机制的影响，AdamW 中的动量和自适应学习率可能与理论假设的梯度流存在差异。
未来方向
作者在论文中提出的未来方向包括将分析扩展到更一般的激活函数（§6.2 已开始）和更深的网络架构。基于本文成果，可延伸的研究方向包括：(1) 将傅里叶分析框架推广到模乘法、置换群上的运算等更复杂的代数结构，Tian（2024）已开始这一方向但缺乏完整的动力学分析；(2) 研究 Transformer 架构中注意力机制如何与傅里叶特征交互，特别是在 few-shot 算术推理任务中；(3) 探索彩票机制在更一般特征学习任务中的适用性，如单索引模型和多索引模型中的特征竞争；(4) 将 grokking 的三阶段理论与实际大模型训练中的相变现象（如涌现能力的突然出现）建立联系；(5) 设计基于理论洞察的训练策略，如利用相位对齐信息初始化网络或设计频率感知的正则化方法来加速泛化。
复现评估
本文的复现条件较为友好。代码已在 GitHub 开源，交互式可视化可通过 HuggingFace Space 访问。实验设置明确：$p=23$, $M=512$, AdamW 优化器，学习率 $\eta = 10^{-4}$，PyTorch 默认初始化后归一化。主要实验在单 GPU 上即可完成，计算需求不高。理论分析的关键假设（小初始化、梯度流、二次激活）都有明确的数学定义。但需要注意：(1) grokking 实验需要权重衰减 2.0 和 75% 训练划分，这些超参数的选择对 grokking 的出现至关重要；(2) 理论证明涉及大量技术引理（附录 B 和 C 共 30+ 页），完整验证需要相当的数学背景；(3) 消融实验（§A.2, §A.3）的结果依赖于特定的 $p$ 和 $M$ 值，跨超参数的系统性验证需要额外工作。总体而言，核心实验复现难度适中，理论验证难度较高。
论文图表
