SPQ：自监督量化的无监督图像检索算法

Self-supervised Product Quantization for Deep Unsupervised Image Retrieval

基于自监督量化的无监督图像检索

提出了一种新的无监督学习方案：交叉量化对比学习，如下图所示：

Fig1 cross quantized contrastive learning

$(t,t^{\prime }\sim \mathcal{T})$为取自集合$\mathcal{T}$的图形变换（数据增强）。

普通的对比学习对特征向量进行投影后计算相似度，SPQ中提出的量化对比学习采用量化方案对特征向量编码后，使用编码之后的结果与原特征向量进行对比学习，并设法使二者之间相似度最小。

---

整体架构

整个算法流程图如下所示：

Fig2 Pipeline of SPQ

交叉相似度使用原始特征$\hat{x}$与经过量化的特征$\hat{z}$计算

具体来说，可以将整个任务描述为：学习一个映射$\mathcal{R}:x\mapsto b$，$\mathcal{R}$定义为整个系统，$x$为一个包含了$N$张训练样本的数据集 X = { x n } n = 1 N \mathcal{X}=\lbrace{x_n}\rbrace^{N}_{n=1} X={xn​}n=1N​中的一张图像，$\hat{b}$是一个$B$位的二值编码 b ^ ∈ { 0 , 1 } B \hat{b} \in \lbrace{0,1}\rbrace^B b^∈{0,1}B。
如图2中所示，SPQ系统（$\mathcal{R}$）含有一个基于CNN的特征提取器$\mathcal{F}(x;{\theta }_{\mathcal{F}})$，这个CNN网络输出一个D维的特征向量：$x\in R^D$。为了加速图像检索数度，SPQ使用一个包含$M$个码本(codebooks)的量化器$\mathcal{Q}(\hat{x};{\theta }_{\mathcal{Q}})$，其中 { C 1 , . . . , C M } ⊂ Q \lbrace{C_1,...,C_M}\rbrace \subset \mathcal{Q} {C1​,...,CM​}⊂Q，每个$C_i$包含有$K$个码字$c\in R^{D/M}$，因此第m个码本可以描述为 C m = { c m 1 , . . . , c m K } C_m=\lbrace{c_{m1},...,c_{mK}}\rbrace Cm​={cm1​,...,cmK​}。深度特征空间为分为若干子空间的笛卡尔乘积，对应子空间的每个码本都表现出数据集$\mathcal{X}$的几个不同的特征。码本中的每个码字都指向分割的深度描述符的其中一个聚类中心，目的是容纳一个频繁出现的局部模式（local pattern）。
在量化过程中，图像之间相似的特征通过分配给相同的码字来实现共享，而可区分的特征则属于不同的码字。

---

自监督训练

在SPQ训练体系中，${\theta }_{\mathcal{F}}$与${\theta }_{\mathcal{Q}}$被定义为特征提取器$\mathcal{F}$与量化器$\mathcal{Q}$的训练参数。由于硬分配量化方案（hard assignment quantization）是很难计算的，所以SPQ采用了软量化器$q_m(.)$定义如下：
$$z_{nm}=\sum _k^K\frac{\exp (-||x_{nm}-c_{mk}|{|}_2^2/{\tau }_q)}{{\sum }_{k^{\prime }}^K\exp (-||x_{n}m-c_{m{k}^{\prime }}|{|}_2^2/{\tau }_q)}{c}_{mk}$$
其中，${\tau }_q$是非负的温度系数，用于控制softmax的输出尺度，而$||.|{|}_2^2$定义为欧几里德距离，用以衡量两个输入之间的相似性。在这种形式下，每个子量化描述符$z_{nm}=q_m(x_{nm};{\tau }_q,C_m)$可以被视为码本$C_m$中的码字的指数加权和。需要注意，这个操作相当于码本中的所有码字都被用来近似量化输出，其中与输出最接近的码字贡献最大。
与之前的深度量化（deep PQ）方案不同的是，SPQ中排除了内部归一化操作（intra-normalization）。在将子量化描述符连接成整体量化描述符$\hat{z}$的过程中，内部归一化可以最小化视觉特征的影响，但是由于SPQ的训练中完全无人工监督，这有助于发现不同的特征。在SPQ中，算法致力于捕捉主要的视觉特征，而不是试图平衡每个码本的影响。
为了同时训练深度描述符和码字，SPQ使用了交叉量化对比学习（cross quantized contrastive learning）模式，尝试对各种视图（转换图像）的深度描述符和乘积量化描述符进行对比。如图2所示，如果视图源自同一幅图像，则深度描述符和量化描述符被视为相关的，反之则视为不相关的。同时，为了增强码字的泛化能力，深度描述符和自身的量化描述符之间的相关性被忽略了，这是因为当子向量与最接近的码字一致性最大化时，其他码字的贡献就减小了。
在一个大小为$N_B的$mini-batch中，我们对其随机应用两次数据增强，以生成新的大小为$2N_B$的数据视图。将源自同一幅图像的视图$(\widetilde{x_i},\widetilde{x_j})$视为相关的，而其他$2(N_B-1)$幅视图视为无关的。基于这个假设，SPQ设计了一个交叉量化损失函数以学习相关的图像对$(i,j)$：
$$l_{(}i,j)=-\log \frac{\exp (\mathcal{S}(i,j)/{\tau }_{cqc})}{{\sum }_{n=1}^{N_B}{1}_{n^{\prime }\ne j}\exp (\mathcal{S}(i,n^{\prime })/{\tau }_{cqc})}$$
其中，$\mathcal{S}(i,j)$定义了$\widehat{x_i}$与$\widehat{z_j}$之间的余弦相似度，使用余弦相似度来度量是为了避免$\widehat{x_i}$与$\widehat{z_j}$之间的范数偏差。

