论文信息: ICLR 2020 (Spotlight); 目前引用量: 329
0. Main Takeaway
通过加入equipartition constraint, 本文将simultaneous clustering and representation learning问题转化为了一个最优输运问题, 并可以通过Sinkhorn-Knopp algorithm算法高效求解. 同时, 作者指出该objective可以被理解为最大化labels和input data indices之间的互信息, 从而可以避免传统算法可能会产生的degenerate solution (即所有data indices被分到同一个label上).
0.1. Math Background
在开始前, 本节列出了有助于理解之后推导的若干数学公式, 如果只想大概理解论文的intuition可以跳过本节.
关于Optimal Transport和Sinkhorn-Knopp algorithm可以参考我之前的文章: Optimal Transport入门简述
关于互信息和熵:
假设我们有随机变量$X,Y$, 联合分布为$p(x,y)$, 我们记$X$的熵为$H(X)$.
则$X$, $Y$之间的互信息为:
$$ \begin{aligned} I(X,Y) &= \int_{x,y} p(x,y) \log\frac{p(x,y)}{p(x)p(y)} \\ &= H(X) - H(X|Y)\quad \text{where } H(X) = -\int_x p(x)\log p(x)\\ &= H(Y) - H(Y|X)\\ \end{aligned} $$
所以$X,Y$之间的互信息可以被理解为在知道了$Y$之后, $X$的熵减. ($H(X)-H(X|Y)$).
$X,Y$的联合分布的熵为:
$$ \begin{aligned} H(X,Y) &= -\int_{x,y} p(x,y)\log p(x,y) \\ &= H(X) + H(Y|X)\\ &= H(X) + H(Y) - I(X,Y) \quad \text{where } I(X,Y) = H(Y)-H(Y|X) \end{aligned} $$
区分: 如果有随机变量$X$, 它有两个概率分布$p(X)$和$q(X)$, 则$p$和$q$之间的cross-entropy为:
$$ H(p,q) = -\int_x p(x) \log q (x) $$
1. Introduction
在CV领域, representation learning一直是一个非常热门的话题, 该任务的目标是学习高维图像输入的有效低维表征, 从而提升模型在下游任务上的表现. 之前的研究表明, 当supervised data充足时, 在ImageNet等带标记数据集上进行classification训练就足以学到较好的抽取表征的pre-trained model. 但是当我们没有足够多带标记数据时, 我们便需要考虑其他方法.
其中, simultaneous clustering and representation learning 是近些年来比较流行的pipeline, 例如DeepCluster. 在面对无标记数据时, 此类算法包括两个阶段, 第一阶段是用cluster算法将数据点分类 (相当于给了label); 第二部分在前面分类好的数据集上进行classification, 从而最终实现representation learning的目的.
但是诸如DeepCluster等算法都存在一个问题, 就是它们可能会产生degenerated solution, 即如果阶段一将所有数据点$\boldsymbol{x}$都分到同一类, 第二阶段的分类任务也只需要输出同样的constant label即可, 这样的一个解虽然也是局部最优解 (loss都将为了0), 但它肯定不是我们想要的.
因此, 本文通过增加了一个equipartition constraint (即希望所有数据点被均匀地map到$K$个类上), 将传统simultaneous clustering and representation learning的loss转化为了一个Optimal Transport问题, 并说明了该objective可以被理解为是在最大化labels和input data indices之间的互信息, 而且可以避免产生上述degenerate solution.
2. Methodology
2.1. Self-labelling as Optimal Transport
在supervised-classification任务中: 假设我们有$N$张图像$\left \{ I_i\right \}$, 每张图像对应$K$个label $\left \{ y_1, \cdots, y_K\right \}$中的一个. 我们通过映射$\Phi$将其映射到特征空间上: $\boldsymbol{x} = \Phi(I)$, $\boldsymbol{x}\in\mathbb{R}^D$. 之后使用映射 (classification head) 在特征空间上完成分类任务 $h: \mathbb{R}^D\mapsto \mathbb{R}^K$.
总的来说, 我们的分类器可以被写为一个复合函数:
$$ p\left(y=\cdot \mid I_{i}\right) = p\left(y=\cdot \mid \boldsymbol{x}_{i}\right)=\operatorname{softmax}\left(h \circ \Phi\left(I_{i}\right)\right) $$
而我们通过最小化cross-entropy loss来jointly优化特征映射$\Phi$和分类器$h$:
$$ E\left(p \mid y_{1}, \ldots, y_{N}\right)=-\frac{1}{N} \sum_{i=1}^{N} \log p\left(y_{i} \mid \boldsymbol{x}_{i}\right) \tag{1} $$
不难看出, 优化式$(1)$需要带标签的数据集, 而当标签不可得时, 我们需要采用self-labelling的mechanism来获取标签 $q(y|\boldsymbol{x}_i)$, 例如使用clustering算法. 则式$(1)$变成:
$$ E(p, q)=-\frac{1}{N} \sum_{i=1}^{N} \sum_{y=1}^{K} q\left(y \mid \boldsymbol{x}_{i}\right) \log p\left(y \mid \boldsymbol{x}_{i}\right) $$
注意到现在loss $E(p,q)$除了和分类结果$p$有关, 还和我们self-labelling的结果$q$有关. 而对于带标记数据集, 相当于$q(y|\boldsymbol{x}_i) = \delta(y-y_i)$. 如果我们同时优化$p$和$q$, 就容易陷入前文所提到的degenerate solution. 所以我们给$q$加上一个equipartition constraint, 即希望数据点被尽可能均匀地分到这$K$个类中.
$$ \min _{p, q} E(p, q) \quad \text { subject to } \quad \forall y: q\left(y \mid \boldsymbol{x}_{i}\right) \in\{0,1\} \text { and } \sum_{i=1}^{N} q\left(y \mid \boldsymbol{x}_{i}\right)=\frac{N}{K} \text {. } $$
我们加入的限制$\sum\limits_{i=1}^{N} q\left(y \mid \boldsymbol{x}_{i}\right)=\frac{N}{K}$其实可以被理解为加入了一条"先验", 即我们在不知道数据标签的真实分布的情况下, 我们假设标签分布是均匀的. 此外我们不难发现, 在固定了$p$时, 单独优化$q$是一个$0-1$整数规划问题, 当样本点数量很大时求解会十分耗时.
我们记$P\in \mathbb{R}^{K\times N}$, $Q \in \mathbb{R}^{K\times N}$, 且$P_{yi} = p(y|\boldsymbol{x}_i)\frac{1}{N}$, $Q_{yi} = q(y|\boldsymbol{x}_i)\frac{1}{N}$ (乘$\frac{1}{N}$是为了让$P$和$Q$变成probability矩阵, 即$\sum\limits_{p\in P}p = 1$). 同时我们记所有满足equipartition constraint的$Q$的可行解空间为$U(r,c)$, 即:
$$ U(r, c):=\left\{Q \in \mathbb{R}_{+}^{K \times N} \mid Q \mathbf{1}_N=r, Q^{\top} \mathbf{1}_K=c\right\} $$
其中$r = \frac{1}{K} \mathbf{1}_K$, $c=\frac{1}{N}\mathbf{1}_N$. 且我们注意到:
$$ \begin{aligned} E(p,q) &= -\frac{1}{N} \sum_{i=1}^{N} \sum_{y=1}^{K} q\left(y \mid \boldsymbol{x}_{i}\right) \log p\left(y \mid \boldsymbol{x}_{i}\right)\\ &= - \sum_{i=1}^{N} \sum_{y=1}^{K} \frac{1}{N}q\left(y \mid \boldsymbol{x}_{i}\right) \log \frac{1}{N}p\left(y \mid \boldsymbol{x}_{i}\right)\times N\\ &= - \sum_{i=1}^{N} \sum_{y=1}^{K} \frac{1}{N}q\left(y \mid \boldsymbol{x}_{i}\right) \log \frac{1}{N}p\left(y \mid \boldsymbol{x}_{i}\right)- \sum_{i=1}^{N} \sum_{y=1}^{K} \frac{1}{N}q\left(y \mid \boldsymbol{x}_{i}\right) \log N\\ &= \langle Q, -\log P\rangle - \log N \end{aligned} $$
其中$\langle \cdot, \cdot \rangle$为Frobenius dot-product, 即矩阵间element-wise乘法. 综上, $(1)$式可以转化为:
$$ \min _{Q \in U(r, c)}\langle Q,-\log P\rangle \tag{2} $$
这便是一个标准的最优输运问题, 其中$-\log P$作为最优输运问题中的代价矩阵, 在优化过程中是固定的. 为了让问题更易于求解, 我们加入正则项$\mathrm{KL}(Q\Vert rc^{\top})$, 从而将$(2)$式中的Wasserstein distance转化为Sinkhorn distance, 从而之后可以通过Sinkhorn-Knopp算法迭代求解:
$$ \min _{Q \in U(r, c)}\langle Q,-\log P\rangle+\frac{1}{\lambda} \mathrm{KL}\left(Q | r c^{\top}\right)\tag{3} $$
而对于$(3)$式, Sinkhorn-Knopp算法证明了它的最优解满足如下形式 (证明参考link):
$$ Q=\operatorname{diag}(\alpha) P^{\lambda} \operatorname{diag}(\beta) $$
其中$P^\lambda$也为$\mathbb{R}^{K\times N}$的矩阵, 且满足$(P^\lambda)_{ij} = (P_{ij})^\lambda$. 而$\alpha\in \mathbb{R}^K$, $\beta \in\mathbb{R}^N$. 之后$\alpha$和$\beta$可以通过迭代求解, 最终目标是使得$Q\in U(r,c)$:
$$ \forall y: \alpha_{y} \leftarrow\left[P^{\lambda} \beta\right]_{y}^{-1} \quad \forall i: \beta_{i} \leftarrow\left[\alpha^{\top} P^{\lambda}\right]_{i}^{-1} $$
注意这里的迭代式是当$r$和$c$都为"均匀向量" ($\forall i,j, r_j = r_j, c_i=c_j$)时的迭代式, 并非通解形式.
在迭代优化$\alpha$, $\beta$时, 一次更新的计算复杂度大约是$\mathcal{O}(KN)$的, 实际计算过程中, 在ImageNet上使用GPU单次计算收敛约只需要2分钟. 而且再下一轮迭代时还可以利用上一轮收敛时的$\alpha$, $\beta$值, 从而达到warm-start的效果.
总结: Self-labelling as Optimal Transport 算法主要包括两个阶段:
针对优化目标 $E(p,q) + \log N = \langle Q,-\log P\rangle$
阶段一 representation learning: 固定label assignments $Q$. 优化分类模型$P_{y i}=\operatorname{softmax}_{y}\left(h \circ \Phi\left(I_{i}\right)\right)$, 优化目标是让cross-entropy尽可能小.
阶段二 self-labelling: 固定分类模型$P$, 通过Sinkhorn-Knopp算法优化label assignments $Q$, 优化目标是在保证标签均匀的前提下, cross-entropy尽可能小.
2.2. Information Theoretical Interpretation of Self-labelling as Optimal Transport
我们定义一个新的随机变量$i$, $i$代表输入数据的下标, 它满足$q(i) = p(i) = \frac{1}{N}$, $p(\boldsymbol{x}|i) = \delta(\boldsymbol{x}- \boldsymbol{x}_i)$. 所以我们有$p\left(y \mid \boldsymbol{x}_{i}\right)=p(y \mid i)$; $q\left(y \mid \boldsymbol{x}_{i}\right)=q(y \mid i)$. 不难发现$P_{yi} = p(y|\boldsymbol{x}_i)\frac{1}{N} = p(y|i) p(i) = p(y,i)$, $Q_{yi} = q(y|\boldsymbol{x}_i)\frac{1}{N}= q(y|i)q(i) = q(y,i)$. 所以我们有:
$$ E(p, q)+\log N= \langle Q,-\log P\rangle =-\sum_{i=1}^{N} \sum_{y=1}^{K} q(y, i) \log p(y, i)=H(q, p) $$
即我们的优化目标其实是联合分布$q(y,i)$和$p(y,i)$之间的cross-entropy, 因此当$p(y,i)=q(y,i)$时, cross-entropy取得极小值, 我们记为$H_q(y,i)$, 即$q(y,i)$的熵:
$$ \min_p H(q, p) = H_q(y,i) = H_{q}(y)+H_{q}(i)-I_{q}(y, i)= \text{const.} -I_{q}(y, i) $$
注意我们假定了$q(i) = p(i) = \frac{1}{N}$, 且equipartition constraint规定了$q(y) = \frac{1}{K}$, 故$H_{q}(y)+H_{q}(i)$都是常数. 所以Self-labelling as Optimal Transport在阶段二中, 在保证了equipartition constraint下去优化$E(p,q) + \log N = \langle Q,-\log P\rangle$的行为, 可以被理解为是: 在显式地保证了equipartition constraint的情况下去最大化标签$y$和输入数据下标$i$之间的互信息.
2.3. Why DeepCluster Avoid Degenerate Solution?
为什么DeepCluster等算法不会产生作者address的Degenerate Solution (即所有样本都被分到一类)的情况呢?
- 当classification model $P$固定, 优化label assignments $Q$时, 由于DeepCluster采取k-means, 所以它无法把不同的特征向量$\boldsymbol{x}_i$拉到一起, 所以只能给他不同的label以降低聚类loss.
- 当label assignments $Q$固定, 优化classification model $P$时, 由于label已经给定, 模型只能调整$\Phi$让具有不同label的图像的特征向量尽可能远离.
但总的来说, 本文相当于给了simultaneous clustering and representation learning一个更well-defined, explainable的objective.
2.4. Data Augmentation and Multiple Heads
在实现中作者发现, 对于数据使用data augmentation对于最终模型的performance有着很大的影响, 和DeepCluster一样, 本文中作者也是用了data transformation来实现数据增广.
$$ P_{y i}=\mathbb{E}_{t}\left[ \operatorname{softmax}_{y} h \circ \Phi\left(t \boldsymbol{x}_{i}\right)\right] $$
其中$t$是随机sample出的图像transformation.
同时, 一个直觉上的想法是, 即使是同一个数据集, 我们也可以有多种同样好的label assignment的方法. 例如可以通过: 颜色, 大小, 视角等各种方式进行label assignment. 为了强调这一点, 我们可以使用multiple heads: $h_1, \cdots, h_T$, 且所有的分类模型共用一个representation model $\boldsymbol{x} = \Phi(I)$, 从而提升$\Phi(\cdot)$抽取的特征的有效性.