超大规模分类（一）：噪声对比估计（Noise Contrastive Estimation, NCE）

NCE损失对应的论文为《A fast and simple algorithm for training neural probabilistic language models》，发表于2012年的ICML会议。

背景

在2012年，语言模型一般采用n-gram的方法，统计单词/上下文间的共现关系，比神经概率语言模型（neural probabilistic language models, NPLMs）效果好。
现在主流的语言模型都是神经概率语言模 型，核心思想是已知上下文$h$，预测下一个词为$w$的概率，通过一定的解码方法（例如greedy search、beam search等），对概率做解码，得到下一个词。Greedy search可以理解为选择概率最大的那个词。
2012年神经概率语言模型效果不好的原因是难训练。一方面自然是硬件的制约，那一年英伟达刚发布GTX680，和现在的A100、H100完全没法比。当时老黄不给力，学术界也没办法；另一方面是算法效率不行，难以进行大规模的分类学习，将”已知上下文$h$，预测下一个词为$w_i$的概率“建模成分类学习任务，目的在于把下一个词分类到词表中的某个词上。
举个例子，已知上下文是“我想去”，需要预测下一个词。词表中有4个词，即['北京'，'上海','天津','广州']，需要把下一个词归类到词表的4个词里。如果词表有10万个词呢？训不动啊~
这就是当时面临的困境。NCE对分类算法做了优化，使得对大词表做分类任务成为可能。

原理

通俗的背景讲完了，接下来谈谈公式化的原理部分。

问题建模

已知上下文$h$，预测下一次词为$w$的概率为：
$$P_{\theta }^h(w)=\frac{exp(s_{\theta }(w,h))}{{\sum }_{w_i}exp(s_{\theta }(w_i,h))}\text{\hspace{1em}(1)}$$
其中，$s_{\theta }(w,h)$表示已知上下文$h$，下一个词为$w$的预测得分；${\sum }_{w_i}$表示词表内的所有词。
一般情况下，$s_{\theta }(w,h)$通过对上下文$h$表征以及词类别$w$表征添加多个全连接层计算得到。最简单的策略，仅对上下文$h$表征$f_h$用一个全连接层$W$做一次映射，再和词类别$w$表征$f_{w_i}$做点积即可。
$s_{\theta }(w,h)=(f_{h}W)\cdot f_w$

难度分析

对公式(1)进行分析，
分子部分$exp(s_{\theta }(w,h))$是好算的，针对单个$w$，只需要计算一次。
分母部分KaTeX parse error: \tag works only in display equations不好算，针对单个$w$，需要计算$exp(s_{\theta }(w_1,h)),exp(s_{\theta }(w_2,h)),...exp(s_{\theta }(w_n,h))$，如果词表中词很多，计算量不小。

目前学术界、工业界对超大规模分类的优化基本上都聚焦在如何优化分母上，例如InfoNCE仅关注batch内的负类样本、KNN softmax对类别聚类，减少类别数目、partial FC对类别做采样以及显存均分来较少计算量、Inf-CL借助FlashAttention的思想，以空间换时间。

优化策略

既然对词表内n个词的大规模分类任务难做，难办，那就掀桌子不办了！！！

将原多分类任务转换成一个更容易实现的任务——新二分类任务。
除了有正常的真实数据之外，从一个噪声分布里采样噪声数据，对真实数据和噪声数据做二分类，可以证明：随着噪声数据越多，转换后任务的优化目标和转换前任务越接近。

新二分类任务

给定上下文$h$后，现在有两个数据分布，一个是真实数据分布$P_d^h(w)$（实际应该写成$P_d(w|h)$，简化形式写成$P_d^h(w)$），另一个是噪声数据分布$P_n(w)$，真实数据和噪声数据的比例是1:k。所以，训练数据的完整分布是$P^h(w)=\frac{1}{k+1}{P}_d^h(w)+\frac{k}{k+1}{P}_n(w)$，训练任务是$D=1$（分辨真实数据）和$D=0$（分辨噪声数据）。
我们希望优化神经网络参数$\theta$，来拟合真实数据分布$P_d^h(w)=P_{\theta }^h(w)$，后者就是我们学到的数据分布$P_{\theta }^h(w)$，于是，训练数据的完整分布写成$P^h(w,\theta )=\frac{1}{k+1}{P}_{\theta }^h(w)+\frac{k}{k+1}{P}_n(w)$

