MoCo: 自监督视觉表示学习的里程碑

MoCo:开创自监督视觉表示学习新纪元

在计算机视觉领域,如何从大量未标注数据中学习到高质量的视觉表示一直是一个重要而富有挑战性的问题。近年来,自监督学习方法在这一领域取得了巨大进展,而Facebook AI Research团队提出的MoCo(Momentum Contrast)无疑是其中的佼佼者。MoCo通过巧妙的对比学习设计和动量编码器的创新,在多个下游任务中取得了突破性成果,成为了自监督视觉表示学习的重要里程碑。

MoCo的核心思想

MoCo的核心思想是通过对比学习来学习视觉表示。具体来说,它将同一图像的不同视图(如不同的数据增强结果)视为正样本对,将不同图像视为负样本对。模型的目标是将正样本对的特征表示拉近,同时将负样本对的特征表示推远。

MoCo框架图

MoCo的创新之处主要体现在以下几个方面:

动量编码器:MoCo引入了一个动量编码器,它是主编码器的缓慢移动平均。这种设计使得模型可以维护一个大而一致的字典,从而提高对比学习的效果。
队列机制:MoCo使用一个队列来存储负样本的特征表示,而不是像之前的方法那样只使用当前mini-batch内的样本。这大大增加了负样本的数量和多样性,有助于学习更好的表示。
动量更新:MoCo采用动量更新的方式来更新key encoder的参数,这使得key encoder的参数变化更加平滑,有助于保持字典的一致性。

MoCo的实现与优化

MoCo的PyTorch实现非常简洁高效。以下是其核心代码片段:

# momentum update of key encoder
self._momentum_update_key_encoder()

# compute query features
q = self.encoder_q(im_q)  # queries: NxC
q = nn.functional.normalize(q, dim=1)

# compute key features
with torch.no_grad():  # no gradient to keys
    k = self.encoder_k(im_k)  # keys: NxC
    k = nn.functional.normalize(k, dim=1)

# compute logits
l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

# contrastive loss
logits = torch.cat([l_pos, l_neg], dim=1)
labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
loss = nn.CrossEntropyLoss().cuda()(logits / self.T, labels)

在MoCo v2中,研究人员进一步对MoCo进行了优化: