GQA-{\mu}P: 群组查询注意力的最大参数化更新

# 分组查询注意力的最大更新参数化
来源：https://arxiv.org/html/2605.15290
Kyle R\. Chickering 加州大学戴维斯分校 & MBZUAI IFM & 王慧娟¹¹ MBZUAI IFM & 吴梦曦¹¹ 南加州大学 & MBZUAI IFM & Alexander Moreno MBZUAI IFM & Muhao Chen 加州大学戴维斯分校 & Xuezhe Ma 南加州大学 & MBZUAI IFM & Daria Soboleva Cerebras & Joel Hestness Cerebras & 刘正宗 MBZUAI IFM & 邢波 卡内基梅隆大学 & MBZUAI IFM

###### 摘要

跨模型架构的超参数迁移显著减少了调整大型语言模型（LLM）所需的计算量。最大更新参数化（μP）通过严谨的数学分析确保了迁移性，但针对新模型架构推导该参数化方法可能具有挑战性。基于 Yang 等人（2023a）的谱特征学习观点，我们取得了两个进展。首先，我们将权重的谱范数条件从启发式方法提升为特征学习的定义，并因此无需借助惰性学习即可推导出 Complete-P 深度和权重衰减缩放规则。其次，我们考虑了一种修正的谱范数，当权重矩阵不是满秩时，它能保留网络权重的有效缩放定律。这使得（据我们所知）首次推导出分组查询注意力（GQA）的 μP 缩放规则。我们展示了学习率在 GQA 重复超参数上的迁移，以及关于权重衰减迁移的实验，从而验证了我们理论推导的有效性。

## 1 引言

最大更新参数化（μP）（Yang and Hu, 2021; Yang et al., 2022）提供了跨模型宽度进行零样本学习率迁移的原则性规则。因此，可以通过扫描一个小的代理模型来确定大型终端模型的超参数。μP 已被用于训练至少 130 亿参数的模型并实现零样本迁移（Blake et al., 2023; Dey et al., 2023; Narayan et al., 2025）。然而，其适用性在很大程度上局限于跨模型宽度的学习率迁移。为了拓宽这一范围，Dey et al.（2025）引入了 Complete-P，将原始规则扩展到权重衰减和模型深度。然而，许多在生产中广泛部署的常见架构仍然缺乏既定的 μP 缩放。

本文旨在通过扩展 Yang 等人（2023a）的谱 μP 框架，使其在推导新型架构的 μP 规则时更具实用性，从而弥补这一差距。作为我们框架实用性的一个例子，我们推导出了（据我们所知）分组查询注意力（GQA）（Ainslie et al., 2023）的首个 μP 缩放。我们的分析揭示了 GQA 带来了先前工作未解决的几个难题。首先，当使用 GQA 时，原始的 μP 实现通过了坐标检查，即该实现惯用的正确性测试。然而，实证分析表明，原始的 μP 实现未能迁移学习率，这似乎与既定理论相矛盾（见图 1 和图 5）。我们通过扩展 Yang 等人（2023a）中引入的谱范数版本的 μP，并证明原始的 μP 实现未能通过更严格的谱范数坐标检查，从而解决了这个问题。其次，GQA 权重矩阵固有的低秩特性扭曲了层输出的预期大小。为了解决这个问题，我们引入了一个新的范数，即期望算子范数，来替代谱 μP 理论中的谱范数，并恢复期望的缩放行为。

我们的主要贡献有三方面：

1. 我们扩展了 Yang 等人（2023a）的谱 μP 理论，使其能够推导出更高级架构（如权重衰减、递归块和 GQA）的 μP。我们的工作提供了（据我们所知）GQA 的 μP 缩放的首次推导。
2. 我们进行了实证分析以验证该理论，并为跨 GQA 设置的学习率迁移提供了实践指导。特别是，我们指出跨不同数量的 GQA 重复进行迁移会导致噪声较大的迁移动态，这表明在尝试迁移学习率时应谨慎。
3. 我们的实验表明，在正确的缩放规则下，权重衰减以及 Wang 和 Aitchison（2024）引入的训练时常数 τ_epoch 似乎都是可迁移的。

参照图注 图 1：标准参数化（左）、朴素 Adam-μP 参数化（中）和我们的 GQA-μP 缩放（右）的比较。对于固定模型大小，我们改变 KV 头的数量。虚线表示每种参数化的平均最优学习率，灰色阴影区域表示最优学习率的标准差。所有模型都训练到每个参数 10 个 token（TPP）。更多细节见附录 B.1.2。

## 2 相关工作

μP 基础：μP 建立在 Yang 的一系列工作之上，发展了张量程序框架（Yang, 2019; 2020a; 2020b; Yang and Hu, 2021; Yang et al., 2022; 2023b）。这一系列工作使用随机矩阵理论仔细分析神经网络在训练过程中的数学性质，同时通过实验证明这些理论方法在现实世界的深度学习中仍然具有价值。在张量程序框架内，Yang 等人（2022）推导了在 SGD 和 Adam 训练下众所周知的宽度 μP 缩放定律。该系列的最后一项工作 Yang 等人（2023b）试图将 μP 扩展到深度缩放。然而，他们无法将其扩展到具有标准隐藏层配置的残差块情况。最后，本工作所呈现的数学框架基于 Yang 等人（2023a），他们展示了一种基于谱范数的 Yang 等人（2022）结果的不同推导方法。

使用 GQA 的模型：分组查询注意力（GQA）（Ainslie et al., 2023）是一种高效的注意力机制，通过在一组查询头之间共享键和值头来减少内存使用。由于其内存效率和模型性能之间的良好平衡，GQA 已被广泛应用于现代大型语言模型中，包括 Mistral 7B（Jiang et al., 2023）、LLaMA 3（Grattafiori et al., 2024）、Qwen3（Yang et al., 2025）和 K2-V2（Liu et al., 2025）。

