Unlocking Feature Learning in Gated Delta Networks at Scale

技术摘要：解锁大规模门控 Delta 网络中的特征学习

问题陈述
大语言模型（LLMs）的扩展面临两个主要瓶颈：训练更大模型的高昂计算成本，以及标准 Transformer 自注意力机制的二次方复杂度。线性循环架构，如门控 Delta 网络（Gated Delta Networks, GDN），已成为高效的替代方案，能够在保持线性时间推理的同时捕捉长程依赖关系。然而，如何有效地扩展这些模型仍存在关键空白。虽然最大更新参数化（Maximal Update Parametrization, µP）已成功实现了标准 Transformer 及某些线性模型在零样本情况下的超参数迁移（特别是学习率），但其在门控 Delta 网络上的应用仍处于空白阶段。核心挑战在于该架构的循环状态动力学——其状态通过依赖数据的门控机制和全矩阵 Delta 规则进行更新，而非现有 µP 推导中所假设的标量循环或标准前馈结构。先前关于结构化状态空间模型（SSMs，如 Mamba）的研究表明，原生 µP 在对角 SSM 上会失效，但 GDN 使用的全矩阵状态更新使得那些特定的修正方法不再适用。

方法论
作者通过严谨地推导出门控 Delta 网络的 µP 公式来解决这一差距。其方法涉及将坐标大小估计（coordinate-size estimates）通过整个前向传播过程进行传递，包括独特的门控机制和循环状态转换，以确定初始化、前向乘数和学习率所需的缩放规则。

关键分析步骤包括：

1. 前向传播分析： 作者推导了投影特征、潜在状态 $S_t$ 以及读出层（readout）的坐标大小。他们发现，查询（queries）与键（keys）的 L2 归一化结合秩一（rank-one）写入更新，会导致特定的方差缩放。至关重要的是，他们提出在 RMSNorm 层之前插入一个 $\sqrt{d}$ 乘数，以确保归一化层的输入保持 $\Theta(1)$ 的坐标大小，从而防止破坏特征学习所需的梯度缩放。
2. 门控机制分析： 论文分析了非线性门控参数 $\alpha_t$（衰减）和 $\beta_t$（写入强度）。不同于标准权重，这些参数是由涉及可训练权重向量（$W_\alpha, W_\beta$）和标量参数（$a_{log}, b$）的标量预激活值导出的。分析表明，由于这些门控组件在状态更新方程中的特定作用，它们并不遵循标准的 µP 缩放法则。
3. 反向传播与学习率推导：
   • AdamW： 由于 AdamW 通过其二阶矩对梯度进行归一化，有效更新幅度在不同权重类别间保持一致，从而简化了学习率要求。
   • SGD： 在纯 SGD 下，原始梯度幅度在不同权重类别之间存在显著差异。作者推导出，门控权重矩阵（$W_\alpha, W_\beta$）需要 $\Theta(1/\sqrt{d})$ 的学习率缩放，而标量门控参数（$a_{log}, b$）则需要 $\Theta(\sqrt{d})$ 的缩放。这些缩放规则偏离了标准 µP 设置，在标准设置中，学习率通常是统一的或遵循简单的 $1/n$ 规则。

核心贡献

1. 理论推导： 本文提供了首个完整的门控 Delta 网络 µP 公式。通过对全前向传播过程进行坐标大小估计，为所有权重类别建立了原则性的初始化方差和前向乘数。
2. 新颖的缩放规则： 作者发现门控组件在 SGD 下需要非标准的学习率缩放：
   • 门控权重矩阵（$W_\alpha, W_\beta$）需要 $\Theta(1/\sqrt{d})$ 缩放。
   • 标量门控参数（$a_{log}, b$）需要 $\Theta(\sqrt{d})$ 缩放。
3. 实证验证： 作者使用 AdamW 和 SGD 优化器，在多个宽度（隐藏维度从 256 到 1563 变化）下预训练了 GDN 语言模型。
   • 结果： 在所提出的 µP 配置下，最优学习率在两种优化器下均能在所有模型宽度间实现零样本迁移。相比之下，标准参数化（Standard Parametrization, SP）和原始 µP 配置均无法实现迁移，其最优学习率随模型宽度增加而显著偏移。
   • 稳定性： 所提出的配置实现了稳定的学习率迁移，证实了理论推导的正确性。

意义与主张
本文声称其分析解决了门控 Delta 网络在大规模特征学习中的参数化开放问题。通过证明所推导的 µP 公式允许实现零样本超参数迁移，该工作验证了理论推导的正确性及其在实践中的效用。作者强调，这是首个针对门控 Delta 网络推导并验证 µP 一致性参数化的工作，这使其区别于以往关于对角 SSM 的研究。研究结果表明，适当的参数化对于扩展线性循环模型至关重要，且 GDN 中特定的门控机制需要独特的缩放规则，这些规则既不同于标准 Transformer，也不同于其他线性架构。作者希望这些推导能为未来研究其他线性或混合架构的缩放法则提供指导。