对比学习可以使用梯度累积吗?
By 苏剑林 | 2021-06-17 | 54330位读者 |在之前的文章《用时间换取效果:Keras梯度累积优化器》中,我们介绍过“梯度累积”,它是在有限显存下实现大batch_size效果的一种技巧。一般来说,梯度累积适用的是loss是独立同分布的场景,换言之每个样本单独计算loss,然后总loss是所有单个loss的平均或求和。然而,并不是所有任务都满足这个条件的,比如最近比较热门的对比学习,每个样本的loss还跟其他样本有关。
那么,在对比学习场景,我们还可以使用梯度累积来达到大batch_size的效果吗?本文就来分析这个问题。
简介 #
一般情况下,对比学习的loss可以写为
\begin{equation}\mathcal{L}=-\sum_{i,j=1}^b t_{i,j}\log p_{i,j} = -\sum_{i,j=1}^b t_{i,j}\log \frac{e^{s_{i,j}}}{\sum\limits_j e^{s_{i,j}}}=-\sum_{i,j=1}^b t_{i,j}s_{i,j} + \sum_{i=1}^b \log\sum_{j=1}^b e^{s_{i,j}}\label{eq:loss}\end{equation}
这里的$b$是batch_size;$t_{i,j}$是事先给定的标签,满足$t_{i,j}=t_{j,i}$,它是一个one hot矩阵,每一列只有一个1,其余都为0;而$s_{i,j}$是样本$i$和样本$j$的相似度,满足$s_{i,j}=s_{j,i}$,一般情况下还有个温度参数,这里假设温度参数已经整合到$s_{i,j}$中,从而简化记号。模型参数存在于$s_{i,j}$中,假设为$\theta$。
可以验证,一般情况下:
\begin{equation}-\sum_{i,j=1}^{2b} t_{i,j}\log p_{i,j} \neq -\sum_{i,j=1}^{b} t_{i,j}\log p_{i,j}-\sum_{i,j=b+1}^{2b} t_{i,j}\log p_{i,j}\end{equation}
所以直接将小batch_size的对比学习的梯度累积起来,是不等价于大batch_size的对比学习的。类似的问题还存在于带BN(Batch Normalization)的模型中。
梯度 #
注意,刚才我们说的是常规的简单梯度累积不能等效,但有可能存在稍微复杂一些的累积方案的。为此,我们分析式$\eqref{eq:loss}$的梯度:
\begin{equation}\begin{aligned}
\nabla_{\theta}\mathcal{L} =&\, -\sum_{i,j=1}^b t_{i,j}\nabla_{\theta}s_{i,j} + \sum_{i=1}^b \nabla_{\theta}\log\sum_{j=1}^b e^{s_{i,j}} \\
=&\, -\sum_{i,j=1}^b t_{i,j}\nabla_{\theta}s_{i,j} + \sum_{i,j=1}^b p_{i,j}\nabla_{\theta} s_{i,j} \\
=&\,\nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j}
\end{aligned}\end{equation}
其中$p_{i,j}^{(sg)}$表示不需要对$p_{i,j}$求$\theta$的梯度,也就是深度学习框架的stop_gradient算子。上式表明,如果我们使用基于梯度的优化器,那么使用式$\eqref{eq:loss}$作为loss,跟使用$\sum\limits_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j}$作为loss,是完全等价的(因为算出来的梯度一模一样)。
内积 #
接下来考虑$\nabla_{\theta}s_{i,j}$的计算,一般来说它是向量的内积形式,即$s_{i,j}=\langle h_i, h_j\rangle$,参数$\theta$在$h_i,h_j$里边,这时候
\begin{equation}\nabla_{\theta}s_{i,j}=\langle \nabla_{\theta}h_i, h_j\rangle + \langle h_i, \nabla_{\theta}h_j\rangle = \nabla_{\theta}\left(\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle\right)\end{equation}
所以loss中的$s_{i,j}$可以替换为$\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle$而效果不变:
\begin{equation}\begin{aligned}
\nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j} =&\, \nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)\left(\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle\right)\\
=&\, 2\nabla_{\theta}\sum_{i,j=1}^b \left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)\langle h_i, h_j^{(sg)}\rangle\\
=&\,\nabla_{\theta}\sum_{i=1}^b \left\langle h_i, 2\sum_{j=1}^b\left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)h_j^{(sg)}\right\rangle
\end{aligned}\label{eq:g}\end{equation}
其中$2\overline{p_{i,j}^{(sg)}}=p_{i,j}^{(sg)} + p_{j,i}^{(sg)}$,第二个等号源于将$\langle h_i^{(sg)}, h_j\rangle$那一项的求和下标$i,j$互换而不改变求和结果。
流程 #
式$\eqref{eq:g}$事实上就已经给出了最终的方案,它可以分为两个步骤。第一步就是向量
\begin{equation}\tilde{h}_i = 2\sum_{j=1}^b\left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)h_j^{(sg)}\label{eq:h}\end{equation}
的计算,这一步不需要求梯度,纯粹是预测过程,所以batch_size可以比较大;第二步就是把$\tilde{h}_i$当作“标签”传入到模型中,以$\langle h_i, \tilde{h}_i\rangle$为单个样本的loss进行优化模型,这一步需要求梯度,但它已经转化为每个样本的梯度和的形式了,所以这时候就可以用常规的梯度累积了。
假设反向传播的最大batch_size是$b$,前向传播的最大batch_size是$nb$,那么通过梯度累积让对比学习达到batch_size为$nb$的效果,其格式化的流程如下:
1、采样一个batch的数据$\{x_i\}_{i=1}^{nb}$,对应的标签矩阵为$\{t_{i,j}\}_{i,j=1}^{nb}$,初始累积梯度为$g=0$;
2、模型前向计算,得到编码向量$\{h_i\}_{i=1}^{nb}$以及对应的概率矩阵$\{p_{i,j}\}_{i,j=1}^{nb}$;
3、根据式$\eqref{eq:h}$计算标签向量$\{\tilde{h}_i\}_{i=1}^{nb}$;
4、对于$k=1,2,\cdots,n$,执行:
$g \leftarrow g + \nabla_{\theta}\sum\limits_{i=(k-1)b+1}^{kb} \langle h_i, \tilde{h}_i\rangle$
5、用$g$作为最终梯度更新模型,然后重新执行第1步。
总的来说,在计算量上比常规的梯度累积多了一次前向计算。当然,如果前向计算的最大batch_size都不能满足我们的需求,那么也可以分批前向计算,因为我们只需要把各个$\{h_i\}_{i=1}^{nb}$算出来存好,而$\{p_{i,j}\}_{i,j=1}^{nb}$可以基于$\{h_i\}_{i=1}^{nb}$算出来。
最后还要提醒的是,上述流程只是在优化时等效于大batch_size模型,也就是说$\langle h_i, \tilde{h}_i\rangle$的梯度等效于原loss的梯度,但是它的值并不等于原loss的值,因此不能用$\langle h_i, \tilde{h}_i\rangle$作为loss来评价模型,它未必是单调的,也未必是非负的,跟原来的loss也不具有严格的相关性。
问题 #
上述流程有着跟《节省显存的重计算技巧也有了Keras版了》介绍的“重计算”一样的问题,那就是跟Dropout并不兼容,这是因为每次更新都涉及到了多次前向计算,每次前向计算都有不一样的Dropout,这意味着我们计算标签向量$\tilde{h}_i$时所用的$h_i$跟计算梯度时所用的$h_i$并不是同一个,导致计算出来的梯度并非最合理的梯度。
这没有什么好的解决方案,最简单有效的方法就是在模型中去掉Dropout。这对于CV来说没啥大问题,因为CV的模型基本也不见Dropout了;对于NLP来说,第一反应能想到的结果就是SimCSE没法用梯度累积,因为Dropout是SimCSE的基础~
小结 #
本文分析了对比学习的梯度累积方法,结果显示对比学习也可以用梯度累积的,只不过多了一次前向计算,并且需要在模型中去掉Dropout。本文同样的思路还可以分析BN如何使用梯度累积,有兴趣的读者不妨试试。
转载到请包括本文地址:https://kexue.fm/archives/8471
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Jun. 17, 2021). 《对比学习可以使用梯度累积吗? 》[Blog post]. Retrieved from https://kexue.fm/archives/8471
@online{kexuefm-8471,
title={对比学习可以使用梯度累积吗?},
author={苏剑林},
year={2021},
month={Jun},
url={\url{https://kexue.fm/archives/8471}},
}
June 18th, 2021
苏神,请问通过multi gpu的 dataparallel 让batch size变大的梯度和单gpu的相同batch size也是不一样的吗?这对对比学习的训练会有不好的影响吗?
多gpu得看你怎么实现,一般框架自带的多gpu也是依赖于独立同分布假设的,不能直接用于对比学习或者BN的模型。
感谢苏神,之前用多gpu并行直接叠加loss发现效果并没有变好,当时就没太想清楚,这篇博客算是找到问题了!
另外想再提一个问题,类似Momentum Contrast for Unsupervised Visual Representation Learning里实现,本质上也是一种增大batch_size的方式,请问增大batch_size在对比学习方面一定会有收益吗?苏神有没有其他的经验?
后面的问题我真的完全没有经验了,抱歉~
June 26th, 2021
是不是只要模型中有dropout就不能这么用了呢,比如BERT
准确来说,有dropout之后,这样累积的梯度不等价于一次性大batch_size,但究竟有没有效果(比如有可能比小batch_size好,比一次性大batch_size差),还是得靠实验来说话。
分享下我这边的实验结论哈,在跨模态的对比学习场景里,模型是BERT+ViT,dropout确实会影响效果,解决方案是,在第一次forward时候每个step人为设置一个rng seed (比如step num或者time之类的),然后记下来,在第二次forward时reset这个seed,这样就能使得总体梯度严格相等(打断点check过梯度误差非常小),整体训练结果也是可以严格对齐的,最终grad accu策略耗时相对原生大batch增加会在20%-30%左右。
此外,也试验过grad checkpointing技术,通过缩小显存占用量来实现增大batch size,耗时略微会比grad accu低一些(大约增加20%左右,可能是因为grad ckpt是在一个step内过完整个batch,所以并行度会略高于multi step的grad accu)。我觉得对于infonce这种loss,两者都是一个可行的解决方案,其中grad accu相对grad ckpt稍微慢些,但是可以支持任意大的batchsize(我们的场景里会用到16384这种batchsize),如果batchsize不是特别大的情况下,也可以通过grad ckpt方案来缩小显存实现把大batch塞进当前计算资源里。
感谢分享。
July 1st, 2021
苏神太强了,说说我的一点经验哈,我在复现simcse的时候,为了实现原文batch_size=512的效果。我是把512的batch拆分成8个batch_size=64的batch,过8次forward,得到512个向量。以同样的方式再过一遍原样本作为对比,得到新的512个向量。然后再用这512x2个向量计算一个simcse的loss,然后再backward求梯度更新。 这样操作其实等价官方代码里用多gpu实现batch_size=512。
这应该叫“forward拆分”吧,这种方式在pytorch这类动态图框架中比较好实现。
你这个我有点不理解原理。
按照我的理解,forward出来的向量是可以存起来,但它已经相当于常量了,应该是没有梯度了,怎么还可以累积几部分再算loss然后backward?除非它每次forward的时候自动将每个输入对每个参数的梯度都缓存起来,但这空间成本未免太大了。
大概理解了,应该是每次forward都建立一个子图,然后拼成一个大图,大图backward的时候,每个子图是串行计算的,所以不会爆显存。
请问苏神这一层说的方法是如何实现的?我还是没能很好理解。如果不清除计算图的话8次64bsz的forward不就相当于一次512bsz的forward吗?这样还是会爆显存的呀,但是如果清除计算图的话就没法对比学习了吧。
pytorch我也不了解,但理论上只要框架愿意,backward也可以串行的。
@刘伟杰|comment-16797
尝试用了一下您的方法来复现simcse,好像无法节省内存;可以麻烦您再详细分享一下吗
November 11th, 2021
good idea,但是,有个小小的问题,公式5是不是不等价的呢?因为pij我觉得是不等于pji的,pij的分母是sum on exp of hi和所有hk的内积,而pji的分母是sum on exp of hj和所有的hk的内积,而这两个分母应该是不一样的,即使分子都是hi和hj的内积,但是他们的分母不同会导致pij != pji。我理解公式5成立的条件是pij等于pji才行?
谢谢指出。想了一下,你是对的,已修正~
June 24th, 2022
还有个问题是对于存在BN的模型,两次前向的统计信息不同,需要以第一次统计信息为准~
October 1st, 2022
[...]对比学习可以使用梯度累积吗? – 科学空间|Scientific Spaces[...]
November 9th, 2022
一共需要两轮,第一轮batch是nb,为了计算h~i,不用保存中间结果;第二轮是分n次forward+backward,为了累计梯度。不知理解是否正确?
对