[論文レビュー] 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は、過去のミニバッチからの埋め込みを保存・動的に更新し、モデルの訓練中に「ゆっくりとした変化（slow drift）」現象を利用することで、単純な対照的損失（contrastive loss）でも最先端の性能を達成する。大規模な画像検索データセットにおいて、R@1を12%〜22.5%向上させ、GPUメモリをわずか0.2 GB追加するだけで実現した。

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.

研究の動機と目的

• ミニバッチサイズに制限を受けるペアベースの深層度量学習における根本的制限を解消すること。
• 訓練中に特徴量がゆっくりと変化する（slow drift）現象により、過去の埋め込みが現在の埋め込みをどれほどよく近似できるかを調査すること。
• 計算コストを増加させることなく、複数のミニバッチにわたる情報豊富なネガティブペアの抽出を可能にするメモリ機構を開発すること。
• XBMを単純な損失関数（例：対照的損失）に統合することで、複雑な最先端手法を上回ることを実証すること。
• 既存のペアベースDMLフレームワークと互換性があり、メモリ効率が高く、即挿入可能なソリューションを提供すること。

提案手法

• 最近のミニバッチからの埋め込みを保存するクロスバッチメモリ（XBM）モジュールを提案し、シンプルなキューイング・デキューイング機構を用いる。
• 「ゆっくりとした変化（slow drift）」現象を活用：過去の埋め込みは、再計算なしで現在のモデル特徴量をよく近似できる。
• 現在のアーキテクチャを、保存済みの過去の埋め込みと接続することで、任意のペアベースDMLフレームワークにXBMを統合可能にする。
• 前方伝搬時に計算された特徴量を用いてメモリバッファを更新するため、追加の推論コストが発生せず、計算効率を確保する。
• 固定サイズのバッファを用いることで、最新の埋め込みのみを保持し、わずか約0.2 GBの追加GPUメモリでメモリ効率を確保する。
• XBMを拡張した損失関数を訓練中に適用し、各アーキテクチャがバッチ内およびバッチ間のネガティブサンプル（メモリに保存されたもの）と比較される。

実験結果

リサーチクエスチョン

• RQ1 訓練の初期段階での過去の埋め込みが、ゆっくりとした特徴量の変化により、現在の埋め込みをどれほどよく近似できるか。
• RQ2 ペアベースの深層度量学習において、クロスバッチのハードネガティブサンプル抽出が性能にどれほど寄与するか。
• RQ3 XBMを搭載した単純な対照的損失が、MS loss や FastAP といった複雑な最先端手法を上回るか。
• RQ4 XBMは、画像検索における異なるバックボーンネットワークやデータセットでどのように性能を発揮するか。
• RQ5 実際の運用において、クロスバッチメモリを維持するためのメモリおよび計算コストはどの程度か。

主な発見

• SOP、In-shop、VehicleID の3つの大規模な画像検索データセットにおいて、XBMはベースラインの対照的損失に対してR@1を12%〜22.5%向上させた。
• XBMを導入した単純な対照的損失は、MS loss [37]、FastAP [2]、MIC [26] といった最先端手法を大きく上回った。
• SOPデータセットでは、XBMを搭載したResNet50がR@1 80.6%を達成し、以前のSOTA（MS loss 78.2%、MIC 77.2%）を上回った。
• In-shopでは、XBMを搭載したResNet50がR@1 91.3%を達成し、以前の最高記録（FastAP 90.9%）を更新した。
• VehicleIDでは、XBMを搭載したResNet50が、スモールセットで94.7%、ミディアムで93.7%、ラージで93.0%を記録し、FastAPおよびMICを上回った。
• XBM機構は、追加GPUメモリをわずか0.2 GBに抑え、既存のDMLフレームワークへの統合に数行のコードで実現可能だった。