训练目标一般是最大化后验概率$P^h(D|w,\theta )$的对数似然期望$E\left[log(P^h(D|w,\theta ))\right]$，需要计算后验概率$P^h(D|w,\theta )$。
$$P^h(D|w,\theta )=P^h(D=1|w,\theta )+P^h(D=0|w,\theta )\text{\hspace{1em}(2)}$$
真实数据分布的后验概率为：
$$\begin{array}{c}\begin{array}{rl}{P}^h(D=1|w,\theta )& =\frac{P^h(w,\theta |D=1)}{P^h(w,\theta )}{P}^h(D=1)\\ & =\frac{P_{\theta }^h(w)}{\frac{1}{k+1}{P}_{\theta }^h(w)+\frac{k}{k+1}{P}_n(w)}\frac{1}{k+1}\\ & =\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\end{array}\end{array}\text{\hspace{1em}(3)}$$
我们来看看等式为什么成立

• 边缘概率$P^h(w,\theta )=\frac{1}{k+1}{P}_{\theta }^h(w)+\frac{k}{k+1}{P}_n(w)$
• 先验概率$P^h(D=1)=\frac{1}{k+1}$，原因是真实数据和噪声数据的比例是1:k。
• 似然函数$P^h(w,\theta |D=1)=P_{\theta }^h(w)$，表明在真实数据分布下，从词表里预测下一个词为$w$的概率是$P_{\theta }^h(w)$，这就是我们想拟合的函数。

类似的，噪声数据分布的后验概率为：
$$\begin{array}{c}\begin{array}{rl}{P}^h(D=0|w,\theta )& =\frac{P^h(w,\theta |D=0)}{P^h(w,\theta )}{P}^h(D=0)\\ & =\frac{P_n(w)}{\frac{1}{k+1}{P}_{\theta }^h(w)+\frac{k}{k+1}{P}_n(w)}\frac{k}{k+1}\\ & =\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\end{array}\end{array}\text{\hspace{1em}(4)}$$

后验概率$P^h(D|w_i,\theta )$的对数似然的期望$E\left[log(P^h(D|w_i,\theta ))\right]$为
$$\begin{array}{c}\begin{array}{rl}{J}^h(\theta )& =E\left[log(P^h(D|w,\theta ))\right]\\ & =E_{P_d^h}\left[log{P}^h(D=1|w,\theta )\right]+E_{P_n}\left[log{P}^h(D=0|w,\theta )\right]\\ & =E_{P_d^h}\left[log\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\right]+E_{P_n}\left[log\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\right]\end{array}\end{array}\text{\hspace{1em}(5)}$$
我们来算一下梯度，等于
$$\begin{array}{c}\begin{array}{rl}\frac{\partial }{\partial \theta }{J}^h(\theta )& =E_{P_d^h}\left[\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)\right]-\\ & k{E}_{P_n}\left[\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)\right]\end{array}\end{array}\text{\hspace{1em}(6)}$$

对(6)式做化简，有
$$\begin{array}{c}\begin{array}{rl}\frac{\partial }{\partial \theta }{J}^h(\theta )& =E_{P_d^h}\left[\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w_i)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)\right]-\\ & k{E}_{P_n}\left[\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)\right]\\ & =\sum _w[P_d^h\cdot \frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)-\\ & k{P}_n\cdot \frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)]\\ & =\sum _w[\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\times \\ & (P_d^h(w)-P_{\theta }^h(w))\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)]\end{array}\end{array}\text{\hspace{1em}(7)}$$

当噪声数据量级巨大，$k\to \infty$ ，$\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\to 1$ ，有
$$\begin{array}{c}\begin{array}{rl}\frac{\partial }{\partial \theta }{J}^h(\theta )& =\sum _w[\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\times \\ & (P_d^h(w)-P_{\theta }^h(w))\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)]\\ & \to \sum _w\left[(P_d^h(w)-P_{\theta }^h(w))\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)\right]\end{array}\end{array}\text{\hspace{1em}(8)}$$

原多分类任务

