在之前的文章《用时间换取效果: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}},
}