Nearest Neighbor Normalization Improves Multimodal Retrieval (2024)

Abstract: Multimodal models leverage large-scale pretraining to achieve strong but still imperfect performance on tasks such as image captioning, visual question answering, and cross-modal retrieval. In this paper, we present a simple and efficient method for correcting errors in trained contrastive image-text retrieval models with no additional training, called Nearest Neighbor Normalization (NNN). We show an improvement on retrieval metrics in both text retrieval and image retrieval for all of the contrastive models that we tested (CLIP, BLIP, ALBEF, SigLIP, BEiT) and for both of the datasets that we used (MS-COCO and Flickr30k). NNN requires a reference database, but does not require any training on this database, and can even increase the retrieval accuracy of a model after finetuning.

## introduction contrastive image and text models are a fundamental building block of large-scale text-to-image or image-to-text retrieval systems @xcite @xcite . these models utilize contrastive loss functions to learn joint text and image embeddings, aligning embeddings for matching text and image pairs while separating embeddings for non-matching pairs. however, contrastive embeddings optimize pretraining objectives such as infonce @xcite rather than downstream retrieval accuracy, so learned embeddings can be suboptimal for retrieval @xcite . many methods for improving contrastive models on downstream retrieval tasks require additional training to adapt models across domains or aggregate information from an external database @xcite @xcite , and others are specialized for individual error categories, such as gender bias @xcite @xcite . recent training-free methods suggest that accuracy can be improved without fine-tuning, which is useful for limited-compute environments and critical for black-box embedding models. such methods typically use a reference database of query and retrieval embeddings to adapt the pretrained model to the downstream retrieval task. for instance, qbnorm and dbnorm normalize scores for each retrieval candidate by computing a softmax over the entire reference database @xcite . these approaches mitigate the hubness problem, where certain retrieval candidates ("hubs") emerge as nearest neighbors for many queries in high-dimensional embedding spaces, leading to incorrect matches @xcite . these methods tend to be computationally impractical, requiring match score calculations for every item in the database and thus scaling linearly with the size of the reference database. distribution normalization (dn) reduces complexity to constant time by using a first-order approximation of softmax normalization @xcite ): text and image embeddings are normalized by subtracting the mean reference embedding. while dn is much faster than qbnorm and dbnorm, this practicality comes at the cost of reduced retrieval accuracy. can sublinear runtime be achieved without sacrificing accuracy? in this paper, we introduce nearest neighbor normalization (nnn), a novel training-free method for contrastive retrieval (figure 1 ). like dn, it adds minimal inference overhead with sublinear time complexity relative to the reference database size-but it also outperforms both qbnorm and dbnorm on retrieval. the key idea is that nnn corrects for the effects of embeddings that are assigned disproportionately high or low retrieval scores, by normalizing per-candidate scores using only the k closest query embeddings from a reference dataset. for example, nnn reduces scores for the image of the surfer in figure 2 (a hub that incorrectly matches a large number of query captions), improving overall accuracy. section 2 provides more details on our approach, and section 3 empirically validates the effect of nnn for a range of models and datasets. overall, we contribute a new and conceptually simple approach for improving contrastive retrieval with little compute overhead. in addition to improving retrieval scores consistently for every model and dataset that we tested, nnn can reduce harmful biases such as gender bias. ## nearest neighbor normalization retrieval models compute a match score s(q, r) between a query q and database retrieval candidate r, and return the highest-scoring candidates. in the case of contrastive multimodal models such as clip, this score is typically the cosine similarity between image and text embeddings @xcite . figure 2 shows how the hubness problem @xcite manifests as a failure mode of contrastive text-to-image retrieval. some images are simply preferred by contrastive models over other images: they have high cosine similarity with a wide array of query captions. to correct for bias towards hubs in image-text retrieval, we propose nnn, an approach that estimates bias for each retrieval candidate using a database of reference queries, d. the bias is then applied as an additive correction to the original match score, then used for retrieval. specifically, given a contrastive retrieval score s(q, r) = q•r, we define the bias b(r) for a retrieval candidate r as a constant multiple (α) of the mean of s(q 1 , r), s(q 2 , r), . . . , s(q k , r), where {q 1 , . . . , q k } = d top k (r) are the k queries from the reference query dataset that have the highest similarity score s(q i , r) with r. namely, if we define the operator argmax k to denote the k arguments for the which a function attains its k maximum values, then we have d top k (r) = arg max k s (q, r) @xmath0, and our bias is computed as: nnn uses the nearest k query embeddings to differentiate similar objects, capturing fine-grained distinctions between retrieval candidates. each retrieval candidate has a constant bias score, so these scores can be computed offline and cached. the debiased retrieval score can then be computed by subtracting the estimated bias from the original score: when using vector retrieval to compute match scores, bias scores are computed in sublinear time and add a constant factor to retrieval runtime; see section 3.1 for further discussion. ## experiments we evaluate nnn on both text-to-image and imageto-text retrieval using a variety of contrastive multimodal models (clip, blip, albef, siglip, beit) @xcite @xcite @xcite on well-established retrieval datasets flickr30k and coco @xcite . we also report the accuracy of dbnorm, the top-performing baseline, using dbnorm's dualis scoring function @xcite . additional dn @xcite , qbnorm @xcite , and dualdis (a similar performing variant of dualis) baselines are discussed in appendix d. ## conclusion we introduce nearest neighbor normalization for contrastive multimodal retrieval. by precomputing bias correction scores using only the k-nearest neighbors, nnn is substantially more efficient while slightly improving accuracy over previous test-time inference methods. we also show that nnn can be used flexibly with arbitrary reference datasets and performs well at reducing gender bias.