论文《Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling》发表于NeurIPS 2024，由腾讯混元及北京大学等机构的研究人员撰写。论文主要研究Adam式优化器中最优学习率与批量大小的关系，通过理论分析和实验验证得出了新的缩放定律，为深度学习训练中的超参数调整提供了重要参考。

  1. 研究背景

    • 优化器发展：深度学习领域中，Adam式优化器（如Adam、Adagrad等）广泛应用，其通过梯度符号更新模型参数，收敛曲线更稳定。
    • 超参数重要性：学习率和批量大小是优化器的关键超参数，影响模型收敛。此前针对SGD式优化器的研究提出了一些学习率与批量大小的关系规则，但这些规则不适用于Adam式优化器。

  2. 理论分析

    • 近似Adam式优化器：Adam式优化器与SGD的主要区别在于使用梯度符号进行更新，随着批量大小增加，更新量期望值趋于饱和。
    • 推导最优学习率：基于给定小批量梯度估计推导最优学习率公式。假设梯度服从高斯分布，得出最优学习率公式，其分子是关于批量大小的一阶函数，分母是二阶函数。
    • 不同批量大小下的最优学习率：当 时，最优学习率 ，在 时达到峰值；当 时，最优学习率 ，且当无限增大时，最优学习率收敛到非零值。
    • 数据/时间效率权衡：Adam优化器在批量大小选择上，数据和时间效率的权衡结论与SGD情况一致， 不仅是最优学习率的局部最大值，也是训练速度和数据效率的平衡点，且随着训练进行， 会逐渐增大。

  3. 实验验证

    • 实验设置：选用4个开源工作负载，包括在Fashion - MNIST上训练5层CNN模型、在Tiny - ImageNet上训练ResNet18模型、在ELI5 - Category数据集上训练dense Transformer模型以及在RedPajama - v2数据集上训练MoE模型。使用Adam优化器，在不同的梯度符号配置和默认超参数下进行实验，实验在NVIDIA A100卡上执行。
    • 变量估计：通过曲线拟合估计 和 的期望值。利用网格搜索结果中达到指定性能的实际步数和处理的训练样本数，进行线性拟合得到 的估计值。
    • 实验结果：在不同工作负载下，实验结果均验证了理论预测。如在CNN - FashionMNIST工作负载中，理论趋势与实际最优学习率表现相符；在ResNet18 - TinyImageNet工作负载中，随着训练进行， 逐渐增大；在DistilGPT2 - Eli5Category工作负载中，不同Adam配置下的实验结果都与理论曲线一致。

  4. 研究讨论

    • 应用价值：该研究有助于深度学习模型训练时更有效地调整超参数，提高收敛速度，避免过度的网格搜索。
    • 研究局限：实际应用中存在多种影响学习过程的因素，该理论基于损失函数的二次近似，可能无法涵盖所有复杂训练配置情况。
    • 未来方向：可探索自适应学习率和批量大小调整策略，研究更高阶近似对模型训练的影响。

  5. 研究结论：论文建立了Adam式优化器中最优学习率与批量大小的缩放定律，理论证明了最优学习率随批量大小先增后减，峰值代表训练速度和数据效率的权衡点，并通过大量实验在多种深度学习模型和数据集上验证了该理论。