我们计算下原多分类任务的对数似然期望和梯度，看看$k\to \infty$ 时的新二分类任务和原多分类任务有什么关系。原多分类任务的优化目标为
$$\begin{array}{c}\begin{array}{rl}{J}^h(\theta )& =E_{P_d^h}\left[log(P_{\theta }^h(w)\right]\\ & =E_{P_d^h}\left[log\left(\frac{exp(s_{\theta }(w,h))}{{\sum }_{w}exp(s_{\theta }(w,h))}\right)\right]\\ & =E_{P_d^h}\left[s_{\theta }(w,h)\right]-E_{P_d^h}\left[log\left(\sum _{w}exp\left(s_{\theta }(w,h)\right)\right)\right]\\ & =E_{P_d^h}\left[s_{\theta }(w,h)\right]-log\left(\sum _{w}exp\left(s_{\theta }(w,h)\right)\right)\end{array}\end{array}\text{\hspace{1em}(9)}$$
等式最后一步成立的原因是$\left[log\left({\sum }_{w}exp\left(s_{\theta }(w,h)\right)\right)\right]$仅和模型预测分布$P_{\theta }^h$有关，和真实数据分布$P_d^h$无关。
对(9)式求梯度，有$$\begin{array}{c}\begin{array}{rl}\frac{\partial }{\partial \theta }{J}^h(\theta )& =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\frac{\partial }{\partial \theta }log\left(\sum _{w}exp\left(s_{\theta }(w,h)\right)\right)\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\frac{1}{{\sum }_{w}exp\left(s_{\theta }(w,h)\right)}\frac{\partial }{\partial \theta }\sum _{w}exp\left(s_{\theta }(w,h)\right)\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\frac{1}{{\sum }_{w}exp\left(s_{\theta }(w,h)\right)}\sum _w\left(s_{\theta }(w,h)\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right)\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\sum _w\frac{s_{\theta }(w,h)}{{\sum }_{w}exp\left(s_{\theta }(w,h)\right)}\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\sum _w{P}_{\theta }^h(w)\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\right]-\sum _w{P}_{\theta }^h(w)\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\\ & =\sum _w{P}_d^h\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)-\sum _w{P}_{\theta }^h(w)\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\\ & =\sum _w(P_d^h(w)-P_{\theta }^h(w))\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)\end{array}\end{array}\text{\hspace{1em}(10)}$$对比公式(8)和公式(10)，很像，但不一样。公式(8)最后是$\frac{\partial }{\partial \theta }log{P}_{\theta }^h(w)$，公式(10)最后是$\frac{\partial }{\partial \theta }{s}_{\theta }(w,h)$，咋回事？

不一样就对了，在NCE中，我们可以将${\sum }_{w}exp\left(s_{\theta }(w,h)\right)$等价成1，那公式(8)和公式(10)就一样了。那为什么可以等价呢？论文的说辞是：$模型参数较多，把正则项当做常数，公式中其他项，比如s_{\theta }，能学到正则项。$（正则项可以理解为${\sum }_{w}exp\left(s_{\theta }(w,h)\right)$），那么${\sum }_{w}exp\left(s_{\theta }(w,h)\right)$是1也好，100也好，都不会对模型收敛有影响。简单起见，当做1就行。

这段说辞还是太抽象了，有没有形象一点的解释？

两个任务为什么可以等价

原多分类任务