CQC Loss的python实现实际如下：

class CQCLoss(T.nn.Module):
    def __init__(self, device, batch_size, tau_cqc):
        super(CQCLoss, self).__init__()
        self.batch_size = batch_size
        self.tau_cqc = tau_cqc
        self.device = device
        self.COSSIM = T.nn.CosineSimilarity(dim=-1)
        self.CE = T.nn.CrossEntropyLoss(reduction="sum")
        # 定义一个minibatch中相关的图像对的mask
        self.get_corr_mask = self._get_correlated_mask().type(T.bool)

    def _get_correlated_mask(self):
        """定义一个minibatch中相关的图像对的mask
        由主对角线和两条副对角线构成
        副对角线的位置与batchsize相关

        Returns:
            一个minibatch中相关的图像对的mask
        """
        diag = np.eye(2 * self.batch_size)
        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
        mask = T.from_numpy((diag + l1 + l2))
        mask = (1 - mask).type(T.bool)
        return mask.to(self.device)

    def forward(self, Xa, Xb, Za, Zb):
        """
        Xa与Xb是同一个batch的图像经过不同数据增强生成的新batch
        Za与Zb是Xa与Xb对应的量化向量
        """
        XaZb = T.cat([Xa, Zb], dim=0)
        XbZa = T.cat([Xb, Za], dim=0)

        Cossim_ab = self.COSSIM(XaZb.unsqueeze(1), XaZb.unsqueeze(0))
        Rab = T.diag(Cossim_ab, self.batch_size)
        Lab = T.diag(Cossim_ab, -self.batch_size)
        Pos_ab = T.cat([Rab, Lab]).view(2 * self.batch_size, 1)
        Neg_ab = Cossim_ab[self.get_corr_mask].view(2 * self.batch_size, -1)
        
        # 计算向量之间的相关性矩阵
        Cossim_ba = self.COSSIM(XbZa.unsqueeze(1), XbZa.unsqueeze(0))
        Rba = T.diag(Cossim_ba, self.batch_size)
        Lba = T.diag(Cossim_ba, -self.batch_size)
        # 寻找所有的相关图像对作为正例    
        Pos_ba = T.cat([Rba, Lba]).view(2 * self.batch_size, 1)
        # 去除所有相关的图像对即为不相关的负例
        Neg_ba = Cossim_ba[self.get_corr_mask].view(2 * self.batch_size, -1)
        # 将所有的正例放在首位
        logits_ab = T.cat((Pos_ab, Neg_ab), dim=1)
        logits_ab /= self.tau_cqc

        logits_ba = T.cat((Pos_ba, Neg_ba), dim=1)
        logits_ba /= self.tau_cqc
        # 定义标签（坐标0为正例）
        labels = T.zeros(2 * self.batch_size).to(self.device).long()

        loss = self.CE(logits_ab, labels) + self.CE(logits_ba, labels)
        return loss / (2 * self.batch_size)

代码中相关性矩阵Cossim_ab是一个$2N\times 2N$大小的方阵，$N$为batchsize，由拼接向量$X_a{Z}_b$与自身进行余弦相似度计算得到。因此其主对角线和batchsize位置的两条副对角线是同一个样本计算相似度得到的值，其他位置为不同图像计算相似度得到的值。
$$\left[\begin{array}{cccccc}{X}_{a_1}\ast X_{a_1}& \dots & X_{a_1}\ast X_{a_n}& X_{a_1}\ast Z_{b_1}& \dots & X_{a_1}\ast Z_{b_n}\\ ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ X_{a_n}\ast X_{a_1}& \dots & X_{a_n}\ast X_{a_n}& X_{a_n}\ast Z_{b_1}& \dots & X_{a_n}\ast Z_{b_n}\\ Z_{b_1}\ast X_{a_1}& \dots & Z_{b_1}\ast X_{a_n}& Z_{b_1}\ast Z_{b_1}& \dots & Z_{b_1}\ast Z_{b_n}\\ ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ Z_{b_n}\ast X_{a_1}& \dots & Z_{b_n}\ast X_{a_n}& Z_{b_n}\ast Z_{b_1}& \dots & Z_{b_n}\ast Z_{b_n}\end{array}\right]$$
因此同一个图像对的样本相似度在矩阵中的位置如下，分别为主对角线与$n=N$的两条副对角线：
$$\left[\begin{array}{cccccc}{X}_{a_1}\ast X_{a_1}& \dots & & X_{a_1}\ast Z_{b_1}& \dots & \\ ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ & \dots & X_{a_n}\ast X_{a_n}& & \dots & X_{a_n}\ast Z_{b_n}\\ Z_{b_1}\ast X_{a_1}& \dots & & Z_{b_1}\ast Z_{b_1}& \dots & \\ ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ & \dots & Z_{b_n}\ast X_{a_n}& & \dots & Z_{b_n}\ast Z_{b_n}\end{array}\right]$$

参考文献

[1] Self-supervised Product Quantization for Deep Unsupervised Image Retrieval
[2] Github官方代码