SwAV:通过对比聚类指派实现无监督视觉特征学习.

SwAV (Swapping Assignments between multiple Views)是一种在线对比学习方法,它不依赖于已有的数据集,而是能够随时根据新的数据进行对比学习。SwAV构造样本的一个增强版本的编码,并用该编码预测样本的另一个增强版本的编码。

给定两个数据样本$x_t,x_s$,构造特征向量$z_t,z_s$,并进一步构造编码$q_t,q_s$。则损失函数定义为特征和编码之间的相似程度:

\[\mathcal{L}_{\text{SwAV}} = \mathcal{L}(z_t,q_s) + \mathcal{L}(z_s,q_t)\]

给定$K$个原型向量(prototype vecotr,即聚类中心) \(C=\{c_1,...,c_K\}\),则可以构造样本特征被划分到不同聚类簇的概率:

\[p_t^{(k)} = \frac{\exp(z_t^Tc_k/ \tau)}{\sum_{k'} \exp(z_t^Tc_{k'}/ \tau)} , p_s^{(k)} = \frac{\exp(z_s^Tc_k/ \tau)}{\sum_{k'} \exp(z_s^Tc_{k'}/ \tau)}\]

则损失函数表示为预测编码和聚类概率之间的交叉熵:

\[\begin{aligned} \mathcal{L}(z_t,q_s) &= -\sum_k q_s^{(k)} \log p_t^{(k)} = -\sum_k q_s^{(k)} \log \frac{\exp(z_t^Tc_k/ \tau)}{\sum_{k'} \exp(z_t^Tc_{k'}/ \tau)} \\ &= -\sum_k [\frac{1}{\tau} z_t^Tc_k q_s^{(k)} - q_s^{(k)}\log \sum_{k'} \exp(z_t^Tc_{k'}/ \tau)] \\ \mathcal{L}(z_s,q_t) &= -\sum_k [\frac{1}{\tau} z_s^Tc_k q_t^{(k)} - q_t^{(k)}\log \sum_{k'} \exp(z_s^Tc_{k'}/ \tau)] \end{aligned}\]

编码$q_t,q_s$的构造被视为最优传输问题,即把$B$个样本划分到$K$个聚类中心。目标函数为最大化特征和原型向量之间的相似度,并引入编码矩阵的平滑性熵约束:

\[\mathop{\max}_{Q \in \mathcal{Q}} Trace(Q^TC^TZ) + \epsilon \mathcal{H}(Q) \\ \mathcal{Q} = \{ Q \in \Bbb{R}_+^{K \times B} | Q1_B=\frac{1}{K} 1_K,Q^T1_K=\frac{1}{B}1_B \}\]

上式约束了编码矩阵的每一行之和为$1/K$,每一列之和为$1/B$,即每个聚类中心被选中的平均次数为$B/K$。

上式通过Sinkhorn-Knopp算法迭代求解:

\[Q^* = Diag(u) \exp(\frac{C^TZ}{\epsilon}) Diag(v)\]

SwAV完整的实现过程如下:

为进一步增强SwAV的性能,作者提出了multi crop方法。在构造正样本对$x_t,x_s$时,通常是先把图像尺寸调整到$256^2$,再从其中随机裁剪两个$224^2$的图像。因为这两张图像都非常大,所以它们重叠的区域也非常多,于是它们就应该代表一个正样本。作者指出,较大的图像块捕捉的是整个场景的特征,还应该裁剪一些较小的图像块以关注物体的局部特征。

具体地,作者裁剪了两个$160^2$的大尺寸图像块和四个$96^2$的小尺寸图像块,在计算损失时只计算大尺寸图像块的编码$q$,总损失记为:

\[\mathcal{L}(z_1,z_2,...,z_{V+2}) = \sum_{i \in \{1,2\} } \sum_{v=1,...,V+2;v \neq i} \mathcal{L}(z_v,q_i)\]

实验结果表明,multi crop方法能够显著提高对比学习的性能,