GQA-μP: The maximal parameterization update for grouped query attention
作者: Kyle R. Chickering, Huijuan Wang, Mengxi Wu, Alexander Moreno, Muhao Chen, Xuezhe Ma, Daria Soboleva, Joel Hestness, Zhengzhong Liu, Eric Xing
分类: cs.LG, cs.AI
发布日期: 2026-05-14
备注: 18 pages
💡 一句话要点
提出GQA-μP:分组查询注意力机制的最大参数化更新方法,实现超参数迁移。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 分组查询注意力 最大参数化更新 超参数迁移 谱特征学习 大型语言模型
📋 核心要点
- 大型语言模型超参数调优需要大量计算资源,而跨模型架构的超参数迁移可以显著减少计算需求。
- 论文基于谱特征学习视角,通过改进谱范数条件,推导出分组查询注意力(GQA)的μP缩放比例。
- 实验证明,该方法能够有效实现学习率和权重衰减等超参数在不同GQA配置间的迁移。
📝 摘要(中文)
本文提出了一种针对分组查询注意力(GQA)的最大参数化更新(μP)方法,旨在解决大型语言模型(LLM)超参数调优所需的巨大计算量问题。该方法基于Yang等(2023a)提出的谱特征学习视角,并进行了两项改进。首先,将权重上的谱范数条件从启发式方法提升到特征学习的定义,从而在不依赖惰性学习的情况下,推导出完整的深度和权重衰减缩放比例。其次,考虑了一种改进的谱范数,当权重矩阵不是满秩时,该范数可以保持网络权重的有效缩放规律。据我们所知,这是首次推导出GQA的μP缩放比例。实验结果表明,我们的理论推导是有效的,能够在GQA重复超参数上实现学习率迁移,并在权重衰减迁移方面取得良好效果。
🔬 方法详解
问题定义:现有的大型语言模型(LLM)超参数调优需要大量的计算资源。虽然超参数迁移是一种有效的解决方案,但如何保证超参数在不同模型架构间的有效迁移是一个挑战。特别是对于分组查询注意力(GQA)这种新型注意力机制,缺乏有效的理论指导其超参数缩放。
核心思路:论文的核心思路是基于最大更新参数化(μP)理论,通过对GQA进行参数化,使得模型在不同规模下具有相似的学习动态,从而实现超参数的有效迁移。关键在于推导出GQA的μP缩放比例,这需要对GQA的权重矩阵进行谱分析,并找到合适的谱范数条件。
技术框架:论文的技术框架主要包括以下几个步骤:1) 基于谱特征学习视角,重新审视权重的谱范数条件,将其提升到特征学习的定义层面。2) 针对GQA的特点,设计了一种改进的谱范数,该范数能够处理权重矩阵非满秩的情况。3) 基于改进的谱范数,推导出GQA的μP缩放比例,包括深度和权重衰减的缩放比例。4) 通过实验验证推导出的缩放比例的有效性,包括学习率和权重衰减的迁移实验。
关键创新:论文的关键创新在于:1) 首次推导出了分组查询注意力(GQA)的μP缩放比例,为GQA的超参数调优提供了理论指导。2) 提出了一种改进的谱范数,能够处理权重矩阵非满秩的情况,扩展了μP理论的应用范围。3) 将谱范数条件从启发式方法提升到特征学习的定义层面,使得理论推导更加严谨。
关键设计:论文的关键设计包括:1) 改进的谱范数:针对GQA中query和key的维度不一致的情况,设计了一种能够保持权重缩放规律的谱范数。2) μP缩放比例:基于改进的谱范数,推导出了GQA的深度和权重衰减的缩放比例,这些比例是实现超参数迁移的关键。3) 实验设计:设计了学习率和权重衰减的迁移实验,验证了推导出的缩放比例的有效性。
🖼️ 关键图片
📊 实验亮点
实验结果表明,基于论文提出的μP缩放比例,学习率可以在不同的GQA重复次数之间有效迁移。此外,权重衰减的迁移实验也取得了良好的效果,验证了理论推导的有效性。这些结果表明,该方法能够显著减少GQA模型超参数调优所需的计算资源。
🎯 应用场景
该研究成果可应用于大型语言模型的快速超参数调优,尤其是在资源受限的情况下。通过μP缩放,可以在小规模模型上进行超参数搜索,然后将最优超参数迁移到大规模模型上,从而显著减少计算成本。此外,该方法还可以促进新型注意力机制的开发和应用。
📄 摘要(原文)
Hyperparameter transfer across model architectures dramatically reduces the amount of compute necessary for tuning large language models (LLMs). The maximal update parameterization (μP) ensures transfer through principled mathematical analysis but can be challenging to derive for new model architectures. Building on the spectral feature-learning view of Yang et al. (2023a), we make two advances. First, we promote spectral norm conditions on the weights from a heuristic to the definition of feature learning, and as a consequence arrive at the Complete-P depth and weight-decay scalings without recourse to lazy-learning. Second, we consider a modified spectral norm that preserves the valid scaling law of network weights when weight matrices are not full rank. This enables (to our knowledge, the first) derivation of μP scalings for grouped-query attention (GQA). We demonstrate the efficacy of our theoretical derivations by showing learning rate transfer across the GQA repetition hyperparameter as well as experiments regarding transfer over weight decay.