区分真实样本
前面的两种是为了去估计配分函数，接下来要介绍的 InfoNCE 虽然带个 NCE，但这个的目的不是要预估配分函数，他是直接像上篇应用 NCE 的方法一样，直接采用自归一化，即：
$$ p(\boldsymbol{w}|\boldsymbol{c}) = \frac{u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})}{1} $$
首先对于每个训练样本来说，我们再从噪声分布$q(x)$中采样$k-1$个样本，我们的目的即是区分出真实样本：
$$ \begin{align} L(\boldsymbol{\theta}) & = -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}; \boldsymbol{w}_{j} \sim q} \log \frac{p(\boldsymbol{w}|\boldsymbol{c})}{p(\boldsymbol{w}|\boldsymbol{c})+\sum_{j=1}^{k-1}q(\boldsymbol{w}_{j})}\\ & = -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D};\boldsymbol{w}_{j} \sim q} \log \frac{u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})}{u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c}) + \sum_{j=1}^{k-1} q(\boldsymbol{w}_{j})} \end{align} $$
最小化上述 loss 到底有什么好处呢？
互信息来访
对于训练样本$i$来说，它是真实样本的概率为：
$$ \begin{align} p(d=i) & = \frac{p(\boldsymbol{w}_{i}|\boldsymbol{c})\prod_{j=1, i\neq j}^{k}q(\boldsymbol{w}_{j})}{\sum_{j=1}^{k}p(\boldsymbol{w}_{j}|\boldsymbol{c})\prod_{l=1,l\neq j}^{k}q(\boldsymbol{w}_{l})} \\ &= \frac{\frac{p(\boldsymbol{w}_{i}|\boldsymbol{c})}{q(\boldsymbol{w}_{i})}}{\sum_{j=1}^{k} \frac{p(\boldsymbol{w}_{j}|\boldsymbol{c})}{q(\boldsymbol{w}_{j})}} \end{align} $$
转换的方法即是上下同除以$\prod_{i=1}^{k} q(\boldsymbol{w}_{i})$
InfoNCE 的目的即为正确区分出真实样本，而真实样本都来自训练集（与噪声集区分），也就是说对于 loss 而言，最优解即为$u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})$正比于 density ratio：
$$ u_{\boldsymbol{\theta}}(\boldsymbol{w}_{i}, \boldsymbol{c}) \propto \frac{p(\boldsymbol{w}_{i}|\boldsymbol{c})}{q(\boldsymbol{w}_{i})} $$
我们知道，互信息的计算如下：
$$ I(\boldsymbol{w}, \boldsymbol{c}) = \text{KL}(p(\boldsymbol{w}, \boldsymbol{c})\|p(\boldsymbol{w})p(\boldsymbol{c})) = \mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}} \log \frac{p(\boldsymbol{w}, \boldsymbol{c})}{p(\boldsymbol{w})p(\boldsymbol{c})} $$
为了和上述 loss 有关，很自然做出以下变化：
$$ I(\boldsymbol{w}, \boldsymbol{c}) = \mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}} \log \frac{p(\boldsymbol{w}|\boldsymbol{c})p(\boldsymbol{c})}{p(\boldsymbol{w})p(\boldsymbol{c})} = \mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}} \log {\color{#337dff}\frac{p(\boldsymbol{w}|\boldsymbol{c})}{p(\boldsymbol{w})}} $$
也就是说如果我们要$u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})$来表示互信息，还得$q(x) = p(x)$，也就是噪声分布即为模型分布（unconditioned）
那么：
$$ u_{\boldsymbol{\theta}}(\boldsymbol{w}_{i}, \boldsymbol{c}) \propto \frac{p(\boldsymbol{w}_{i}|\boldsymbol{c})}{{\color{#337dff}p(\boldsymbol{w}_{i})}} $$
此时完整的 loss 即为：
$$ L(\boldsymbol{\theta }) = -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}; \boldsymbol{w}_{j} \sim q} \log \frac{u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})}{\sum_{j=1}^{k} u_{\boldsymbol{\theta}}(\boldsymbol{w}_{j}, \boldsymbol{c})} $$
当我们要使得训练样本$\boldsymbol{w}$是真实样本的概率变得最大，即一定程度最大化$u_{\boldsymbol{\theta}}(\boldsymbol{w}, \boldsymbol{c})$，然后最小化$u_{\boldsymbol{\theta}}(\boldsymbol{w}_{j}, \boldsymbol{c})$，即最大化$\boldsymbol{w}, \boldsymbol{c}$之间的互信息，最小化负样本与$\boldsymbol{c}$之间的互信息
注意：现在用的时候并非像一开始通过真实样本的概念，而是可以将$\boldsymbol{w}, \boldsymbol{c}$看成是一对正样本，噪声样本看作负样本
Hardness-Aware
同时 InfoNCE 会着重关照那些跟真实样本$\boldsymbol{c}$更相似的负样本，这类也叫做「Hard-Negative」，而相应地会轻视一些与$\boldsymbol{c}$本来就没有那么相似的负样本，即「Hardness-Aware」
$$ \begin{align} L(\boldsymbol{\theta }) & = - \mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D};\boldsymbol{w}_{j}\sim q} \log \frac{u_{\boldsymbol{\theta }}(\boldsymbol{w}, \boldsymbol{c})}{u_{\boldsymbol{\theta }}(\boldsymbol{w}, \boldsymbol{c})+\sum_{j=1}^{k-1} u_{\boldsymbol{\theta }}(\boldsymbol{w}_{j}, \boldsymbol{c})} \\ &= -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D};\boldsymbol{w}_{j}\sim q} \log \frac{\exp(s_{\boldsymbol{w}, \boldsymbol{c}} )}{\exp(s_{\boldsymbol{w}, \boldsymbol{c}})+\sum_{j=1}^{k-1}\exp(s_{\boldsymbol{w}_{j}, \boldsymbol{c}})} \\ &= -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D};\boldsymbol{w}_{j} \sim q} \left[s_{\boldsymbol{w}, \boldsymbol{c}} - \log\left( \exp(s_{\boldsymbol{w}, \boldsymbol{c}})+\sum_{j=1}^{k-1}\exp(s_{\boldsymbol{w}_{j},\boldsymbol{c}}) \right)\right] \end{align} $$
对于$\boldsymbol{w}_{j}$而言，与$\boldsymbol{c}$的相似性越大，则在损失函数中占比越大，可以看成是惩罚越大，然而仅凭这个没办法做到上述性质，比如下面的 loss 也可以：
$$ L_{\text{simple}}(\boldsymbol{\theta }) = \mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}; \boldsymbol{w}_{j} \sim q} \left[-s_{\boldsymbol{w}, \boldsymbol{c}} + \sum_{j=1}^{k-1} s_{\boldsymbol{w}_{j}, \boldsymbol{c}}\right] $$
然而实验证明，这个 loss 效果并不是很好
接着，我们来求下损失函数相对于$s_{\boldsymbol{w}_{j}, \boldsymbol{c}}$的偏导：
不妨就对单个样本$\boldsymbol{w}, \boldsymbol{c}$来看：
$$ \frac{ \partial L_{\text{simple}} }{ \partial s_{\boldsymbol{w}_{j}, \boldsymbol{c}} } = 1 $$
也就是说 loss 对于所有$s_{\boldsymbol{w}_{j}, \boldsymbol{c}}$的变化一样敏感，即「一视同仁」，没有给损失函数的优化注入其他信息
$$ \frac{ \partial L }{ \partial s_{\boldsymbol{w}_{j},\boldsymbol{c}} } = \frac{\exp(s_{\boldsymbol{w}_{j}, \boldsymbol{c}})}{\exp(s_{\boldsymbol{w},\boldsymbol{c}} )+ \sum_{j=1}^{k-1}\exp(s_{\boldsymbol{w}_{j}, c})} $$
而对于 InfoNCE 而言，对于所有负样本而言，分母是固定的，而「分子越大」，即$s_{\boldsymbol{w}_{j}, \boldsymbol{c}}$越大，则此时 loss 对于$s_{\boldsymbol{w}_{j}, \boldsymbol{c}}$的变化则越敏感
总结说就是，InfoNCE 这种 loss 在优化时会注入一种偏好，更偏向于把与真实样本相似的噪声样本给区分开来
温度系数
一般现在用的时候，还会加入一个调节因子：$\tau$
$$ L(\boldsymbol{\theta})= -\mathbb{E}_{\boldsymbol{w}, \boldsymbol{c} \sim \mathcal{D}; \boldsymbol{w}_{j} \sim q} \log \frac{\exp(s_{\boldsymbol{w}, \boldsymbol{c}}/\tau )}{\exp(s_{\boldsymbol{w}, \boldsymbol{c}}/\tau)+\sum_{j=1}^{k-1}\exp(s_{\boldsymbol{w}_{j}, \boldsymbol{c}}/\tau)} $$
那么：
$$ \frac{ \partial L }{ \partial s_{\boldsymbol{w}_{j},\boldsymbol{c}} } = \frac{1}{\tau}\frac{\exp(s_{\boldsymbol{w}_{j}, \boldsymbol{c}}/\tau)}{\exp(s_{\boldsymbol{w},\boldsymbol{c}}/\tau )+ \sum_{j=1}^{k-1}\exp(s_{\boldsymbol{w}_{j}, c}/\tau)} $$
推论 1：当$\tau$很小时，比如$0.02$那么由上式可得 loss 会格外关注于将与真实样本相似的噪声样本给区分开。换句话说，$\tau$越小，对 hard-negative 的惩罚越大
下面我们引入一个 kernel 来衡量一个分布的均匀性（uniform），$f(\boldsymbol{x})$代表模型输出的表示
$$ L_{\text{uniformity}}(f;t) = \log \mathbb{E}_{\boldsymbol{x}, \boldsymbol{y} \sim \mathcal{D}} \left[e^{-t\|f(\boldsymbol{x}) - f(\boldsymbol{y})\|_{2}^{2}}\right](t \in \mathbb{R}_{+}) $$
作者采用不同的$\tau$来对比模型产生表示的均匀性
当$\tau$比较小时，模型产生的表征空间更加均匀，随着$\tau$的增大，均匀性被破坏
而对于一个好的对比学习的表征空间应该满足「locally clustered，globally uniform」，当表征均匀分布在球面上时，此时用一个超平面（线性分类器）便可分开，比如下图
推论 2：当$\tau$较小时，用 infoNCE loss 训出的模型表征空间比较均匀，当$\tau$增大时，表征空间的均匀性被逐步破坏
上代码
当下的算力已经足够支撑直接用概率而非自归一化了，其实会发现 InfoNCE loss 跟多分类交叉熵损失一样：
import torch.nn.functional as F
def infoNCE(q_vectors, c_vectors, labels, tau=1):
"""Compute the InfoNCE Loss in
http://arxiv.org/abs/1807.03748.
Params:
- q_vectors: Tensor. Shape: (bs, d_model). Vector presentation of query.
- c_vectors: Tensor. Shape: (num_pos_neg_contexts, d_model). Vector presentation of context.
- labels: Tensor. Shape: (bs, ). Indices for the positive contexts.
- tau: Scalar. The parameter to control penalty on hard negatives (smaller => harder).
A good initialization value can be 0.02 or 0.05.
Examples:
>>> q_vectors = torch.randn((2, 4))
>>> c_vectors = torch.randn((4, 4)) # each query has 4 contexts
>>> labels = torch.tensor([1, 3])
>>> print(infoNCE(q_vectors, c_vectors, labels))
"""
scores = (q_vectors @ c_vectors.T) / tau
return F.cross_entropy(scores, labels)