$$\begin{array}{c}\begin{array}{rl}{J}^h(\theta )& =E_{P_d^h}\left[log(P_{\theta }^h(w)\right]\\ & =E_{P_d^h}\left[log\left(\frac{exp(s_{\theta }(w,h))}{{\sum }_{w}exp(s_{\theta }(w,h))}\right)\right]\end{array}\end{array}\text{\hspace{1em}(11)}$$
该任务的对数似然期望见公式(11)，$log$函数曲线如下：

如果$log(P_{\theta }^h(w)=exp(s_{\theta }(w,h))\in [0,+\infty ]$，$J^h(\theta )=E_{P_d^h}\left[log(P_{\theta }^h(w)\right]$不存在极值，无法收敛。
如果对$log(P_{\theta }^h(w)=exp(s_{\theta }(w,h))\in [0,+\infty ]$进行归一化，$log(P_{\theta }^h(w)=\left[log\left(\frac{exp(s_{\theta }(w,h))}{{\sum }_{w}exp(s_{\theta }(w,h))}\right)\right]\in (0,1)$，$J^h(\theta )=E_{P_d^h}\left[log(P_{\theta }^h(w)\right]$存在极值，具备收敛条件。

现二分类任务

从公式(5)可知，
$$\begin{array}{c}\begin{array}{rl}{J}^h(\theta )& =E\left[log(P^h(D|w,\theta ))\right]\\ & =E_{P_d^h}\left[log{P}^h(D=1|w,\theta )\right]+E_{P_n}\left[log{P}^h(D=0|w,\theta )\right]\\ & =E_{P_d^h}\left[log\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\right]+E_{P_n}\left[log\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}\right]\\ & =E_{P_d^h}\left[log(\sigma (\Delta ))\right]+k{E}_{P_n}\left[log(1-\sigma (\Delta ))\right]\end{array}\end{array}\text{\hspace{1em}(12)}$$
，其中$\Delta =log{P}_{\theta }^h(w)-logk{P}_n(w)$，将公式(5)推导成具备$\sigma$的公式(12)，原因在于求导方便，$\frac{\partial }{\partial x}\sigma (x)=\sigma (x)(1-\sigma (x))$，将公式(5)推导成公式(12)的过程是：
$$\begin{array}{c}\begin{array}{rl}\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}& =\frac{1}{1+\frac{k{P}_n(w)}{P_{\theta }^h(w)}}\\ & =\frac{1}{1+exp(log(\frac{k{P}_n(w)}{P_{\theta }^h(w)}))}\\ & =\frac{1}{1+exp(logk{P}_n(w)-log{P}_{\theta }^h(w))}\\ & =\frac{1}{1+exp(-(log{P}_{\theta }^h(w)-logk{P}_n(w)))}\\ & =\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w))\end{array}\end{array}\text{\hspace{1em}(12)}$$
$$\begin{array}{c}\begin{array}{rl}\frac{k{P}_n(w)}{P_{\theta }^h(w)+k{P}_n(w)}& =1-\frac{P_{\theta }^h(w)}{P_{\theta }^h(w)+k{P}_n(w)}\\ & =1-\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w))\end{array}\end{array}\text{\hspace{1em}(13)}$$
于是，计算对数似然均值（公式(12)）对$log{P}_{\theta }^h(w)$的一阶导，有
$$\begin{array}{c}\begin{array}{rl}\frac{\partial J^h(\theta )}{\partial log{P}_{\theta }^h(w)}& =\frac{\partial J^h(\theta )}{\partial \Delta }\frac{\partial \Delta }{\partial log{P}_{\theta }^h(w)}\\ & =\frac{\partial J^h(\theta )}{\partial \Delta }\\ & =\frac{\partial }{\partial \Delta }\left\{E_{P_d^h}\left[log(\sigma (\Delta ))\right]+k{E}_{P_n}\left[log(1-\sigma (\Delta ))\right]\right\}\\ & =E_{P_d^h}\left[\frac{\partial }{\partial \Delta }log(\sigma (\Delta ))\right]+k{E}_{P_n}\left[\frac{\partial }{\partial \Delta }log(1-\sigma (\Delta ))\right]\\ & =E_{P_d^h}\left[1-\sigma (\Delta )\right]+k{E}_{P_n}\left[-\sigma (\Delta )\right]\\ & =\sum _w{P}_{\theta }^h(w)(1-\sigma (\Delta ))-k{P}_n(w)\sigma (\Delta )\end{array}\end{array}\text{\hspace{1em}(14)}$$
如果$P_{\theta }^h(w)=P_d^h(w)$，对数似然均值达到极大值（这个是废话，因为训练目标就是希望$P_{\theta }^h(w)\to P_d^h(w)$，并且在优化策略章节开始部分，我们就让$P_{\theta }^h(w)=P_d^h(w)$）其中$P_d^h(w)$表示真实分布。
我们再计算对数似然均值（公式(12)）对$log{P}_{\theta }^h(w)$的二阶导，有：
$$\begin{array}{c}\begin{array}{rl}\frac{{\partial }^2{J}^h(\theta )}{\partial lo{g}^2{P}_{\theta }^h(w)}& =\frac{{\partial }^{2}J(\theta )}{\partial {\Delta }^2}\\ & =\frac{\partial }{\partial \Delta }\left\{E_{P_d^h}\left[1-\sigma (\Delta )\right]+k{E}_{P_n}\left[-\sigma (\Delta )\right]\right\}\\ & =E_{P_d^h}\frac{\partial }{\partial \Delta }\left[1-\sigma (\Delta )\right]+k{E}_{P_n}\frac{\partial }{\partial \Delta }\left[-\sigma (\Delta )\right]\\ & =E_{P_d^h}[-\sigma (\Delta )(1-\sigma (\Delta ))]+k{E}_{P_n}[-\sigma (\Delta )(1-\sigma (\Delta ))]\end{array}\end{array}\text{\hspace{1em}(14)}$$
因为$[-\sigma (\Delta )(1-\sigma (\Delta ))]$始终小于0，所以二阶导始终小于0，说明新二分类任务的对数似然均值是关于$log{P}_{\theta }^h(w)$的凸函数，有唯一极大值。所以极大值一定是$P_{\theta }^h(w)=P^h(w)$。
最重要的是，整个推导过程对是否需要归一化没有要求，既然没有要求，直接让${\sum }_{w}exp\left(s_{\theta }(w,h)\right)=1$

