摘要：在人工智能技术飞速发展的今天，深度学习模型的训练成本正在迅猛增长。这项来自北京大学智能科学与技术学院的唐安达、董一鸣、曾宇涛、周勋及林宙辰教授领导的研究团队，在2025年5月发表于arXiv（arXiv:2505.24452v1）的论文《Stepsize an
想象一下，你计划去一个陌生的城市旅行，但只有固定的三天时间。如何在这有限的时间里游览最多的景点？是快速浏览所有景点还是精选几个深入游览？深度学习训练面临类似的挑战：当计算资源有限，只能进行固定次数的训练迭代时，如何调整学习率以获得最佳模型性能？
随着模型规模和数据集的不断扩大，训练深度神经网络的计算成本已经成为研究人员和企业面临的重大瓶颈。虽然硬件技术在不断进步，但模型训练的资源需求增长得更快。这种情况下，预算迭代训练（budgeted-iteration training）变得尤为重要，它的目标是在预先确定的迭代预算内实现最佳学习效果。
学习率调度作为深度学习训练的核心组成部分，直接影响不同网络和任务的性能，特别是在预算迭代场景中。然而，目前的学习率调度设计主要基于启发式方法，缺乏理论基础，而且寻找最佳学习率调度往往需要大量的试错实验，使训练过程效率低下。
北京大学的研究团队通过构建一个新颖的训练预算感知优化框架，明确考虑了对景观曲率变化的鲁棒性，推导出了UBA调度器。这个调度器只需一个超参数φ就能在灵活性和简单性之间提供完美平衡，免去了针对每个网络进行数值优化的麻烦。更重要的是，团队建立了φ与条件数之间的理论联系，为他们的方法提供了解释和合理性。
研究团队进行了大量实验，结果显示UBA在多种视觉和语言任务中，跨越不同网络架构（如ResNet、OLMo）和规模，在不同训练迭代预算下，始终优于常用的调度方法。这意味着，无论你是训练一个小型图像分类模型还是大型语言模型，无论你有一天还是一周的训练时间，UBA都能帮助你获得最佳性能。
一、UBA调度器的诞生背景
随着人工智能技术的飞速发展，模型规模和数据集的不断扩大，训练深度神经网络的计算成本已经成为一个严峻的挑战。想象一下，几年前训练一个模型可能只需要几个小时，而现在，训练一个大型语言模型可能需要数周甚至数月的时间，消耗的电力足以支撑一个小镇的日常用电。
在这种情况下，预算迭代训练的概念应运而生。所谓预算迭代训练，简单来说就是在严格限定的训练迭代次数下，尽可能地提高模型性能。这就像是你只有三天假期，却想尽可能多地游览一座城市的景点——你需要精心规划每一天的行程，以确保在有限的时间内获得最佳体验。
学习率调度是深度学习训练中的关键组成部分，它决定了模型在训练过程中"学习"的速度。想象学习率就像是学习舞蹈时的节奏：开始时节奏可能较快，让你迅速掌握基本动作；随着学习的深入，节奏会逐渐放慢，让你更精细地打磨每个动作的细节。
然而，现有的学习率调度方法大多是基于经验或启发式设计的，缺乏坚实的理论基础。比如，有些方法建议在训练开始时使用较大的学习率，然后逐渐减小；有些方法则推荐周期性地调整学习率。但这些方法往往需要大量的试错实验，才能找到适合特定任务的最佳调度策略，这在计算资源有限的情况下显得尤为低效。
北京大学的研究团队敏锐地发现了这一问题，他们提出了一个自然而关键的问题：是否存在一种理论上有根据的、统一的调度方法，能够在不同任务、网络、规模和训练预算下都表现出色，从而避免繁琐的试错过程？
二、统一预算感知优化框架的构建
为了回答这个问题，研究团队首先构建了一个统一的预算感知训练优化框架。这个框架不同于传统的优化方法，它明确考虑了训练过程中景观曲率变化的鲁棒性。
想象一下地形的变化。在平坦的草原上行走，你可以大步前进；而在崎岖的山地上，你需要小心翼翼地调整步伐。深度学习的优化景观也是如此——有些区域平坦，有些区域陡峭。传统的学习率调度方法往往没有充分考虑这种地形变化，而研究团队的框架则将其纳入考量。
具体来说，他们定义了一个有限优化问题：在给定的神经网络和数据集下，找到一组学习率，使得在最坏情况下（即最具挑战性的优化景观）仍能在固定的训练迭代次数内最小化损失函数。这就像是设计一条hiking路线，即使在最崎岖的地形下，也能在规定时间内到达目的地。
团队采用二阶泰勒展开来近似损失函数，并考虑了Hessian矩阵的特征值分布，这些特征值反映了优化景观的曲率。他们将问题建模为一个最小-最大优化问题：
min_{η1,η2,...,ηT-1} max_{f∈F} L(f(WT, ξ)) s.t. Wt+1 = Wt - ηt?L(f(Wt, ξ)) t = 0, 1, 2, ..., T - 1
其中，η_t是第t次迭代的学习率，T是最大训练迭代次数，F是函数类，L是损失函数。
通过一系列的数学推导，团队将这个问题简化为寻找一组学习率，使得在最坏情况下（即对任何可能的Hessian特征值），权重更新的放大因子最小。这就像是设计一个万能的徒步策略，无论遇到什么地形，都能以最稳定的步伐前进。
三、UBA调度器的导出与特性
为了解决这个优化问题，研究团队采用了交替的投影梯度法，在学习率变量和景观参数之间迭代优化。但是，每次应用他们的方法到新的网络时，都需要重新进行数值优化，这显然不够实用。
为了避免这种重复的数值优化，团队提出了一个通用的参数函数来近似数值解。经过仔细的曲线拟合和分析，他们得到了统一预算感知（UBA）调度器的表达式：
ηt = (ηmax - ηmin) * [2(1 + cos((2(t-Tk)-1)π/(2(Tk+1-Tk)) + (k-1)π))] / [2φ + (2-φ)(1 + cos((2(t-Tk)-1)π/(2(Tk+1-Tk)) + (k-1)π))] + ηmin
这个公式看起来可能有点复杂，但它的核心思想很简单：学习率的变化由一个余弦函数控制，而参数φ决定了这个变化的速度和形状。想象φ就像是调节音乐节奏的旋钮——不同的φ值会产生不同的学习"节奏"。
UBA调度器有几个显著的特点：
首先，它只需一个超参数φ，大大简化了调优过程。这就像是只需调整一个旋钮，就能适应不同的舞蹈节奏，而不需要精通各种复杂的舞步。
其次，研究团队证明了φ与条件数（优化难度的度量）之间存在理论联系。条件数越大，优化难度越大，需要的φ值也越大。这为参数选择提供了理论指导，而不仅仅是凭经验。
第三，UBA调度器可以通过简单的参数调整，近似模拟现有的各种学习率调度方法，如余弦退火、步进衰减、循环调度等。但更重要的是，它在训练收敛性和最终精度上超越了这些方法。
此外，研究团队还证明了不同φ值下的收敛性，并通过理论分析和实验结果为φ的选择提供了实用指南。
四、UBA调度器的实验验证
为了验证UBA调度器的有效性，研究团队进行了全面的实验，涵盖了视觉和语言任务，跨越不同的网络架构和规模，在不同的训练迭代预算下进行测试。
在视觉分类任务中，团队使用了CIFAR10/100和ImageNet数据集，分别采用VGG16、ResNet18、ResNet34、ResNet50等不同架构。他们独立训练了模型，使用最大训练轮次的25%、50%和100%作为固定的epoch预算，而不是重用或插值较长训练运行的结果。
实验结果表明，UBA在所有训练预算下都展现出最强的性能，不仅在小规模基准测试（使用ResNet18/34的CIFAR10/100）上取得优异结果，也在大规模基准测试（使用ResNet50的ImageNet）上保持一致的优势。这种一致性的改进突显了UBA跨模型和数据集规模的泛化能力。
值得注意的是，UBA不仅在100%训练预算下优于基线，在25%和50%迭代预算下也表现出色，证明了它在计算资源受限场景中的效率。更重要的是，虽然第二好的调度方法会根据数据集、架构和训练预算的不同而变化，但UBA始终保持稳定的优越性。这对于寻求可靠调度方法而不需要进行广泛方法选择的实践者来说，UBA提供了一个默认的强大选择。
在语言模型评估方面，团队在OLMo模型上进行了测试，该模型是一个基于仅解码器变换器架构的真正开放的语言模型。他们调整了模型大小和训练步骤以探索不同的预算约束，评估了UBA在跨越四种参数规模的OLMo网络上的表现：36M、73M、151M和300M，涵盖了从普通到大规模的模型。
结果显示，UBA在所有规模中约50%的基准测试上达到了最先进的性能，同时在平均得分上始终优于基线。它在SciQ-73M(+1.7)和ARC-E-300M(+2.63)上表现出显著的改进，突显了它在不同基准测试中增强泛化能力的能力。
更令人印象深刻的是，UBA在整个训练过程中始终实现更低的训练损失和验证损失，表明它具有高效的训练能力和下游性能增强能力。
五、参数φ的影响与跨优化器性能
为了更深入地理解UBA调度器，研究团队进行了一系列消融研究，探索了参数φ的敏感性以及跨优化器的性能。
首先，团队在CIFAR100数据集和ImageNet数据集上使用AdamW优化器进行了实验。结果表明，UBA在SGD和AdamW优化器上都实现了最先进的验证准确率，在CIFAR100-ResNet34和ImageNet-ResNet50基准测试中始终优于基线。这证明了虽然UBA是从标准梯度下降动态理论推导出来的，但它能有效地泛化到现代优化器，如AdamW，尽管后者引入了额外的动量和自适应机制。
接着，团队对方程中的关键参数φ进行了敏感性分析，评估了φ∈{0.25, 0.5, 1.0, 2.5, 5, 10}在不同优化场景中的表现。实验结果揭示了一个有趣的二分法：
在使用AdamW时，较小的φ值（如φ=0.5）产生了更好的性能；而在使用SGD时，较大的φ值（如φ=5）表现更优。团队将这种现象归因于AdamW的预条件效应。
与SGD不同，AdamW通过用第二矩估计的平方根缩放梯度来调整学习率。在梯度较大的区域（表明周围景观较陡峭），分母也较大，有效地减小了学习率。这种自适应行为模拟了预条件，隐式降低了优化景观的条件数。
研究团队的理论框架建立了参数φ与局部极小值周围景观条件数之间的联系：条件数较低的问题（优化难度较小）适合较小的φ，而条件数较大的问题（优化难度较大）则适合较大的φ。由于AdamW的预条件效应自然缓解了病态条件，因此使用较小的φ值的调度更受青睐。相比之下，SGD缺乏这种自适应机制，因此在这种情况下更适合使用较大的φ值。
这种理论与实验之间的一致性凸显了根据优化器特性调整φ的重要性。因此，研究团队建议在面对具有挑战性的优化难度时选择较高的φ值，而在优化难度较小的场景中使用较低的φ值。
六、多阶段设计与局限性
UBA调度器具有周期性阶段基础的学习率调整设置，其中第k阶段的学习率由损失景观的第k个局部最小值动态确定。在原始概念中，这种设计捕获了连续局部最小值周围的优化动态，超参数φ与每个阶段的难度相关。
研究团队进行了实验，通过改变K值来验证他们的调度策略。他们观察到，当在每个阶段大致设置φ为常数时，随着训练的进行，性能会下降。这表明固定的φ策略无法适应多阶段调度的变化景观。由于优化难度应该在训练过程中减小，团队随着阶段的增加而降低φ。通过在每个阶段精细评估φ，他们实现了性能提升，确认了动态控制的必要性。
这些结果突显了一个基本的权衡：当多阶段调度的每个阶段的φ值可以精细评估时，多阶段调度比单阶段调度更好地捕获损失景观的非平稳行为。然而，为多阶段调度选择最优φ值仍然不是一件容易的事，这促使未来工作关注自动化的景观感知φ调整。
尽管存在这些开放性问题，UBA的有效性、实现简单性和调整容易性使其成为深度学习实践者的必备工具。
七、总结与未来展望
北京大学研究团队构建了一个统一的预算感知训练优化框架，该框架能够自然适应景观曲率变化，增强了训练的鲁棒性。基于这个框架，他们提出了统一预算感知（UBA）调度器，这是一种理论上有根据的学习率调度方法，在不同的架构和任务下，在不同的受限训练预算下，始终优于常用的调度方法。
UBA调度器只需一个超参数φ，在灵活性和简单性之间提供了完美平衡。团队建立了φ与条件数之间的理论联系，为参数选择提供了理论指导。他们还证明了不同φ值下的收敛性，并通过理论分析和实验结果为φ的选择提供了实用指南。
理论和实验结果都表明，参数φ与训练过程的优化难度相关，这影响了φ的最优选择和UBA的性能。然而，φ与优化难度之间的明确关系仍然没有被充分探索，也没有建立评估指标来量化优化难度。这些限制激发了未来关于优化难度感知φ调整的工作。
尽管存在这些开放性问题，UBA的有效性、实现简单性和调整容易性使其成为深度学习实践者的必备工具。这项研究为学习率调度领域提供了新的理论基础和实践指导，有望促进更多关于学习率调度的研究。
对于普通开发者和研究人员来说，UBA调度器提供了一个简单而强大的工具，可以在计算资源有限的情况下，通过优化学习率调度来最大化模型性能。无论是训练一个小型图像分类模型还是大型语言模型，无论是有一天还是一周的训练时间，UBA都能帮助你获得最佳性能。
这就像是无论你去哪座城市旅行，无论你有多少时间，都能帮你规划出最佳行程的智能向导——只需要你根据旅行的难度（城市的复杂程度）调整一个参数，就能获得最佳体验。
未来，随着深度学习模型和数据集规模的不断增长，计算资源的有效利用将变得越来越重要。UBA调度器为这一挑战提供了一个有力的解决方案，有望在预算有限的训练场景中发挥重要作用。
来源：至顶网一点号