Breaking the Batch Barrier (B3) of Contrastive Learning via Smart Batch Mining

TL;DR

This paper introduces B3, a novel batch construction method for contrastive learning that uses a pretrained teacher model and community detection to create high-quality, diverse batches, leading to state-of-the-art results with smaller batch sizes.

Contribution

B3 is a new batch construction strategy that improves contrastive learning by curating batches with strong negatives using a community detection approach, outperforming existing methods.

Findings

B3 achieves state-of-the-art results on MMEB benchmark.

Models trained with B3 perform well with much smaller batch sizes.

B3 generalizes across domains and weak teachers.

Abstract

Contrastive learning (CL) is a prevalent technique for training embedding models, which pulls semantically similar examples (positives) closer in the representation space while pushing dissimilar ones (negatives) further apart. A key source of negatives are 'in-batch' examples, i.e., positives from other examples in the batch. Effectiveness of such models is hence strongly influenced by the size and quality of training batches. In this work, we propose 'Breaking the Batch Barrier' (B3), a novel batch construction strategy designed to curate high-quality batches for CL. Our approach begins by using a pretrained teacher embedding model to rank all examples in the dataset, from which a sparse similarity graph is constructed. A community detection algorithm is then applied to this graph to identify clusters of examples that serve as strong negatives for one another. The clusters are then used to construct batches that are rich in in-batch negatives. Empirical results on the MMEB multimodal embedding benchmark (36 tasks) demonstrate that our method sets a new state of the art, outperforming previous best methods by +1.3 and +2.9 points at the 7B and 2B model scales, respectively. Notably, models trained with B3 surpass existing state-of-the-art results even with a batch size as small as 64, which is 4-16x smaller than that required by other methods. Moreover, experiments show that B3 generalizes well across domains and tasks, maintaining strong performance even when trained with considerably weaker teachers.