Love Me, Love My Label: Rethinking the Role of Labels in Prompt Retrieval for Visual In-Context Learning

1. Problem Statement

Visual In-Context Learning (VICL) allows foundation models to perform various tasks (e.g., segmentation, detection) by steering them with demonstrative prompts (image-label pairs). A critical bottleneck in VICL is prompt retrieval: selecting the most relevant prompt from a database to guide the model for a specific query.

• Current Limitation: Existing state-of-the-art methods (e.g., SupPR, Partial2Global) focus primarily on visual similarity between the query image and prompt images. They largely ignore the label information associated with the prompts.
• The Consequence: This "label-agnostic" approach often retrieves prompts that are visually similar but label-inconsistent. For example, a query image of a "cat" might retrieve a prompt image containing a cat and a flower, but annotated with the label "flower." This inconsistency confuses the VICL model, degrading performance.
• The Hypothesis: The authors demonstrate that higher label consistency between the query and the retrieved prompt is positively correlated with better VICL performance. However, a challenge exists: at inference time, the query label is unavailable (unobserved), making it difficult to enforce label consistency during retrieval.

2. Methodology: LaPR (Label-aware Prompt Retrieval)

The authors propose LaPR, a framework designed to explicitly incorporate label information into the prompt selection process, even when the query label is unknown.

A. Core Architecture

LaPR utilizes a Mixture-of-Experts (MoE) mechanism combined with a Query-Adaptive Router to handle the missing query label.

1. Label-Aware Prompt Embeddings:

   • Prompts are encoded as joint representations by fusing the image features ($z^I$) and label features ($z^L$) into a single vector ($z = z^I + z^L$).
   • A set of Prompt Experts ($\bar{E}_k$) processes these joint embeddings to capture distinct "modes" (e.g., specific semantic features like "long tail" or "sharp beak").

2. Label-Aware Query Embeddings (Handling Unknown Labels):

   • Since the query label is unknown, the system cannot directly match labels. Instead, it uses a Router ($R$) to infer the implicit query label.
   • Query Experts ($E_k$): Similar to the prompt side, multiple experts encode the query image into different mode-specific representations.
   • Router: The router takes the query embedding and outputs a probability distribution ($\pi_q$) over the $K$ modes. This distribution acts as mixture weights, effectively estimating which modes (and thus which implicit labels) are relevant to the current query.
   • The final query embedding is a weighted sum of the expert outputs based on $\pi_q$.

3. Retrieval Process:

   • The system computes similarity between the label-aware query embedding and the query-relevant prompt embeddings (which are also weighted by the router's inferred mode distribution).
   • The prompt with the highest similarity score is selected.

B. Optimization Strategy: Alternating Training

To prevent the router and experts from interfering with each other's learning objectives, LaPR employs a decoupled alternating optimization scheme. Each mini-batch undergoes two successive updates:

1. Expert Step (Performance-Guided):

   • Goal: Strengthen the ability of experts to capture specific modes that lead to good VICL performance.
   • Method: The router is frozen. Experts are trained using a VICL Performance-Guided Contrastive Loss ($\mathcal{L}_{PG}$).
   • Mechanism: The system runs the VICL model (MAE-VQGAN) on candidate prompts to get actual task scores (e.g., IoU). Prompts yielding high scores are treated as positives, and low scores as negatives. The experts learn to encode features that correlate with high task performance.

2. Router Step (Label-Guided):

   • Goal: Train the router to correctly infer the mixture weights that align with the ground-truth query label.
   • Method: The experts are frozen. The router is trained using a Label-Guided Contrastive Loss ($\mathcal{L}_{LG}$) and a Load-Balancing Loss ($\mathcal{L}_{LB}$).
   • Mechanism:
     • Label Guidance: Prompts with labels consistent with the ground-truth query label are treated as positives. The router learns to assign higher weights to the experts corresponding to these consistent modes.
     • Load Balancing: A KL-divergence loss ensures all experts are utilized evenly across the batch, preventing "expert collapse" (where only one expert is used).

3. Key Contributions

1. First Label-Aware Retrieval in VICL: The paper is the first to explicitly address label inconsistency in prompt retrieval, shifting the paradigm from purely visual similarity to image-label consistency.
2. Novel MoE Architecture: Introduces a dual-sided MoE design where:
   • Prompts use joint image-label embeddings.
   • Queries use a router to infer implicit labels and dynamically select relevant modes.
3. Decoupled Optimization: Proposes a stable alternating training strategy that separates the learning of "what makes a prompt good" (Experts) from "how to identify the query's mode" (Router).
4. State-of-the-Art Performance: Demonstrates significant improvements across multiple tasks and datasets.

4. Experimental Results

The authors evaluated LaPR on three downstream tasks: Foreground Segmentation, Single-Object Detection, and Image Colorization.

• Quantitative Performance:
  • Segmentation (Pascal-5i): LaPR achieved 41.36% mIoU (average), outperforming the previous best retrieval-based method (SupPR, 35.56%) by ~5.8% and the reranking-based RH-Partial2Global (39.02%) by 2.34%.
  • Detection (Pascal VOC): LaPR achieved 32.01% mIoU, surpassing RH-Partial2Global (30.94%) by 1.07%.
  • Colorization: Achieved competitive results with lower MSE.
• Cross-Fold Transferability: LaPR showed superior generalization when trained on one fold and tested on others, outperforming SupPR by 14.8% and Partial2Global by 13.9% in average mIoU. This suggests the learned modes are finer-grained and more robust than simple category labels.
• Feature Extractor Independence: The method remained top-performing even when swapping the CLIP encoder for DINOv2, proving the framework's robustness.
• Qualitative Analysis: Visualizations (Figure 4) showed that while label-agnostic methods often retrieved visually similar but semantically wrong prompts (e.g., retrieving a "TV" prompt for a "Cat" query), LaPR consistently retrieved prompts with matching labels, leading to accurate predictions.

5. Significance and Conclusion

This paper fundamentally rethinks the role of metadata (labels) in visual retrieval. It proves that ignoring labels in prompt selection is a suboptimal strategy for Visual In-Context Learning.

• Theoretical Insight: It establishes that label consistency is a stronger predictor of VICL success than visual similarity alone.
• Practical Impact: LaPR provides a robust, generalizable framework that can be plugged into existing VICL pipelines to significantly boost performance without requiring changes to the backbone foundation model.
• Future Direction: The work opens new avenues for research into "label-aware" retrieval systems and the utilization of implicit semantic structures in multimodal learning.

In summary, LaPR solves the "label inconsistency" problem in VICL by using a smart, router-guided mixture of experts to infer query intent and retrieve prompts that are not just visually similar, but semantically and label-consistent.