规模化释放门控 Delta 网络中的特征学习能力 Unlocking Feature Learning in Gated Delta Networks at Scale
推导 Gated Delta Networks 的 µP 参数化法则，实现零样本超参数迁移。
前置知识
Maximal Update Parametrization (μP)
μP 是一种神经网络参数化方法，基于 Tensor Programs 框架推导。它的核心思想是在无限宽度极限下，通过合理设置权重初始化方差和学习率缩放规则，使得每层的特征更新保持有界且非零消失（即特征更新量级为常数）。在标准参数化（如 He/Kaiming 初始化）下，为了避免激活发散，学习率必须以 O(1/d) 缩放，这迫使网络进入 NTK（懒惰训练） regime，特征无法有效学习。μP 通过宽度相关的参数化，允许在小模型上调优的超参数零样本迁移到大模型，从而极大地降低了大规模模型训练的超参数搜索成本。
本文的核心目标就是为 Gated Delta Network 推导 μP 参数化公式，理解 μP 的基本原理（特别是坐标尺度传播和特征学习条件）是读懂本文理论推导部分的前提。
Gated Delta Network (GDN)
GDN 是一种高效的线性 Transformer 变体，结合了 Delta Rule（通过减去当前键值对的预测误差来更新快速权重矩阵）和 Mamba-2 的数据依赖门控机制。其潜状态更新规则为 St = St-1(αt(I - βt kt kt^T)) + βt vt kt^T，其中 αt ∈ (0,1) 是遗忘门控，βt ∈ (0,1) 是写入强度，St ∈ R^(d×d) 是全矩阵潜状态。输出直接通过 ot = St qt 读取。GDN 通过外积更新而非标量递归，在保持线性时间推理复杂度（O(1) 每步）的同时，具有强大的长序列建模能力。
本文的研究对象正是 GDN。理解其状态更新方程、门控机制（αt, βt）以及与 Mamba 等对角 SSM 的区别，对于理解为何标准 μP 无法直接应用以及本文推导的特殊缩放规则至关重要。
Coordinate Size / 坐标大小
在 μP 理论中，坐标大小用于量化向量或矩阵中每个元素的量级。一个向量 v ∈ R^d 具有 Θ(d^a) 大小的坐标，如果其 L2 范数的平方除以 d 为 Θ(d^(2a))，即每个元素的方差为 Θ(d^(2a))。特别地，如果其 L2 范数为 Θ(√d)，则称其具有 Θ(1) 坐标大小。在推导 μP 缩放规则时，关键在于通过前向传播（从输入到输出）和反向传播（从损失到参数）传播这些坐标尺度估计，确保所有中间激活保持 Θ(1) 坐标大小，从而支持稳定的特征学习。
本文的理论推导完全基于坐标大小估计。无论是分析潜状态 St、门控参数 αt 还是梯度大小，都依赖于坐标尺度的传播。掌握这一概念是理解本文为何能得出特定缩放规则的数学基础。
研究动机
现有的 Maximal Update Parametrization (μP) 理论虽然在标准 Transformer 上取得了巨大成功，但在扩展到线性模型（特别是具有结构化状态转移的架构）时面临根本性挑战。核心问题在于这些模型的递归状态是通过序列维度更新的，这种动态机制不符合 μP 在标准前馈或基于注意力的架构中的推导假设。此前针对 Mamba 等对角状态空间模型（SSM）的研究表明，标准 μP 和谱缩放条件都无法支持其中的特征学习。Gated Delta Networks 与对角 SSM 有根本区别，因为其状态通过外积 Delta 规则（vt kt^T）更新，且包含额外的数据依赖标量门控，使得现有的 SSM 特定分析无法直接套用，导致 GDN 在大规模训练时的参数化问题完全未被解决。
本文的目标是本文的目标是为 Gated Delta Networks 正式推导完整的 μP 参数化公式，涵盖所有权重类别的初始化方差、前向乘数和学习率缩放规则。具体而言，研究旨在确定投影权重（Wq, Wk, Wv）、门控权重矩阵（Wα, Wβ）以及标量门控参数（alog, b）所需的特殊缩放行为。通过理论推导，期望发现并纠正标准 μP 在此类架构上的失效点，最终验证推导出的公式能够在 AdamW 和 SGD 优化器下实现零样本学习率迁移，从而解决 GDN 在规模化训练时的参数化难题。
与已有工作不同的是，本文的独特切入点在于首次针对基于外积 Delta 规则更新的全矩阵递归状态进行 μP 分析，而非此前研究关注的对角标量递归。工作不仅处理了线性 RNN 的递归特性，还深入剖析了 αt 和 βt 非线性门控机制对坐标尺度传播的影响，这是此前 Vankadara et al. (2024) 针对 Mamba 的分析所未曾涉及的领域。此外，本文还提出了在 RMSNorm 前插入 √d 乘数的架构微调建议，以确保隐式重缩放不会破坏 μP 所需的梯度尺度，填补了混合线性架构特征学习理论的空白。
核心方法
方法基于 Tensor Programs 框架，通过在 GDN 的完整前向传播、门控机制和递归状态动力学中传播坐标尺度估计，推导满足 μP 特征学习条件的缩放规则。分析假设隐藏状态 xt 满足其 L2 范数为 Θ(√d) 且其更新的 L2 范数也为 Θ(√d)，即 xt 具有 Θ(1) 坐标大小且其更新具有相同量级。在此假设下，通过逐步计算查询、键、值、潜状态、输出以及门控参数的坐标大小，推导出各权重矩阵的初始化方差和前向乘数。接着，在反向传播中，通过分析梯度传播路径（特别是 BPTT 尾部累积项），推导出满足特征学习条件（即参数更新后预激活的变化为 Θ(1)）所需的学习率缩放。
核心创新在于发现了 GDN 中门控机制的特殊性导致其参数缩放规则偏离标准 μP。研究表明，为了维持特征学习，门控权重矩阵 Wα 和 Wβ 的 SGD 学习率需要非标准的 Θ(1/√nℓ-1) 缩放，而标量门控参数 alog 和 b 也需要 Θ(√nℓ-1) 缩放。这与标准 μP 中隐藏权重通常只需 Θ(1) 或 Θ(nℓ-1) 缩放有本质区别。此外，分析建议在 RMSNorm 之前插入 √d 乘数，以确保输入到 RMSNorm 的已经是 Θ(1) 坐标大小，避免隐式重缩放破坏 μP 所需的梯度缩放。这些发现揭示了 GDN 架构中存在的独特的缩放规律。
方法步骤详情
首先推导投影特征（查询、键、值）的坐标尺度，在短卷积和 L2 归一化后，qt 和 kt 具有 Θ(1/√d) 坐标大小。接着分析潜状态 St 的坐标尺度，证明若遗忘门 αt 和写入门 βt 严格位于 (0,1) 区间，则 St 的稳态方差与单次写入更新 Ut = βt vt kt^T 匹配，即 Θ(1/√d)。然后推导读取输出 ot = St qt，发现其坐标大小为 Θ(1/√d)，因此需在 RMSNorm 前乘以 √d 以满足 μP 条件。最后分析 SGD 下的反向传播梯度，假设有效记忆长度 Leff = O(1)，推导出投影权重 Wq, Wk, Wv 的学习率为 Θ(1)，而门控权重 Wα, Wβ 和标量参数的学习率则需 Θ(1/√nℓ-1) 以满足特征学习条件。
技术新颖性
本文的技术新颖性在于首次成功将 μP 理论扩展到了具有全矩阵递归状态和非线性门控的线性架构。与之前仅处理对角 SSM 的工作不同，本文通过传播坐标尺度估计，处理了外积更新（vt kt^T）和遗忘/写入门控（αt, βt）的交互。研究揭示了 GDN 架构中存在的特殊的缩放规律：门控权重的梯度尺度与其他隐藏权重不同，需要特定的学习率调整。此外，本文提出的在 RMSNorm 前插入 √d 乘数的架构微调建议，也是一种新颖的解决隐式重缩放问题的方案。这不仅验证了 μP 理论在复杂架构中的普适性，也为理解其他混合线性模型的缩放行为提供了新的理论工具。
实验结果
实验结果强有力地验证了本文推导的 μP 参数化公式的正确性和实用价值。在 AdamW 优化器下，测试了宽度 d ∈ {256, 512, 1024, 1536} 的模型（参数量从约 21M 到 342M），移位验证损失曲线显示最优学习率在所有宽度下高度一致（例如宽度 256 的最优移位损失为 3.4525，宽度 1536 为 2.7554，虽然绝对损失值不同但最优 LR 位置完全对齐），实现了零样本学习率迁移。相比之下，标准参数化 (SP) 的最优学习率随宽度显著偏移，无法迁移。在 SGD 优化器下（测试宽度 d ∈ {256, 512, 768, 1024}），标准参数化和原始 μP 配置均无法实现宽度间的学习率迁移（最优 LR 发生偏移），而本文的 μP 配置则成功实现了不同宽度下最优学习率的完美对齐。这些结果不仅在理论上验证了推导的正确性，也证明了其在 AdamW 和 SGD 两种优化器下的实用效率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Language Model Pre-training (FineWeb-Edu 100B) | Shifted Validation Loss | Optimal LR consistent across widths (AdamW ~6-7e-3, SGD ~0.2) | Standard Parametrization (SP): Optimal LR shifts significantly with width | Enables zero-shot hyperparameter transfer from small (d=256) to large (d=1536) models, saving massive tuning costs. |
局限与改进
作者承认分析中暂时忽略了 SiLU 激活的影响，假设其被抑制以保持高斯性。在预激活远离零时，这种近似质量会下降，但在 μP 初始化下仍然足够准确。此外，分析假设有效记忆长度 Leff = O(1)，即门控值 αs 衰减足够快。如果模型用于长上下文场景且 αt 接近 1，则 Leff 可增长为 O(T)，此时 μP 学习率处方应相应按 1/Leff 缩放。作者还提到，原始 SGD 优化器在实验中表现出极大的不稳定性，因此使用了带 Nesterov 动量的 SGD。我的观察是，论文主要关注 8 层模型的实验，对于更深层数或不同数据分布的泛化能力未做充分验证，且实验仅基于单个 NVIDIA H100 GPU，可能无法完全反映更大规模训练下的稳定性。
独立分析的弱点
分析的稳健性依赖于短有效记忆假设。在需要极长上下文依赖的任务中（如长文档摘要或代码生成），αt 可能接近 1，导致 BPTT 尾部累积项不再是直接梯度的 O(1) 倍，这将破坏当前的缩放法则。未来的工作需要针对 Leff = O(T) 的场景重新推导缩放规则。此外，对于门控参数的非线性变换（如 αt = e^(gt)），分析仅通过一阶导数近似处理其在初始化时的行为，对于训练后期这些参数偏离初始化较远时的动态缺乏理论保证。另外，实验中发现原始 SGD 优化器极不稳定，需依赖 Nesterov 动量，这可能限制了 μP 在某些特定优化场景下的直接应用。
未来方向
作者希望本文的推导能启发对其他线性或混合架构缩放定律的进一步研究。特别是基于 Delta Rule 的其他变体（如 Irie et al., 2021 的形式化）或结合了复杂结构化状态空间的混合模型（如 RWKV 或 HGRN）。此外，探索在长上下文场景（Leff 较大）下的修正缩放规则，以及将 μP 扩展到更多优化器（如 Lion、Adafetalon 或 Sophia）在 GDN 上的应用，都是有价值的方向。基于本文成果，还可以研究如何将这种参数化方法应用于其他具有复杂递归机制的神经架构（如某些图神经网络中的记忆更新机制或强化学习中的循环网络），以及如何将坐标尺度估计技术与自动化机器学习（AutoML）相结合，自动发现新架构的缩放规则。
复现评估
论文承诺代码可在 GitHub 获取 (https://github.com/lauyikfung/gated_delta_net_mup)，并提供了针对 Amazon Trainium 芯片的优化代码。实验使用了单个 NVIDIA H100 80GB GPU，在 FineWeb-Edu 100B 数据集上训练 20k 步，全局批大小为 480，序列长度为 1024（总计约 9.83B tokens）。算力需求（单个 H100）对于中等规模的实验是可接受的，但完整的参数扫描（每种宽度 5-7 个学习率）仍需较多计算资源（可能需要数天）。配置细节（如层深 8 层、注意力头数 6、初始化方案、优化器参数）在论文中有详细描述，复现难度中等。然而，由于实验涉及特定的学习率网格搜索和自定义的 μP 初始化实现，完全复现结果可能需要对代码库的深入理解。
论文图表
