前言

笔者在之前的大规模对比学习训练过程（训练CLIP）中，发现在混合精度训练时候，对比学习的交叉熵损失（带温度系数）容易出现的一个小问题，特此笔记下，希望对读者有所帮助。如有谬误请联系指出，本文遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明并且联系作者。

在对比学习训练过程中，经常会采用带温度系数的交叉熵损失，如式子(1-1)所示

其中的 为温度系数， 和 表示query和key，其中的 表示正样本，而 表示第 个样本，其中query和key的定义见之前博文 [1]，里面有更为细致的介绍，该损失函数在CLIP训练中有着应用，见博文[2]。而其中的温度系数 很关键，是用于控制整个学习任务的『难易』程度的，这里我们举个例子。用以下代码为例，分两种情况， 当训练处于初始阶段时候， 打分接近均匀分布，那么假设 ，其中0.2为正样本打分，其他为负样本打分，注意到此时label位置的打分（也即是第0维打分0.2）是比其他位置的打分还要小的。我们选择四组温度系数 ，如以下代码所示，此时我们发现温度系数 越小，其输出的损失值越大。这其实很好理解，由于 对打分进行了扩大，而这种扩大是非线性的，参考指数函数的曲线，如Fig 1.所示，当 线性扩大的时候，其指数是有着更为高阶的增大速度的。因此当 的时候，np.exp(logit/t)为[ 54.59815003 403.42879349 148.4131591 148.4131591 ]，对比其原始的打分[0.2 0.3 0.25 0.25]，我们从式子(1-1)中已经知道损失函数是正样本打分/样本打分总和的相反数，而当通过指数函数进行放大后，其损失就从原始的

变为

显然损失就被放大了，更易于模型刚开始的优化。而目标位置和其他位置的打分差别越大，其损失差别就会通过指数函数放大得更大，此时，我们认为越小的温度系数，模型训练越『简单』，可以认为此时我们将温度系数调小，是调小了整个训练任务的学习难度。

def cross_entropy(logit, t, ind_pos=0):
    print(logit)
    print(np.exp(logit/t))
    return -np.log(np.exp(logit[ind_pos]/t) / np.exp(logit/t).sum())
  
orig_score = [0.2,0.3,0.25,0.25]
orig_score = np.array(orig_score)
y0p05 = cross_entropy(orig_score, 0.05)
print(y0p05) 
# 输出 loss: 2.626523375036445
# logit: [0.2  0.3  0.25 0.25]
# np.exp(logit/t): [ 54.59815003 403.42879349 148.4131591  148.4131591 ]

y0p5 = cross_entropy(orig_score, 0.5)
print(y0p5)
# 输出 loss: 1.4887933201471417
# logit: [0.2  0.3  0.25 0.25]
# np.exp(logit/t): [1.4918247  1.8221188  1.64872127 1.64872127]

y1p0 = cross_entropy(orig_score, 1.0)
print(y1p0)
# 输出 loss: 1.4369192960265726
# logit: [0.2  0.3  0.25 0.25]
# np.exp(logit/t): [1.22140276 1.34985881 1.28402542 1.28402542]

y2p0 = cross_entropy(orig_score, 2.0)
print(y2p0)
# 输出 loss: 1.4114506070510497
# logit: [0.2  0.3  0.25 0.25]
# np.exp(logit/t): [1.10517092 1.16183424 1.13314845 1.13314845]

exp_function

Fig 1. 指数函数会将打分进行非线性扩大。

然而，当训练到一定程度时候，模型的打分已经比较置信了，假设为 ，那么同样采用以上代码，我们有输出。此时我们发现温度系数越大，其损失越大了，原因和之前分析的一样，由于指数函数的非线性造成的。此时反而应该适当增大温度系数才能让模型继续学习下去，此时我们能认为模型已经『学习』到较为简单的表征了，需要加大训练难度才能让模型学习到更为细致（或者困难）的表征，此时我们将温度系数调大，可以认为是加大了整个任务的『难度』。总得来说，在训练不同阶段通过调节温度系数的大小可以控制整个任务的难度，从而让模型学习出不同粒度的表征，可以看成是一种『coarse to fine，从粗到细』的方法。

# logit: [0.7  0.1  0.05 0.15]
# np.exp(logit/t): [1.20260428e+06 7.38905610e+00 2.71828183e+00 2.00855369e+01]
# 输出 loss: 2.5105927394272577e-05
====
# logit: [0.7  0.1  0.05 0.15]
# np.exp(logit/t): [4.05519997 1.22140276 1.10517092 1.34985881]
# 输出 loss: 0.6453200240879728
====
# logit: [0.7  0.1  0.05 0.15]
# np.exp(logit/t): [2.01375271 1.10517092 1.0512711  1.16183424]
# 输出 loss: 0.9737318345317211
====
# logit: [0.7  0.1  0.05 0.15]
# np.exp(logit/t): [1.41906755 1.0512711  1.02531512 1.07788415]
# 输出 loss: 1.1702870665310683
====

