大规模门控Delta网络中的特征学习解锁

TL;DR 首次为 Gated Delta Network 推导 μP 缩放规则，实现学习率跨宽度零样本迁移，并在语言模型预训练上验证了其在 AdamW 与 SGD 下的稳定性。

问题

大型语言模型（LLM） 的训练和扩展需要巨额算力,这催生了两大方向:一是设计 次二次方复杂度的替代架构 (如状态空间模型、线性注意力),二是发展 原理性超参数调优方法 (如最大更新参数化 μP),以实现小模型调参后直接迁移至大模型的零样本超参数迁移,从而大幅降低实验成本。

目前,μP 已成功应用于标准 Transformer,但其向 线性模型 的扩展——尤其是那些具备 结构化状态转移 和 复杂门控机制 的架构——仍几乎空白。Gated Delta Networks (GDN) 一类的模型通过门控循环状态动态实现了高效的长序列建模,但其内部各组件(查询、键、值投影,门控标量,循环状态)的 坐标大小 (coordinate size) 在深度和宽度变化下的传播规律尚未被刻画。若直接沿用标准参数化 (SP),不同宽度模型的最优学习率随宽度剧烈漂移,导致小模型实验结论无法可靠外推,限制了探索这类高效架构的规模上限的迭代速度。

这一问题的难点在于,GDN 前向传播同时涉及 非线性门控、随时间展开的隐藏状态累加 以及 多分支信息混合,使得逐层分析坐标大小并推导出保持稳定特征学习的初始化与学习率缩放规则高度非平凡。业界对 线性复杂度的下一代基础模型 寄予厚望,若能建立其 μP 理论,意味着研究者可以在低成本的小模型上搜索超参数,然后直接用于大规模训练,极大利好资源受限团队的创新效率,并加速新架构的工业化验证。

类比而言,这就像当初 Transformer 的 μP 让算力一般的研究者也能在玩具尺寸上找到最佳超参,然后直接训练出 GPT-3 级别的模型;而 GDN 的 μP 则为 状态空间模型与线性注意力 这类有望替代 Transformer 的架构补上了同样的 “缩放秘钥”,使其从实验性原型走向可规模化部署的关键一步。

核心洞察

• 这项工作首次将 μP 理论系统性地扩展到带有门控机制和循环状态的高效线性模型，填补了亚二次架构超参数迁移的理论空白。不同于以往 μP 仅验证于标准 Transformer 或简单线性层，本文通过对 Gated DeltaNet 的坐标大小在前向传播、门控及状态更新中进行严格逐项估计，推导出一套完整的宽度缩放规则。这使得在语言模型预训练中，学习率等关键超参数能够零次迁移至不同宽度的模型，避免了逐宽度手动调参的巨大算力消耗。它为 Mamba、RWKV 等更复杂的线性 RNN 架构提供了可复制的方法论，对工程中快速扩展高效模型具有重要指导意义。
• 门控标量 \(\alpha_t\), \(\beta_t\) 的初始化与缩放规则被首次详细分析并验证，揭示了它们在稳定特征学习中扮演的关键角色。在标准参数化下，这些门控因不恰当的初始化会导致梯度消失或更新不稳定，从而破坏超参数迁移。本文通过推导出匹配宽度变化的门控缩放因子，使得无论采用 SGD 还是 AdamW 优化器，均能在宽度放大时保持一致的训练动态。这一发现突破了以往 μP 对门控机制的简化或忽略，为实际训练中引入门控的线性模型提供了可直接应用的缩放公式，显著降低了从原型到大规模模型迁移的风险。

方法

该方法将 Maximal Update Parametrization (μP) 理论系统扩展到 Gated Delta Network，一种使用门控线性递归单元替代自注意力的高效次二次方架构。输入为序列 token 嵌入，输出为下一 token 预测 logits。关键推导流程如下：

前向坐标大小传播

从输入嵌入的方差出发，逐层估计 投影特征（查询、键、值）、潜在状态 及 读出层 在初始化时的坐标级大小（coordinate size）。特别对 门控标量 (α, β) 做一阶分析，捕获其受输入宽度影响的方式。该传播严格穿过门控激活、残差连接和状态更新方程，确保每一中间量的数量级与宽度关系可追踪。

反向梯度缩放

基于前向坐标大小，在 SGD 设定下推导各参数梯度的坐标级缩放。包括 读出误差 反向传播至状态和投影，以及 门控权重（Wα, Wβ）与标量参数（a_log, b）的梯度。核心在于关联梯度大小与参数坐标大小，使更新量在宽度改变时保持特征学习（feature learning）稳定。

缩放规则导出

综合前向与反向分析，为每个参数指定 初始化方差 与 学习率 的宽度缩放指数。例如，主投影权重 (Wq, Wk, Wv, Wo) 采用与 Transformer μP 一致的初始化与学习率缩放；门控参数则需根据其特定动态设置不同指数。推导覆盖 SGD，并分析确认这些规则在 AdamW 的尺度不变区（scale-invariant regime）中也适用，确保自适应优化器下同样实现学习率迁移。

输出验证

最终缩放规则被用于不同宽度（从数百万到十亿参数量级）的语言模型预训练实验，验证学习率直接迁移后的训练损失高度一致，而标准参数化（SP）则无法迁移。

与同类工作的差异：首次将 μP 理论应用于具有 结构化状态转移和显式门控机制 的线性递归模型，解决了前辈工作仅针对标准 Transformer 或简单线性层的局限，为高效次二次方架构的超参数调优提供了原则性方案。

实验

实验设计

论文以 Gated Delta Network 为对象，系统验证了所推导的 μP 缩放规则 在语言模型预训练中的有效性。实验对比了两种参数化方式：

