[论文解读] Cross-Batch Memory for Embedding Learning

Xun Wang, Haozhi Zhang+2 more|arXiv (Cornell University)|Dec 14, 2019

Domain Adaptation and Few-Shot Learning参考文献 12被引用 14

一句话总结

该论文提出了一种交叉小批量记忆（XBM）机制，这是一种增强记忆的机制，能够存储并动态更新来自过去小批量的嵌入表示，从而在深度度量学习中实现跨小批量的难负样本挖掘。通过利用‘缓慢漂移’现象——即嵌入表示在训练过程中逐渐变化——XBM使简单的对比损失即可达到最先进性能，在仅增加0.2 GB额外GPU显存的前提下，使大规模图像检索数据集上的R@1指标提升12%至22.5%。

ABSTRACT

Mining informative negative instances are of central importance to deep metric learning (DML), however this task is intrinsically limited by mini-batch training, where only a mini-batch of instances is accessible at each iteration. In this paper, we identify a "slow drift" phenomena by observing that the embedding features drift exceptionally slow even as the model parameters are updating throughout the training process. This suggests that the features of instances computed at preceding iterations can be used to considerably approximate their features extracted by the current model. We propose a cross-batch memory (XBM) mechanism that memorizes the embeddings of past iterations, allowing the model to collect sufficient hard negative pairs across multiple mini-batches - even over the whole dataset. Our XBM can be directly integrated into a general pair-based DML framework, where the XBM augmented DML can boost performance considerably. In particular, without bells and whistles, a simple contrastive loss with our XBM can have large R@1 improvements of 12%-22.5% on three large-scale image retrieval datasets, surpassing the most sophisticated state-of-the-art methods, by a large margin. Our XBM is conceptually simple, easy to implement - using several lines of codes, and is memory efficient - with a negligible 0.2 GB extra GPU memory. Code is available at: https://github.com/MalongTech/research-xbm.

研究动机与目标

• 解决基于样本对的深度度量学习中的根本性限制：即难负样本挖掘受限于小批量大小。
• 探究由于训练过程中的特征缓慢漂移，过去嵌入是否可近似当前嵌入。
• 开发一种记忆机制，实现在不增加计算成本的前提下，跨多个小批量挖掘具有信息量的负样本对。
• 证明将XBM集成到如对比损失等简单损失函数中，可超越复杂的最先进方法。
• 提供一种内存高效、即插即用的解决方案，兼容现有的基于样本对的DML框架。

提出的方法

• 提出一种跨小批量记忆（XBM）模块，利用简单的入队-出队机制存储近期小批量的嵌入表示。
• 利用‘缓慢漂移’现象：过去嵌入在很大程度上仍能代表当前模型的特征，从而实现无需重新计算的近似。
• 通过将当前锚点与存储的过去嵌入连接，将XBM集成到任意基于样本对的DML框架中，用于损失计算。
• 使用前向传播过程中计算的特征更新内存缓冲区——无需额外推理开销——确保计算效率。
• 采用固定大小的缓冲区，仅保留最近的嵌入表示，从而实现内存高效，额外GPU显存仅约0.2 GB。
• 在训练过程中应用XBM增强的损失，其中每个锚点同时与内存中存储的跨小批量负样本及小批量内负样本进行比较。

实验结果

研究问题

• RQ1 由于训练过程中的缓慢特征漂移，是否可有效利用早期训练迭代的过去嵌入来近似当前嵌入？
• RQ2 在基于样本对的深度度量学习中，跨小批量挖掘难负样本能在多大程度上提升性能？
• RQ3 在XBM增强下，简单的对比损失是否可超越如MS损失或FastAP等复杂最先进方法？
• RQ4 XBM在不同主干网络和图像检索数据集上的表现如何？
• RQ5 在实际应用中，维护跨小批量记忆的内存与计算成本是多少？

主要发现

• 在三个大规模图像检索数据集（SOP、In-shop和VehicleID）上，XBM使基线对比损失的R@1指标提升12%至22.5%。
• 引入XBM后，简单的对比损失在性能上大幅超越最先进方法，如MS损失[37]、FastAP[2]和MIC[26]。
• 在SOP数据集上，XBM结合ResNet50的R@1达到80.6%，超越此前最先进方法（MS损失为78.2%，MIC为77.2%）。
• 在In-shop数据集上，XBM结合ResNet50的R@1达到91.3%，超过此前最佳结果（FastAP为90.9%）。
• 在VehicleID数据集上，XBM结合ResNet50在小集上达到94.7% R@1，中集为93.7%，大集为93.0%，优于FastAP与MIC。
• XBM机制仅增加0.2 GB额外GPU显存，且仅需几行代码即可集成到现有DML框架中。