摘要：近日，由浙江大学、西湖大学、香港科技大学和北京大学的研究团队联合发表了一篇创新性论文《Taming LLMs by Scaling Learning Rates with Gradient Grouping》。这项研究由西湖大学的刘子成教授担任通讯作者，论文第
一、大型语言模型的训练难题与传统解决方案
想象一下，你是一位教育者，面对一个有着数千名学生的超大班级。每个学生（就像模型中的参数）都有着不同的学习速度和能力。如何高效地教授这样一个庞大且差异巨大的班级？这就是大语言模型（LLM）训练面临的核心挑战。
目前的大语言模型往往包含数十亿甚至数千亿个参数，这些参数分布在不同的网络层次和结构中，就像那个超大班级的学生来自不同的年级和专业。虽然Adam和AdamW等自适应优化器能够为每个参数单独调整学习速度（就像为每个学生定制个性化教学计划），但这种做法面临着两个主要问题：一是消耗大量内存资源（想象要为每个学生准备完全不同的教材需要多少纸张），二是难以准确估计每个参数最佳的学习速率（难以精确评估每个学生的最佳学习节奏）。
行业内已有多种应对方案。一种是参数高效微调（PEFT）方法，如LoRA，它通过只训练少量参数来减少计算负担，就像只选择班级里的少数代表进行重点培训。另一种是优化器状态压缩方法，如量化或梯度统计近似，就像将详细的学生档案简化为概要信息。然而，这些方法往往会带来性能下降或在不同任务中表现不一致的问题。
最近的研究发现，大语言模型中不同层的参数（比如注意力层和MLP层）展现出不同但内部一致的优化行为模式。基于这一发现，Adam-mini等方法将模型参数分为预定义的组，每组使用一个平均学习率，而不是为每个参数单独设置。这种方法虽然减少了计算负担，但也可能损失优化精度。
二、SGG方法：集体与个性的平衡艺术
西湖大学团队提出的解决方案名为"梯度分组缩放"（Scaling with Gradient Grouping，简称SGG），它采取了一种巧妙的平衡策略：既保留了参数级别的个性化优化，又引入了组级别的集体约束。
想象一下，我们不是简单地将班级分成固定的小组，而是根据学生们的实时学习状态动态调整分组，然后为每个组制定不同的教学策略，同时仍然关注每个学生的个体差异。这就是SGG的核心思想。
具体而言，SGG的工作流程包括两个关键步骤：
首先，SGG会在每一层网络中动态地将参数（确切地说是梯度动量向量）聚类成若干组。这种聚类是基于参数优化行为的相似性，而不是预先定义的固定分组。就像一位敏锐的教师根据学生的实时表现动态调整学习小组，而不是按照固定的座位表或学号分组。
其次，SGG为每个聚类组计算特定的缩放因子，用于调整该组内所有参数的学习率。这种缩放是基于该组与该层平均水平以及整个模型全局统计数据的偏差程度来确定的。可以想象为：如果一组学生的学习速度明显快于班级平均水平，教师会适当放慢他们的学习节奏，让他们能更深入地消化知识；反之，对于学习较慢的小组，则会适当加快节奏，帮助他们赶上进度。
重要的是，SGG并不是替换现有的自适应学习率，而是对其进行缩放调整。这就像教师在保留个性化教学计划的同时，增加了班级层面和学校层面的整体协调机制，既尊重个体差异，又确保整体进度的一致性。
通过实验，研究团队发现，使用动量向量而非原始梯度进行聚类，以及采用层内参数与层平均值的偏差中位数（Median of Deviation to Average，简称MDA）作为缩放依据，能够取得最佳效果。
三、实验验证：让数据说话
SGG的效果如何？研究团队在多种大语言模型任务上进行了全面验证，包括C4数据集上的预训练、GLUE基准测试上的监督微调、常识推理任务上的参数高效微调，以及基于人类偏好的直接优化（DPO）等任务。
在C4预训练任务中，SGG与Adam结合使用时，在各种模型规模（从6000万到10亿参数）上都取得了显著的性能提升。例如，在60M参数的模型上，Adam+SGG比原始Adam降低了3.75%的验证困惑度；在10亿参数的模型上，降低了1.26%。更令人惊讶的是，SGG与LoRA结合使用时，使得低秩预训练的性能首次达到了全秩训练的水平，在130M参数模型上验证困惑度降低了惊人的10.30%。
在GLUE基准测试中，AdamW+SGG在全秩训练中平均提升了1.00%的准确率，在低秩训练（LoRA）中提升了更多，达到1.27%。特别是在MNLI和RTE等任务上，性能提升尤为显著。
在常识推理任务上，SGG使LoRA的平均准确率提高了2.9%，在特定任务如OBQA上甚至提高了4.2%。这一性能水平超过了许多现有的PEFT方法，如Prefix、Series和Parallel，甚至接近或超过了最新的DoRA、GaLore和Fira方法。
在DPO任务中，AdamW+SGG在LoRA训练中将Qwen2.5 0.5B模型的准确率提高了1.80%，达到72.02%，不仅超过了基线方法，甚至超过了其全秩训练的对应版本（71.85%）。
在多模态大语言模型（MLLM）验证中，SGG同样表现出色。在LLaVA-v1.5上，AdamW+SGG在视觉问答基准测试上平均提高了1.0%的准确率，在VizWiz上甚至提高了3.0%。当与Adafactor结合时，SGG也能带来0.6%的平均收益。在LoRA微调中，SGG提供了平均1.0%的性能提升；在8位量化LoRA（Q-LoRA）中，也带来了0.6%的提升。
四、SGG的强大稳定性与实用性
除了性能提升外，SGG还展现出了令人印象深刻的稳定性。在不同批次大小（从128到4096）和学习率（从1e-5到1e-1）的实验中，SGG始终保持稳定的验证损失，即使在极端条件下（如批次大小4096和学习率0.1）也能正常工作。这表明SGG能有效地缓解梯度异常值，动态调整学习率，确保在各种配置下的训练稳定性。
在实际应用方面，SGG的设计非常灵活，可以无缝集成到现有的优化器和PEFT技术中，无需更改训练流程或模型架构。研究团队提供了CPU、GPU和混合实现版本，以满足不同的需求。虽然在GPU上进行在线聚类会增加显存负担，但将聚类索引和缩放因子存储在CPU上的实现方式几乎不会增加GPU内存消耗，且对训练时间的影响可以忽略不计。
在超参数选择方面，SGG也表现出了良好的鲁棒性。聚类数量K可以根据任务特点设为2或3；重聚类间隔T可以设为总训练迭代次数的5%左右；缩放衰减率β?=0.99是一个稳健的选择。
五、SGG的工作原理深度解析
为什么SGG能够如此有效地提升大语言模型的训练效果？研究团队的分析揭示了几个关键因素：
首先，大语言模型中的梯度分布往往呈现聚类现象，不同层内的参数展现出不同但内部一致的优化行为。SGG通过动态聚类捕捉这种自然分组特性，而不是使用预定义的固定分组。
其次，虽然参数在组内表现出相似性，但它们仍然存在显著的参数级差异。SGG不是简单地用单一学习率替换组内所有参数的学习率，而是保留参数级自适应学习率，同时应用组级缩放进行调整，从而兼顾了个体差异和集体一致性。
第三，SGG的全局视角帮助缓解了大语言模型训练中常见的问题，如不同层之间训练动态的严重差异、梯度爆炸/消失等问题。通过引入全局中位数偏差作为参考，SGG促进了训练的同质化，抑制了可能导致破坏性更新的发散行为。
最后，SGG的动态聚类策略比固定分组更灵活，能够适应训练过程中参数行为的变化，就像一位优秀教师能够根据学生的实时表现调整教学策略一样。
六、研究意义与未来展望
这项研究的意义远不止于提出一种新的优化器包装器。它揭示了一种全新的思路：通过梯度分组约束来缩放自适应学习率，而不是简单地替换它们。这种方法在保持参数级适应性的同时，引入了组级和全局约束，实现了微观优化和宏观协调的平衡。
SGG的成功表明，大语言模型训练中存在大量冗余信息，通过适当的分组和缩放策略，可以显著提高训练效率和效果。特别是它使低秩预训练首次达到了全秩训练的性能水平，这对于资源受限场景下的大模型训练具有重要意义。
展望未来，SGG提供的框架具有高度灵活性，可以探索不同的分组和缩放策略。例如，可以尝试更精确的在线聚类、基于启发式的静态分区，或者甚至是学习型的分组函数。此外，SGG的思想也可能扩展到更广泛的应用场景，如图像生成、多模态学习、视觉骨干网络和专家混合体系结构等。
对于大语言模型训练的从业者来说，SGG提供了一种几乎"免费"的性能提升方案，无需更改模型架构或训练流程，就能获得更好的训练效果和更快的收敛速度。这种简单易用但效果显著的方法，有望在未来的大模型训练中得到广泛应用。
总的来说，这项由西湖大学团队领导的研究不仅提供了一种实用的技术解决方案，还为我们理解和优化大语言模型训练过程提供了新的视角。它再次证明，有时候最有效的改进不是来自复杂的算法创新，而是来自对问题本质的深刻理解和巧妙的平衡策略。
来源：至顶网一点号