然而手动在不同阶段调节温度系数很麻烦，而且可能并不是最优的，因此就有人尝试将温度系数设置为可学习的（learnable temperature），通过梯度更新温度系数。这一点在全精度训练下并没有太大问题，但是为了加大训练的batch size，通常会在半精度场景中进行训练（auto mixed precision,amp），此时容易出现下溢出的问题，损失函数如Fig 2.所示，训练过程会突然崩溃，此来龙去脉且听笔者道来。

amp_train_loss_acc

Fig 2. 在加入了带可学习的温度系数损失函数后进行从头训练，会经常出现不稳定的现象。

我们以paddle的半精度训练为例子，如文档[3]所示，半精度即是FP16。如Fig 3.所示，此时实际情况中模型训练中的某些变量, 比如grad (特别是 activation 的 grad), 可能会因小于 FP16的精度低而变成0 （红色线条以左）; 另一方面在FP16 的表示范围的中有很大的一部分(从蓝色线最大值往右) 却没有被利用到，导致整个分布（绿色方块）只有一小部分（红线-蓝线之间）落在FP16的表示范围内。此时对梯度的数值进行整体的一个放大（或者缩小），能够更充分的利用FP16 的表示范围。AMP 会在反向开始前对 loss 进行放大，并在执行任何梯度相关操作(e.g. gradient-clip, update) 之前对 gredient 进行缩放恢复原来的大小。

通常来说，框架会提供动态损失放缩（Dynamic loss scaling）的机制，此时框架会检测scale大小是否合适，当scaling up 不足时，仍会有部分较小变量会被表示成 0而损失精度；当scaling up 过度时，变量超过FP16表示范围出现 nan or inf， 同样造成精度损失。此时动态损失放缩会采用自动检测梯度值的方法：

当连续incr_every_n_steps(int)个batch 中所有的gradient 都在FP16 的表示范围, 将scaling factor 增大incr_ratio(float)倍;当有连续decr_every_n_nan_or_inf(int)个batch 中gradient 里出现 nan / inf时, scaling factor 缩小 decr_ratio(float)倍.

如Fig 4.所示，在 Dynamic loss scaling 中，框架在每一个 iteration 都会依据当前 gradients 是否出现 nan or inf 还有用户设置的 Dynamic loss scaling 参数来动态调整 loss scaling factor 的大小，将gradient 尽量保持在 FP16 的表示范围之内。

loss_scale_amp

Fig 3. 如果不进行loss scale，那么会有很大一部分的数值在FP16中溢出。

loss_scale_dynamic

Fig 4. 正常情况下loss scale是波动上下变化的，框架会自动检测并且调整。

然而，如果在半精度场景中把温度系数设置为可学习的，通过之前的讨论我们知道温度系数控制了整个学习任务的难易程度，同时其也可以控制损失大小，也就意味着在训练到某个阶段，模型可能会倾向于连续降低温度系数以进一步降低损失（此时继续优化模型参数本身需要更多的iteration，但是『优化』温度系数会获得短期更快的损失下降）。此时温度系数会持续下降，如Fig 5.右图所示，而温度系数一直下降意味着 的值会迅速增大，而其指数函数 增大得更快，导致loss scale持续下降（否则就超过了FP16的表示范围），知道loss scale降到了0都无法挽救为止，此时训练将会出现nan或者inf，导致训练过程崩塌。

temperature_loss_scale

Fig 5. loss scale和可学习温度系数的大小变化曲线图。

怎么缓解这个问题呢？笔者实践中是将温度系数固定在0.05开始模型从头的训练，等模型训练到差不到的时候才开放温度系数的可学习性。

Reference

[1]. https://blog.csdn.net/LoseInVain/article/details/119515146, 《MoCo 动量对比学习——一种维护超大负样本训练的框架》

[2]. https://blog.csdn.net/LoseInVain/article/details/119516894, 《CLIP-对比图文多模态预训练的读后感》

[3]. https://fleet-x.readthedocs.io/en/stable/paddle_fleet_rst/fleet_collective_training_speedup_with_amp_cn.html, 《自动混合精度练加速分布式训练》

公众号后台回复“CVPR 2022”获取论文合集打包下载～