Leveraging Data to Say No: Memory Augmented Plug-and-Play Selective Prediction

TL;DR

This work tackles selective prediction for vision-language foundation models, extending abstention techniques to open-set tasks like captioning. It introduces PaPSP as a lightweight, training-free confidence scorer and MA-PaPSP as a memory-augmented enhancement that uses a retrieval dataset to stabilize embeddings and calibrate scores. By employing proxy embeddings and contrastive normalization against hard negatives, MA-PaPSP achieves consistent improvements over PaPSP across selective captioning, image-text matching, and classification, with gains amplified by larger SP-VLMs. The approach offers practical benefits for diverse domains, demonstrating robust performance without retraining and highlighting the importance of retrieval data coverage in open-set selective prediction.

Abstract

Selective prediction aims to endow predictors with a reject option, to avoid low confidence predictions. However, existing literature has primarily focused on closed-set tasks, such as visual question answering with predefined options or fixed-category classification. This paper considers selective prediction for visual language foundation models, addressing a taxonomy of tasks ranging from closed to open set and from finite to unbounded vocabularies, as in image captioning. We seek training-free approaches of low-complexity, applicable to any foundation model and consider methods based on external vision-language model embeddings, like CLIP. This is denoted as Plug-and-Play Selective Prediction (PaPSP). We identify two key challenges: (1) instability of the visual-language representations, leading to high variance in image-text embeddings, and (2) poor calibration of similarity scores. To address these issues, we propose a memory augmented PaPSP (MA-PaPSP) model, which augments PaPSP with a retrieval dataset of image-text pairs. This is leveraged to reduce embedding variance by averaging retrieved nearest-neighbor pairs and is complemented by the use of contrastive normalization to improve score calibration. Through extensive experiments on multiple datasets, we show that MA-PaPSP outperforms PaPSP and other selective prediction baselines for selective captioning, image-text matching, and fine-grained classification. Code is publicly available at https://github.com/kingston-aditya/MA-PaPSP.

Figures (27)

• Figure 1: PaPSP uses an external representation model and the CLIP score to enable selective prediction for VLM tasks like captioning without training. MA-PaPSP augments this model with an external dataset, which is leveraged to estimate proxy embeddings of greater stability and better calibrated contrastive scores. The figure shows an example where PaPSP fails but MA-PaPSP succeeds at rejecting an incorrect caption for the image shown. Also shown is the Cider-4 score between predicted and ground truth captions.
• Figure 2: Left: VLM problems. a) instability of representations: the representations of images (orange) and texts (blue) of the same concept can vary significantly, leading to unreliable similarity scores. b) poor calibration: distances between concepts of identical similarity (red ellipses) vary across the VLM embedding ${\cal F}^e$. Right:PaPSP solutions: c) proxy embeddings of a query $\bf q$ (orange) average multiple nearest neighbor representations from a retrieval dataset (blue) to produce a more stable representation (red), closer to the concept ground-truth (green). d) contrastive scores normalize similarity scores between image and predicted caption by those between the image and a set of hard-negatives, to ensure consistency across the space.
• Figure 2: Selective prediction AURC ($\downarrow$). S denotes the use of a small SP-VLM ($\text{SigLIP}_{\text{B/16}}$), while L denotes the larger $\text{SigLIP}_{\text{SO-400M}}$. 'Gain' is the % improvement of MA-PaPSP over PaPSP. 'Gain (B)' shows the same for SigLIP$_{\text{B/16}}$.
• Figure 3: CLIP scores between class labels and images of the class.
• Figure 3: Impact of retrieval set type on the AURC ($\downarrow$) of MA-PaPSP for captioning, classification, and ITM. Random dataset: MS-COCO for captioning, Flowers for classification, and SugarCrepe for ITM. Blue denotes improved performance of out-of domain vs in-domain and boldface the best results.