Source-linked AI summary
Cross-Batch Memory for Embedding Learning
Xun Wang, Haozhi Zhang, Weilin Huang, Matthew R. Scott
TL;DR
Pair-based deep metric learning needs informative negative pairs, but mini-batch training limits the available examples. The paper uses slowly drifting embeddings in a cross-batch memory to mine negatives across recent batches, and reports substantial retrieval gains with modest memory overhead.
Problem
Mini-batch training restricts access to informative negative pairs, which are critical for training pair-based deep metric learning models.
Method
Cross-batch memory stores embeddings from recent mini-batches and connects current anchors to them for cross-batch hard-negative mining.
Results
A contrastive loss with XBM surpasses state-of-the-art methods by a large margin on three large-scale image-retrieval datasets.
Takeaways & Limitations
XBM provides a simple, directly integrable, and memory-efficient mechanism for improving hard-negative mining in pair-based DML.
Abstract
from arXiv · showhide
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.
1. Introduction
Pair-based deep metric learning depends on informative negative pairs, but mini-batch training restricts the available examples. The paper proposes cross-batch memory to reuse slowly drifting embeddings across batches, improving hard-negative mining without simply enlarging the batch.
- Pair-based DML learns from pair-wise similarities and includes contrastive, triplet, lifted-structure, n-pairs, and multi-similarity losses.
- Larger mini-batches improve pair-based DML because the number of negative pairs grows quadratically with batch size.
- Simply enlarging mini-batches is constrained by GPU memory, computational cost, and the need for difficult cross-device synchronization.
- XBM exploits slow embedding drift by storing recent mini-batch features and connecting current anchors to embeddings from multiple past batches.
- XBM integrates with existing pair-based methods, uses about 0.2 GB of extra GPU memory, and yields large retrieval gains across three large-scale datasets.
- Hard-negative mining is limited because each iteration accesses only a mini-batch, while many available pairs become uninformative as learning progresses.
3. Cross-Batch Memory Embedding Networks
Pair-based DML depends on informative negative pairs, but mini-batches restrict the available mining pool. XBM exploits slow feature drift to reuse past embeddings, expanding negative mining with a simple, memory-efficient module.
- Pair-based DML: Pair-based DML methods frame training as weighting pairwise similarities, with performance heavily dependent on mining informative negative pairs.Contrastive, triplet, and multi-similarity losses can be described within the GPW framework.
- Pair-based DML: Larger mini-batches improve pair-based DML, but GPU memory and multi-GPU communication limit this straightforward solution.The limitation follows because the number of possible training pairs depends on mini-batch size.
- Cross-Batch Memory Module: The memory bank is updated by enqueuing the current mini-batch and dequeuing the oldest, then comparing current anchors against stored embeddings with a pair-based loss.The memory-augmented loss uses similarities between current embeddings and stored past embeddings.
- Slow Drift Phenomena: XBM uses slow feature drift to retain past mini-batch embeddings and mine informative examples across multiple mini-batches.Past features can approximate current-model features because embeddings drift slowly after the early training phase.
4. Experiments and Results
Experiments show that XBM consistently improves pair-based DML while reducing dependence on large mini-batches, with modest memory and computational costs. Across three retrieval benchmarks, XBM-augmented contrastive loss surpasses state-of-the-art methods by substantial margins.
- Memory Ratio: 71.7%→78.2% recall@1 lets XBM with mini-batch 16 outperform non-memory training with mini-batch 256 while saving GPU memory.This demonstrates that cross-batch memory can substitute for much larger mini-batches in the reported setting.
- Mini-batch Size: Increasing mini-batch size from 16 to 256 improves XBM by 3.2%, versus 25.1% for original contrastive loss.The reported gap indicates that XBM substantially reduces the effect of mini-batch size on performance.
- General Pair-based DML: XBM improves contrastive loss and MS loss consistently across all three benchmarks.The memory module raises contrastive loss from 64.0% to 77.8% and MS loss from 69.7% to 76.2%.
- Memory and Computational Cost: The module requires 0.2 GB of GPU memory and 34% extra training time relative to forward and backward computation.The reported memory cost covers caching the whole training set, while affinity computation increases linearly with memory size.
- Quantitative and Qualitative Results: XBM-augmented contrastive loss surpasses state-of-the-art methods on all three datasets by large margins.Reported gains include SOP 77.2%→80.6% over MIC, In-shop 88.2%→91.3% over MIC, and VehicleID 81.9%→92.5% over A-BIER.
- Quantitative and Qualitative Results: Qualitative results indicate that XBM can retrieve correct images using discriminative details, although some retrieved results remain visually closer to the query than baseline results.The paper presents both successful examples and bad cases in Figure 7.
5. Conclusions
XBM provides a simple, memory-efficient cross-batch mining mechanism that substantially improves pair-based DML retrieval performance. It can collect hard negatives across multiple mini-batches and even the whole dataset, while extending beyond DML is suggested by the authors.
- XBM dynamically updates recent mini-batch embeddings, enabling sufficient hard-negative pairs across multiple mini-batches or the whole dataset.
- Figure 7 compares the top four retrieved images without and with the memory module, marking correct results green and incorrect results purple.
- XBM integrates directly into general pair-based DML and significantly improves existing methods on image retrieval.
- The authors suggest that XBM may extend to other machine-learning tasks because slow drift is reported beyond DML.
1. Results on More Datasets
The study evaluates XBM on CUB, Cars, and DeepFashion2, including small fine-grained and large-scale retrieval datasets. Additional training tricks help small datasets but were excluded from the reported XBM results to isolate its effectiveness.
- XBM is evaluated on CUB, Cars, and DeepFashion2, spanning small fine-grained datasets and a large-scale clothing dataset.CUB has 11,788 images across 200 classes; DeepFashion2 contains 216K clothing images and over 686K training pairs.
- Freezing batch normalization or reducing pretrained-backbone learning rates can increase Recall@1 by more than 2% on CUB and Cars.
- These additional tricks show no effect on large-scale datasets such as SOP, In-shop, and VehicleID, which contain sufficient data to mitigate overfitting.
- Reported XBM performance was trained without these additional tricks to demonstrate the module’s actual effectiveness.
2. Memory Update
The memory-update study compares enqueue-dequeue, moving-average, and back-propagation strategies. Results indicate that update behavior and dataset scale affect performance, while the simple XBM update avoids substantial computational and memory costs.
- The default XBM update enqueues latest features while dequeuing the oldest, and the study compares it with moving-average and back-propagation updates.
- Larger moving-average momentum benefits small datasets but can impair performance on large-scale datasets.
- Moving-average updates may aggregate embeddings from different augmentations when feature drift is small.
- Back-propagation requires substantially more memory and computation for gradients, yet does not yield clear improvements across all datasets.
- The findings suggest that past mini-batches can largely represent the current embedding-space distribution because embeddings drift slowly.
3. Feature Drift on General Tasks
Slow embedding drift is observed beyond pair-based DML, including image recognition. ResNet50 embeddings trained with cross-entropy on ImageNet remain relatively stable even across large training intervals.
- Slow drift occurs in other machine-learning tasks, including image recognition.
- ImageNet embeddings remain within a relatively small distance even across an interval of Δt = 1000.
- ResNet50 global-pooling embeddings drift at a slow rate when trained with cross-entropy loss on ImageNet.
4. Hyperparameters
This section lists the key hyperparameters used in the experiments and situates them in comparisons involving memory-augmented models.
- Table 7 lists the key hyperparameters applied in the experiments.
- XBM achieves outstanding performance on large-scale datasets and comparable results on small-scale datasets without training tricks or large training iterations.
- The listed settings concern training memory-augmented models compared with state-of-the-art methods.
5. Proof of Lemma 1
The proof compares gradients from accurate and approximated losses, then characterizes their error. The error can be strictly controlled by small feature drift.
- The proof computes gradients for the accurate and approximated losses.
- It then derives an expression for the gradient error between the two losses.
- Because 2 is usually less than 1, the gradient error can be strictly controlled by small feature drift.
6. Visualization.
The visualization examines XBM’s mined negatives during training and retrieved examples during testing. The figures show diverse, visually related, fine-grained, and sometimes near-identical negatives.
- The visualization samples examples from SOP, In-shop, and VehicleID under specified selection rules.
- Figure 2 shows hard negatives mined from memory with similarities above 0.5.
- The mined negatives are diverse, visually related, and fine-grained, supporting informative negative-pair construction.
- Figure 3 presents the top 10 negatives for anchors whose hardest-sample similarity exceeds 0.8.
- Some displayed negatives are extremely similar to, or exactly the same items as, their corresponding anchors.