SakanaAI/AI-Scientist项目：分层学习率对Grokking现象的影响研究

申梦珏Efrain

于 2025-06-03 09:12:30 发布

阅读量346

点赞数 5

CC 4.0 BY-SA版权

本文链接：https://blog.csdn.net/gitblog_00259/article/details/148394272

SakanaAI/AI-Scientist项目：分层学习率对Grokking现象的影响研究

AI-Scientist The AI Scientist: Towards Fully Automated Open-Ended Scientific Discovery 🧑‍🔬 项目地址: https://gitcode.com/gh_mirrors/ai/AI-Scientist

引言

在深度学习模型训练过程中，学习率的选择对模型性能有着至关重要的影响。传统方法通常对所有层使用统一的学习率，但近年来，分层学习率策略逐渐受到关注。本文基于SakanaAI/AI-Scientist项目中的实验数据，深入探讨分层学习率对Transformer模型中Grokking现象的影响。

什么是Grokking现象？

Grokking是指神经网络在长时间训练后突然"顿悟"的现象，表现为验证准确率在经过一段时间的停滞或波动后突然大幅提升。这种现象在数学运算等任务中尤为常见。

实验设计

实验采用Transformer模型，在四种不同任务上进行测试：

除法运算(x_div_y)
减法运算(x_minus_y)
加法运算(x_plus_y)
排列组合(permutation)

实验对比了以下几种学习率配置方案：

基准测试(Run 0)

统一学习率：所有层使用相同的学习率

分层学习率方案

Run 1：激进分层
- 嵌入层：1e-4
- 底层Transformer：1e-3
- 高层Transformer：1e-2
Run 2：温和分层
- 嵌入层：5e-4
- 底层Transformer：1e-3
- 高层Transformer：2e-3
Run 3：优化分层
- 嵌入层：7e-4
- 底层Transformer：1.5e-3
- 高层Transformer：3e-3
Run 4：微调分层
- 嵌入层：8e-4
- 底层Transformer：2e-3
- 高层Transformer：3e-3

实验结果分析

基准测试表现

基准测试中，模型在简单任务(x_div_y, x_minus_y, x_plus_y)上最终达到了100%的训练和验证准确率，但在复杂的permutation任务上表现不佳，验证准确率仅为3.6%。

分层学习率的影响

收敛速度：
- 所有分层学习率方案都显著加快了模型达到99%验证准确率的速度
- 例如在x_div_y任务中，Run 1仅需906步，而基准测试需要4200步
最终性能：
- 激进分层(Run 1)虽然加速了收敛，但牺牲了最终准确率
- 优化后的分层(Run 3)不仅保持了快速收敛，还提高了最终准确率
- 在permutation任务中，Run 3将验证准确率从3.6%(基准)提升至99.95%
任务差异性：
- 简单运算任务对学习率变化相对稳健
- 复杂任务(如permutation)对学习率配置更为敏感

关键发现

分层学习率能有效加速Grokking：
- 通过为不同层分配不同学习率，可以显著减少模型达到高准确率所需的训练步数
平衡是关键：
- 过于激进的分层(Run 1)会导致训练不稳定
- 过于保守的分层(Run 2)无法充分发挥分层优势
- Run 3找到了最佳平衡点
复杂任务受益更明显：
- 在permutation任务中，优化后的分层学习率带来了质的飞跃

最佳实践建议

基于实验结果，我们推荐以下分层学习率配置策略：

嵌入层：
- 使用中等偏低的学习率(7e-4至8e-4)
- 这些层通常需要更稳定的更新
底层Transformer：
- 使用中等学习率(1.5e-3至2e-3)
- 这些层负责基础特征提取
高层Transformer：
- 使用较高学习率(3e-3)
- 这些层负责高级特征组合和输出

技术实现要点

在PyTorch中实现分层学习率，可以使用参数组(parameter groups)功能：

optimizer = torch.optim.Adam([
    {'params': model.embedding.parameters(), 'lr': 7e-4},
    {'params': model.transformer.layers[0].parameters(), 'lr': 1.5e-3},
    {'params': [*model.transformer.layers[1].parameters(), 
                *model.output.parameters()], 'lr': 3e-3}
])