老学庵

天行健,君子以自强不息;地势坤,君子以厚德载物!

0%

LLM之余弦退火学习率

  学习率这个概念在非线性优化中经常出现,在深度学习中模型在反向传播阶段严重依赖于损失函数梯度的链式传播,为了更好的控制参数更新的步长,引入了学习率的概念: \[ w_{new} = w_{old} - \eta \ \Delta J(w) \]

  其中的\(\eta\)为学习率,直观来说,非线性优化的过程就像下山,学习率就是步长,步子迈的太大或太小都不太理想。当步子迈得的太大了,你可能就直接跨过山脚,对应到模型训练中,就是学习率太高,每次更新参数时,跳过了最小值,导致在最优值附近来回振荡,无法收敛;步子太小,则需要走很久,才能到达山脚,对应到模型训练中,就是学习率太低,模型收敛的速度很慢,需要较长时间的训练才能收敛到最优值。举这个经典的下山的例子是为了说明学习率在模型训练中的重要意义。   在深度学习的模型训练中学习率有很多使用策略,包括不限于:

  • 固定学习率

  顾名思义就是在模型训练过程中使用固定学习率来进行模型训练,这样会难以平衡初期的收敛速度和后期的精细调整,一般很少使用,只在一些简单问题上应用。

  • 阶梯式衰减学习率(step decay)

  指学习率在模型训练的特定步骤中衰减,会导致学习率变化不平滑,模型训练不稳定。

  • 指数衰减学习率(exponential decay)

  指学习率在模型训练过程中按指数函数衰减,由于指数函数的高增长特性,前期往往能快速收敛,但是后期由于学习率衰减过快,导致训练停滞,收敛速度很慢。

  • 线性预热+线性衰减

  指学习率在模型训练过程中先线性增加再线性较少,此方案在大语言模型的预训练被广泛使用,缺点在于是线性变化,变化不太平滑。

  • 余弦退火学习率(Cosine Annealing)

  余弦退火学习率的公式如下: \[ \eta_{t} = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})(1+\cos(\frac{t}{T}\pi)) \]   其中,\(\eta_{t}\)是第\(t\)步的学习率,\(\eta_{min}\)是训练过程中能接受的最小学习率,\(\eta_{max}是最大学习率\)\(t\)为当前的训练步树,\(T\)该epoch中的总训练步数。从公式可以看出,在训练的开始阶段,学习率解决最大值,模型的收敛速度较快,到达一定训练步数后,学些率开始减少,模型会在更小的参数空间内搜索最优值,随着训练步数的增长,学习率逐渐逼近最小值,模型在最优值附近的参数空间内搜索。

  同时学习率是按照余弦曲线的方式进行更新,变化平滑的同时,不会出现变化幅度不大。目前基于余弦退火学习率衍生出一些更精细的学习率调度算法,如:

  • 带热重启的余弦退火

      学习率在训练过程中周期性回到最大值,然后再次降低,公式如下:

    \[ \eta_{t} = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})(1+\cos(\frac{T_{cur}}{T_i}\pi)) \]   其中,\(T_{cur}\)是自上次热重启之后的训练步数,\(T_i\)是当前周期的总步数。这种学习率调度方式有助于跳出局部最优解,探索更广阔的参数空间

  • 带预热的余弦退火

      带预热的余弦退火通过两阶段学习率调整实现:

    • ​预热阶段

      训练初期学习率从较小值线性或非线性增长至预设的最大值,避免随机初始化权重下的大幅度参数更新引发振荡。

    • 余弦退火阶段

        预热完成后,学习率按余弦函数周期性衰减,从最大值逐渐降至最小值,数学公式如下: \[ \eta_{t}= \begin{cases} \eta_{start} + \dfrac{t}{T_{warmup}}(\eta_{max} - \eta_{start}) \qquad 0 \leq t < T_{warmup} \\ \eta_{max} + \frac{1}{2}(\eta_{max} - \eta_{min})(1+cos(\frac{t}{T_{cos}}\pi)) \qquad t \geq T_{warmup} \end{cases} \] \(T_{warmup}\)为预热步数,\(T_{cos}\)为退火周期

  • 循环余弦退火(cyclical cosine annealing)

      顾名思义,就是在模型训练过程中,学习率的变化在余弦退火的基础上存在多个余弦周期,通过重启(Restart)机制恢复学习率至最大值,形成多周期探索-利用循环,可以平衡梯度在参数空间的探索与利用,避免剧烈波动。

  余弦退火学习率策略通过平滑地调整学习率,帮助模型在训练初期快速收敛,在训练后期精细调整参数,从而找到更好的局部最优解。它的提出和发展极大地推动了深度学习模型训练的稳定性和效率,已成为现代深度学习训练的标准配置之一。

  虽然它不是万能的解决方案,但在大多数深度学习任务中,特别是大型语言模型的训练中,余弦退火都能提供显著的性能提升。通过理解其工作原理和适用边界,我们可以更好地利用这一强大工具,提高模型训练效果。

编程实现

  • demo 1

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    from torch.optim.lr_scheduler import CosineAnnealingLR
    import torch

    # 创建模型和优化器
    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 创建余弦退火调度器
    # T_max是半个周期的长度,通常设为总训练步数
    # eta_min是最小学习率
    scheduler = CosineAnnealingLR(optimizer, T_max=1000, eta_min=0.0001)

    # 在训练循环中使用
    for epoch in range(num_epochs):
    for batch in data_loader:
    # 训练步骤
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    # 更新学习率
    scheduler.step()

  • demo 2

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    import math
    def get_cosine_lr(current_step, total_steps, max_lr, min_lr):
    """计算当前步骤的余弦退火学习率"""
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * current_step / total_steps))

    # 在训练循环中使用
    max_lr = 0.001
    min_lr = 0.0001
    total_steps = 1000

    for step in range(total_steps):
    # 计算当前学习率
    current_lr = get_cosine_lr(step, total_steps, max_lr, min_lr)

    # 更新优化器的学习率
    for param_group in optimizer.param_groups:
    param_group['lr'] = current_lr

    # 正常的训练步骤
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

  • demo 3

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

    # 创建带热重启的余弦退火调度器
    # T_0是第一次重启前的迭代次数
    # T_mult是控制重启周期如何变化的因子
    scheduler = CosineAnnealingWarmRestarts(
    optimizer,
    T_0=1000, # 第一个周期的长度
    T_mult=2, # 每次重启后周期长度翻倍
    eta_min=0.0001 # 最小学习率
    )

    # 在训练循环中使用
    for epoch in range(num_epochs):
    for batch in data_loader:
    # 训练步骤
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    # 更新学习率
    scheduler.step()