Composing Global Solutions to Reasoning Tasks via Algebraic Objects in Neural Nets¶
会议: NeurIPS 2025
arXiv: 2410.01779
代码: facebookresearch/luckmatters
领域: optimization
关键词: 代数结构, 半环, 环同态, 模加法, Fourier 基, 全局解构造, sum potentials, 两层网络, grokking, 推理任务
一句话总结¶
揭示两层二次激活网络在 Abelian 群推理任务上训练时权重空间具有半环代数结构，提出 CoGS 框架通过环运算将部分解组合为全局最优解，约 95% 梯度下降解与理论构造精确匹配。
研究背景与动机¶
- LLM 推理能力不稳定：大语言模型在基础推理任务上常犯令人意外的错误（如 reversal curse），推理能力的本质仍是开放问题
- 模加法是理解推理的经典探针：预测 \(a+b \mod d\) 结构简单但展现出 grokking 等复杂学习行为，是研究推理机制的理想切入点
- 已知 Fourier 基现象缺乏统一框架：前人发现训练后的网络使用 Fourier 基表示模加法，也手工构造过解析解，但缺少系统性理论框架来解释和推广这些结果
- 现有构造依赖无限宽度假设：Gromov (2023) 的解析构造需假设网络无限宽，与实际有限宽度训练不符
- 梯度下降为何偏好特定解：训练为何收敛到低阶 Fourier 解而非完美记忆解？已有工作未给出理论解释
- 权重空间内部结构未被探索：几何深度学习利用数据对称性，但从未打开黑盒分析训练过程中权重空间本身的代数结构
方法详解¶
框架总览：CoGS (Composing Global Solutions)¶
针对两层二次激活 (\(\sigma(x)=x^2\)) 网络在 Abelian 群乘法预测任务上的 \(L_2\) 损失优化问题，CoGS 发现：(1) 不同宽度网络的权重空间 \(\mathcal{Z}\) 具有半环代数结构；(2) 损失函数由 sum potentials (SP) 构成，SP 是环同态映射；(3) 利用环同态性质，从仅满足部分损失约束的"部分解"出发，通过环加法和环乘法代数组合出全局最优解。
关键设计 1：权重空间的半环结构¶
- 做什么：在不同宽度 \(q\) 的两层网络权重集合 \(\mathcal{Z} = \bigcup_{q \geq 0} \mathcal{Z}_q\) 上定义加法和乘法运算
- 核心思路：加法定义为隐藏节点维度的拼接 (concatenation)，乘法定义为 Khatri-Rao 积（Kronecker 积沿隐藏维），两者满足交换律和分配律，构成交换半环 \(\langle \mathcal{Z}, +, * \rangle\)
- 设计动机：将不同宽度的网络统一在同一代数框架中，是后续所有组合构造的基石；传统工作只分析固定宽度网络，错失了跨宽度的丰富结构
关键设计 2：Sum Potentials 作为环同态¶
- 做什么：证明 \(L_2\) 损失中的核心项——sum potentials \(r(\mathbf{z}) = \sum_j \prod_{p,k} z_{pkj}\)——是从权重半环 \(\mathcal{Z}\) 到复数域 \(\mathbb{C}\) 的环同态映射
- 核心思路：环同态满足 \(r(\mathbf{z}_1 + \mathbf{z}_2) = r(\mathbf{z}_1) + r(\mathbf{z}_2)\) 和 \(r(\mathbf{z}_1 * \mathbf{z}_2) = r(\mathbf{z}_1) \cdot r(\mathbf{z}_2)\)。这意味着若部分解 \(\mathbf{z}_1\) 使某些 SP 为 0、\(\mathbf{z}_2\) 使另一些 SP 为 0，则环乘 \(\mathbf{z}_1 * \mathbf{z}_2\) 可使两者的并集为 0
- 设计动机：将高度非线性的全局优化问题转化为在代数结构上的组合问题，避免直接求解复杂非凸优化
关键设计 3：多项式构造低阶全局解¶
- 做什么：在半环 \(\mathcal{Z}\) 上构造多项式 \(\mathbf{z} = \mathbf{u}^L + \mathbf{c}_1 * \mathbf{u}^{L-1} + \cdots + \mathbf{c}_L\)，从单频生成元 \(\mathbf{u}\) 系统产生部分解，再组合为全局解
- 核心思路：选取对称性好的 order-1 生成元（如使用 3 次或 4 次单位根 \(\omega_3, \omega_4\)），使其在 SP 上产生大量重复值（非 1），从而以低次多项式覆盖更多约束。最终构造出 order-6 Fourier 解 \(\mathbf{z}_{F6}\)（每频率 6 个隐藏节点）和 order-4/6 混合解 \(\mathbf{z}_{F4/6}\)（总阶更低）
- 设计动机：追求低阶解（少隐藏节点），因为梯度动力学在 weight decay 下天然偏好低阶——代数连接的解在拓扑上也连通（Theorem 5），高阶的完美记忆解 (\(d^2\) 阶) 不被动力学青睐
关键设计 4：梯度动力学与过参数化分析¶
- 做什么：分析为何训练偏好 Fourier 解而非完美记忆解，以及过参数化的作用
- 核心思路：(1) Theorem 5 证明若全局解 \(\mathbf{z} = \mathbf{y} * \mathbf{z}'\)，则 \(\mathbf{z}\) 与更低阶的 \(\mathbf{z}'\) 之间存在零损失路径，weight decay 自然推向低阶端；(2) Theorem 6 证明宽度趋于无穷时 SP 的动态解耦，解释过参数化改善训练的原因
- 设计动机：弥补静态解结构分析与实际训练动态之间的鸿沟，为 grokking 等现象提供部分解释
损失函数与训练¶
- 损失：投影 \(L_2\) 损失 + 二次激活函数，损失可解析分解为各频率 \(k\) 上的 sum potentials 函数（Theorem 1）
- 训练：Adam 优化器，学习率 0.01，训练 10000 epochs，配合不同强度的 weight decay 正则化
- 数据：Abelian 群乘法全量合成数据，90%/10% 训练/测试划分
实验关键数据¶
表 1：梯度下降解的阶分布（\(q=512\), weight decay \(5 \times 10^{-5}\)）¶
| 群大小 \(d\) | 不可分解 % | Order-4 (\(\mathbf{z}_{\nu=i} * \mathbf{z}_\xi\)) | Order-6 (\(\mathbf{z}_\nu * \mathbf{z}_{syn}\)) | 其他 % |
|---|---|---|---|---|
| 23 | 0.0% | 47.07% | 39.80% | 1.82% |
| 71 | 0.0% | 72.57% | 21.14% | 2.29% |
| 127 | 1.50% | 82.96% | 14.13% | 0.66% |
→ \(d\) 越大 order-4 解占比越高，与 \(\mathbf{z}_{F4/6}\) 混合解理论一致（仅需 1 个 order-6 频率即可构成全局解）
表 2：理论匹配精度¶
| 指标 | 结果 |
|---|---|
| 可分解率 | ~95%（分解误差 ~0.04，解范数量级为 1） |
| CoGS 结构预测成功率 | ~98%（剩余 ~2% 因训练不充分） |
| \(d=127\) 时 order-6 频率数 | 实测 ~1.26，理论预期 1 |
- Weight decay 越大，解分布越向低阶集中（直至模型崩溃），验证 Theorem 5 的"Occam's Razor"效应
- 移除 \(R_*\) 约束后 order-3 解即可满足，验证了部分解组合的层次结构
亮点¶
- 首次发现训练中的代数结构：揭示权重空间的半环结构和损失中的环同态性质，是全新的分析视角
- 理论与实验高度吻合：约 95% 的梯度下降解可精确匹配理论构造，而非仅统计近似
- 不依赖无限宽度假设：比 Gromov (2023) 更简洁实用，order-6 Fourier 解仅需有限宽度
- 解释了多个实践现象：weight decay 偏好简单解、过参数化有利、完美记忆不被梯度偏好
局限性¶
- 仅限二次激活函数：对 SiLU 等实用激活需 Taylor 展开产生高阶 SP，框架推广尚需验证
- 仅处理 Abelian 群：非交换群（如置换群）的推理任务未涉及
- 限于两层网络：深层网络、Transformer 架构的推广路径不明
- 合成任务评估：仅在模加法等合成基准上验证，与真实推理任务的关联仍需探索
- 未完全解释 grokking：动力学分析为粗粒度特征，从记忆到泛化的相变过程未被完整刻画
与相关工作的对比¶
- Gromov (2023)：手工构造 Fourier 解但依赖无限宽度近似 → CoGS 无需该假设且更系统
- Morwani et al. (2023)：用 max-margin + \(L_{2,3}\) 范数分析代数任务 → CoGS 用标准 \(L_2\) 损失并发现半环结构
- Nanda et al. (2023)：通过机械可解释性提取电路 → CoGS 从代数角度解释电路为何呈现 Fourier 结构
- 几何深度学习 (Bronstein et al., 2021)：在架构中引入对称性 → CoGS 发现训练过程中权重空间自发涌现的代数结构
- Li et al. (2024), Liu et al. (2022)：构造 Transformer 权重实现自动机 → 未验证构造与梯度下降解一致
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ — 首次揭示训练中权重空间的半环代数结构，开创全新分析范式
- 实验充分度: ⭐⭐⭐⭐ — 多种 \(d\) 值下 95% 精确匹配，但仅限合成模加法任务
- 写作质量: ⭐⭐⭐⭐ — 理论推导严谨清晰，但代数记号密集、门槛较高
- 价值: ⭐⭐⭐⭐ — 为理解网络推理能力提供全新代数视角，未来可能影响训练算法设计