• 标准参数化 (Standard Parametrization)：模型宽度变化时，学习率等超参数无法直接迁移。
• μP 参数化 (本文方案)：根据前向坐标尺寸估计、门控机制与循环状态动力学推导出的缩放规则。

训练采用 AdamW 和 SGD 两种优化器，在不同宽度下进行语言模型预训练，观察学习率迁移的稳定性。

关键发现

• 在 标准参数化 下，改变模型宽度后，原本最优的学习率在不同宽度时性能急剧下降，无法实现零样本超参数迁移。
• 采用本文的 μP 配置，学习率在多个宽度上展现出高度一致性，损失曲线几乎重合，表明特征学习行为在宽度变化时保持稳定。
• 该稳定性在 AdamW 和 SGD 下均成立，验证了理论分析的普适性，突破了先前 μP 主要围绕标准 Transformer 的局限。

与基线的对比解读

标准参数化的失效源于其无法维持前向/反向传播中激活与梯度的坐标尺度一致性；而本文通过严格传播 坐标尺寸估计 (coordinate-size estimates)，精确设定了投影矩阵、门控权重与标量参数的初始化方差和学习率缩放规则，使得每一层的更新幅度独立于宽度。这是首个针对 带有结构化状态转移和复杂门控的线性架构 成功实现 μP 的工作，为后续探索 sub-quadratic 架构 的超参数迁移奠定了理论基础，工程上可直接降低不同规模模型的调参成本。

行业影响

落地场景

Gated Delta Networks (GDN) 作为 sub-quadratic 线性注意力架构，其 μP 缩放规则 的建立直接解决了此类模型从实验规模向工业级参数扩展时的超参数漂移问题。核心落地场景包括：

• 大语言模型预训练：任何需要训练百亿到千亿参数 LLM 的团队，可基于小规模代理模型（如 100M）进行超参数搜索，而后通过缩放法则零-shot 迁移至目标宽度，省去大规模网格搜索。
• 长上下文推理与文档处理：GDN 的隐式状态记忆机制适合处理超长文本（如法律合同分析、论文摘要、代码库理解），μP 确保在不同批次大小与序列长度下训练稳定。
• 边缘/端侧设备上的流式语言模型：得益于线性复杂度，GDN 可在资源受限设备上做增量推理；通过 μP 缩放规则，可直接在小尺寸代理上调试学习率，再部署到不同芯片的定制化宽度型号。

商业价值

• 降本：将超参数搜索成本从 O(n^2) 降低为 O(1)。传统大模型训练前通常需要多次不同规模实验来寻找合适学习率，μP 允许仅在一次小模型上完成调参后直接迁移，节省数百万美元的计算预算（尤其对千亿参数以上模型）。
• 加速产品迭代：模型架构更新或引入新门控机制时，工程师无需重新扫描超参数空间，可直接基于历史缩放规则快速启动训练，将实验周期从数周缩短到天级别。
• 提升多平台兼容性：同一个训练配方可直接应用于不同宽度模型（适用于云、边缘、移动端），降低维护成本。

与现有工作流的接口

• 训练框架：该工作可直接集成到 PyTorch / JAX 训练代码中，仅需在模型定义时应用 μP 初始化与学习率缩放规则。作者已提供 GitHub 实现（gated_delta_net_mup），可作为 mup 库的扩展或独立 module。
• 与现有 Transformer 共存：GDN 可作为 Transformer 的替代组件用于编解码器，或混合使用（如编码器用 GDN 处理长输入，解码器用标准注意力）。缩放规则的推导方式（通过坐标尺寸估计）也可复用到其他类似状态空间模型和门控线性单元，方便工业界统一超参数策略。
• 与自动化 ML 管道整合：该缩放规则可嵌入到自动化调度系统（如 Kubernetes 上的训练 Job），根据目标宽度自动调整学习率和初始化，无需人工干预。

具体用案

1. 内容平台的推荐与摘要系统：视频/新闻平台需对海量用户评论和文章做实时分析，使用 GDN 做特征提取和生成摘要。随着平台扩容，模型宽度需从实验室的 128 维扩展到 1024 维，μP 规则保证直接复用早期小模型找到的最优学习率，避免因超参数不当导致训练崩溃或资源浪费。
2. 金融领域的高频因子挖掘：量化交易公司使用线性注意力网络处理多模态时序数据（如订单簿 + 新闻流）。GDN 的递推状态天然适合流式处理，μP 缩放能力使得在不同交易品种或时间粒度上调整模型宽度时，无需重新进行昂贵的超参数搜索，加速模型迭代。

局限

• **局限于特定架构**：本文仅为 **Gated Delta Network** 推导了 μP 缩放规则，该方法论虽然可能推广至其他线性注意力或状态空间模型（如 Mamba、RetNet），但论文未给出直接证据或通用框架。实际应用中，若需部署到不同架构，仍需重新进行繁琐的坐标大小分析，限制了零成本迁移的便利性。
• **实验规模与任务覆盖有限**：语言模型预训练实验仅在小宽度（最大 1536）和浅层网络（5 层）上进行，未验证在数十亿参数工业级模型上的效果。另外，论文未评估下游任务微调或适应性场景，缺乏对最终应用性能的闭环验证，因此在大规模生产环境中的收益仍需进一步观察。
• **优化器与二阶效应未深入**：分析围绕 SGD 和 AdamW 的一阶前向 / 反向缩放，但未考虑学习率 warmup、衰减调度、二阶矩累积稳定性或混合精度训练下的交互。门控参数的动态分析也依赖强假设（如稳定状态），实际训练中门控饱和或梯度消失 / 爆炸可能导致迁移失败，而论文未提供鲁棒性保障。