μP 的扩展：Yang 等人（2022）提出的原始 μP 公式仅适用于对固定深度、固定批量大小的神经网络进行宽度缩放。虽然这已经是一个强大的工具，但后续的作者试图扩展 μP 的原理以涵盖原始公式未涉及的情况。Dey 等人（2023）使用 μP 进行了大规模验证实验，并发现学习率可以跨批量和数据集大小迁移的经验证据。Dey 等人（2023）提出了权重衰减、Adam ε 和深度的 μP 类型缩放。他们对深度缩放的贡献最为显著，因为他们的实证发现与 Yang 等人（2023b）提出的缩放规则相矛盾。然而，他们广泛的实证分析表明他们推导的缩放规则是正确的。我们在本文第 3.3 节中使用所概述的框架得出了相同的缩放规则。最近，Mlodozeniec 等人（2025）通过将 SDE 参数化扩展到涵盖批量大小的超参数迁移，以及展示逐层学习率调优的价值，扩展了 Dey 等人（Dey et al., 2025）的工作。

Blake 等人（2023）在低精度 LLM 大规模训练的背景下应用了 μP。他们使用 ABC 参数化应用 μP 缩放规则，同时保持网络中所有层的单位方差，他们称之为单位缩放-μP。此外，他们通过实验验证了在受控条件下，学习率迁移在跨数据集、批量大小、深度和训练迭代次数时仍然存在。Narayan 等人（2025）提出了一个不同且更简化的单位缩放-μP 版本，他们也证明该版本适用于使用 μP 训练低精度网络。

最后，一项相关且后续的工作 Zheng 等人（2026）提出了一个与我们的相似的理论框架。我们的工作和他们的工作都持有这样的观点：谱范数为推导 μP 提供了比张量程序（Yang and Hu, 2021）更具原则性的替代方案。然而，我们的工作在范围和动机上有所不同。Zheng 等人（2026）系统地将他们的框架应用于在宽度和深度缩放下的一类广泛的优化器，而我们的工作则确定了期望算子范数是正确处理秩退化权重所必需的，并利用这个范数首次推导了 GQA 的 μP。

## 3 推导新颖的最大更新参数化

考虑神经网络中按层 ℓ 索引的一组权重矩阵 W^ℓ ∈ R^{n_ℓ × m_ℓ}。Yang 等人（2023a）证明了对网络权重矩阵施加的条件意味着特征学习（从而学习率迁移），正如 Yang 等人（2022）所定义的那样（见方程 (3)）。对于初始权重 W_0^ℓ 和迭代权重 W_t^ℓ = W_0^ℓ + ∑_{k=1}^t ΔW_t^ℓ，其中 ΔW_t^ℓ = W_t^ℓ - W_{t-1}^ℓ，Yang 等人（2023a）指出初始化和更新都必须满足：

‖W_0^ℓ‖ = Θ(√n_ℓ / √m_ℓ),   ‖ΔW_t^ℓ‖ = Θ(√n_ℓ / √m_ℓ),      (1)

其中 ‖W‖ := sup_{‖x‖_2=1} ‖Wx‖_2 是通常的谱（或诱导）范数。这种关于特征学习的谱视角是强大的，我们引入了三个微小但重要的修改，使得我们能够将 Yang 等人（2023a）的方法扩展到涵盖像 GQA 这样的新型架构。

参照图注 图 2：当输入是从 N(0,1) 独立同分布采样时，谱范数未能准确捕捉低秩矩阵行为的演示。r 是键值头重复次数，r=1 对应于没有 GQA 的设置。每个点是对 1000 次独立生成的 A 矩阵取平均，阴影带显示 ±1 个标准差。

**新范数下的分析：** 谱范数可以解释为由算子 φ: R^m → R^n 引起的输入向量的最大形变。对于满秩算子，例如密集的前馈层，随机矩阵理论表明谱范数的量化值是渐近可达的。在经典的 n×n 随机矩阵 A 的情况下，当 n→∞ 时，我们有尖锐的渐近关系 ‖A‖ = 2√n。

然而，对于像 GQA 中使用的那种退化秩矩阵，谱范数在实践中并非渐近可达。原因在于，正如张量程序（Yang and Hu, 2021）所示，训练期间 GQA 层的输入是独立同分布的，因此对于退化秩矩阵，导致这种“最大形变”的向量出现的概率为零！这种差异的可视化见图 2。相反，我们应该使用一种能反映训练过程中实际遇到的形变的大小概念。

为此，令 Ω 为输入向量的概率分布。我们定义*期望算子范数*如下¹¹：从技术上讲，我们定义的 ‖A‖_{E,Ω,p} 在不对 Ω 施加进一步约束的情况下只是一个半范数。特别地，如果 supp Ω ≠ R^n，那么所有随机向量 x~Ω 都可能位于 A 的零空间中。这种边缘情况在神经网络训练中不会发生。

‖A‖_{E,Ω,p} := E_{x~Ω} [ ‖Ax‖_p / ‖x‖_p ]。      (2)

在本文中，我们采用约定 ‖A‖_E = ‖A‖_{E, N(0,1), 2}，其中 x~N(0,1) 具有独立同分布的条目。关键的是，当 A 是方阵且具有独立同分布条目时，它以概率 1 满秩，我们得到渐近关系 ‖A‖_E = Θ(‖A‖)。引理 2 中提供了一个证明。

**算子范数聚焦的特征学习：** Yang 等人（2023a）表明约束权重矩阵的谱范数意味着 Yang 和 Hu（2021）意义上的特征学习，其中特征学习被定义为当