LLaMA-6B训练过程中损失呈阶梯式下降的原因探究

一、学习率调度策略的主导作用 #

这是最普遍的触发因素：

• 若采用阶梯式学习率衰减（StepLR）或带重启的余弦退火调度（CosineAnnealingWarmRestarts），单轮次内学习率维持在低固定值，模型参数更新幅度微弱，损失几乎停滞；进入新轮次时，学习率被重置（或大幅上调后逐步回落），模型获得足够更新动力，损失出现断崖式下降。
• 若设置了按轮次的学习率衰减规则（如每轮乘以固定衰减系数），新轮次初始学习率的跳变也会直接触发损失骤降。

二、数据加载与训练流程的影响 #

• 若训练数据按轮次重新洗牌且存在数据分布差异：某一轮次的数据集偏易（模型已能较好拟合），损失自然停滞；新轮次加载了更具挑战性的数据子集，模型被迫学习新模式，损失随参数更新快速下降。
• 部分训练框架在轮次切换时会重置梯度累积计数器或优化器状态（如自定义调度逻辑），也可能导致参数更新幅度突然变化。

三、LLaMA架构的参数更新特性 #

LLaMA-6B的Transformer架构存在参数更新的滞后性：

• 底层嵌入层、注意力层参数可能在单轮次内已收敛，但上层Feed-Forward网络（FFN）参数需更多梯度累积才能更新；当轮次切换带来学习率或数据分布变化时，上层参数开始大幅更新，带动整体损失骤降。
• 模型中的RMSNorm归一化层虽能稳定梯度，但也可能让模型陷入局部最优区间，只有学习率变化打破平衡后，才能跳出局部最优，损失快速下降。

四、优化器动量的累积效应 #

• 采用AdamW这类带动量的优化器时，动量参数会累积之前的梯度信息：若某一轮次内梯度普遍偏小，动量会逐渐趋近于0，参数更新停滞；进入新轮次后，新数据带来的大梯度重新激活动量，参数更新幅度突然增大，损失骤降。

内容的提问来源于stack exchange，提问作者Jing zhao