代码实现

从公式(12)，我们可以知道：$\Delta =log{P}_{\theta }^h(w)-logk{P}_n(w)$
$$\begin{array}{c}\begin{array}{rl}{J}^h(\theta )& =E\left[log(P^h(D|w,\theta ))\right]\\ & =E_{P_d^h}\left[log\sigma (\Delta )\right]+k{E}_{P_n}\left[log(1-\sigma (\Delta ))\right]\\ & =E_{P_d^h}\left[log\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w))\right]+\\ & k{E}_{P_n}\left[log(1-\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w)))\right]\\ & =\sum _w\left\{P_d^h\left[log\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w))\right]\right\}+\\ & k\sum _w\left\{P_n\left[log(1-\sigma (log{P}_{\theta }^h(w)-logk{P}_n(w)))\right]\right\}\\ & \to log(\sigma (log{P}_{\theta }^h(w_0)-logk{P}_n(w_0))+\\ & \sum _{i=1}^k\left[log(1-\sigma (log{P}_{\theta }^h(w_i)-logk{P}_n(w_i)))\right]\\ & =log(\sigma (s_{\theta }(w_0,h)-logk{P}_n(w_0))+\\ & \sum _{i=1}^k\left[log(1-\sigma (s_{\theta }(w_i,h)-logk{P}_n(w_i)))\right]\end{array}\end{array}\text{\hspace{1em}(15)}$$
具体实现时，正样本项仅考虑目标class，负样本项随机选择k个样本，通过蒙特卡洛来模拟抽样。
那最终损失函数代码应该怎么写呢？

$$\begin{array}{c}\begin{array}{rl}loss& =-J^h(\theta )\\ & =-log(\sigma (s_{\theta }(w_0,h)-logk{P}_n(w_0)))-\\ & \sum _{i=1}^k\left[log(1-\sigma (s_{\theta }(w_i,h)-logk{P}_n(w_i)))\right]\end{array}\end{array}\text{\hspace{1em}(16)}$$

公式(16)中有四个项输入，分别是

• $s_{\theta }(w_0,h)$，目标class的logit
• $P_n(w_0)$，目标class的噪声分布
• $s_{\theta }(w_i,h)$，噪声class的logit
• $P_n(w_i)$，噪声class的噪声分布

from torch import randn, tensor, log, multinomial
import torch.nn.functional as F
from einops import repeat
import torch
import mathbs,k=2,8
num_classes=16#构造噪声：按照类别的频率采样
#（噪声分布约等于实际数据分布，两个分布越接近，nce效果越好）
classes=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
class_freq=tensor([20,10,30,5,45,56,76,43,23,11,34,5,6,54,23,7])
class_probs=class_freq/class_freq.sum()
noise_classes=multinomial(class_probs, num_classes)#模型预测的logits
logits=randn(bs, num_classes)
#2个样本的标签
labels=tensor([2, 4])#目标class的logit
true_class_logits=logits.take_along_dim(labels[:, None], dim=1)#目标class的噪声分布
true_class_noise=class_probs[labels]
#噪声class的logit
logits_k = repeat(logits, '(b 1) h -> (b k) h', k=k)
noise_class_logits = logits_k.take_along_dim(noise_classes.reshape(bs * k, -1), dim=1)
#噪声class的噪声分布
noise_class_noise=class_probs[noise_classes]#nce loss计算
true_class_loss = -torch.log( F.sigmoid(true_class_logits - torch.log(k*true_class_noise))).mean()
noise_class_loss = -torch.log( 1-F.sigmoid(noise_class_logits - torch.log(k*noise_class_noise))).mean()loss = true_class_loss+noise_class_loss
print("nce loss is {:.4f}".format(loss))