©PaperWeekly 原创 · 作者 | 苏剑林
单位 | 追一科技
研究方向 | NLP、神经网络
在之前的文章 用时间换取效果:Keras 梯度累积优化器 中,我们介绍过“梯度累积”,它是在有限显存下实现大 batch_size 效果的一种技巧。一般来说,梯度累积适用的是 loss 是独立同分布的场景,换言之每个样本单独计算 loss,然后总 loss 是所有单个 loss 的平均或求和。然而,并不是所有任务都满足这个条件的,比如最近比较热门的对比学习,每个样本的 loss 还跟其他样本有关。
那幺,在对比学习场景,我们还可以使用梯度累积来达到大 batch_size 的效果吗?本文就来分析这个问题。
简介
一般情况下,对比学习的 loss 可以写为:
这里的 b 是 batch_size;是事先给定的标签,满足,它是一个 one hot 矩阵,每一列只有一个 1,其余都为0;而是样本 i 和样本 j 的相似度,满足,一般情况下还有个温度参数,这里假设温度参数已经整合到中,从而简化记号。模型参数存在于中,假设为。
可以验证,一般情况下:
所以直接将小 batch_size 的对比学习的梯度累积起来,是不等价于大 batch_size 的对比学习的。类似的问题还存在于带 BN(Batch Normalization)的模型中。
梯度
注意,刚才我们说的是常规的简单梯度累积不能等效,但有可能存在稍微复杂一些的累积方案的。为此,我们分析式(1)的梯度:
其中表示不需要对求的梯度,也就是深度学习框架的 stop_gradient 算子 。上式表明,如果我们使用基于梯度的优化器,那幺使用式 (1) 作为 loss,跟使用作为 loss,是完全等价的(因为算出来的梯度一模一样)。
内积
接下来考虑的计算,一般来说它是向量的内积形式,即,参数在里边,这时候:
所以 loss 中的可以替换为而效果不变:
第二个等号源于将那一项的求和下标 i,j 互换而不改变求和结果。
流程
式(5)事实上就已经给出了最终的方案,它可以分为两个步骤。第一步就是向量:
的计算,这一步不需要求梯度,纯粹是预测过程,所以 batch_size 可以比较大;第二步就是把当作“标签”传入到模型中,以为单个样本的 loss 进行优化模型,这一步需要求梯度,但它已经转化为每个样本的梯度和的形式了,所以这时候就可以用常规的梯度累积了。
假设反向传播的最大 batch_size 是 b,前向传播的最大 batch_size 是 nb,那幺通过梯度累积让对比学习达到 batch_size 为 nb 的效果,其格式化的流程如下:
1. 采样一个 batch 的数据,对应的标签矩阵为,初始累积梯度为 g=0;
2. 模型前向计算,得到编码向量以及对应的概率矩阵;
3. 根据式(6)计算标签向量;
4. 对于,执行:
5. 用 g 作为最终梯度更新模型,然后重新执行第 1 步。
总的来说,在计算量上比常规的梯度累积多了一次前向计算。当然,如果前向计算的最大 batch_size 都不能满足我们的需求,那幺也可以分批前向计算,因为我们只需要把各个算出来存好,而可以基于算出来。
最后还要提醒的是,上述流程只是在优化时等效于大 batch_size 模型,也就是说的梯度等效于原 loss 的梯度,但是它的值并不等于原 loss 的值,因此不能用作为 loss 来评价模型,它未必是单调的,也未必是非负的,跟原来的 loss 也不具有严格的相关性。
问题
上述流程有着跟《节省显存的重计算技巧也有了 Keras 版了》 [1] 介绍的“重计算”一样的问题,那就是跟 Dropout 并不兼容,这是因为每次更新都涉及到了多次前向计算,每次前向计算都有不一样的 Dropout,这意味着我们计算标签向量时所用的跟计算梯度时所用的并不是同一个,导致计算出来的梯度并非最合理的梯度。
这没有什幺好的解决方案,最简单有效的方法就是在模型中去掉 Dropout。这对于 CV 来说没啥大问题,因为 CV 的模型基本也不见 Dropout 了;对于 NLP 来说,第一反应能想到的结果就是 SimCSE 没法用梯度累积,因为 Dropout 是 SimCSE 的基础。
小结
本文分析了对比学习的梯度累积方法,结果显示对比学习也可以用梯度累积的,只不过多了一次前向计算,并且需要在模型中去掉 Dropout。本文同样的思路还可以分析 BN 如何使用梯度累积,有兴趣的读者不妨试试。
参考文献
[1] https://kexue.fm/archives/7367
Be First to Comment