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

Abstract: Visual in-context learning (VICL) enables visual foundation models to handle multiple tasks by steering them with demonstrative prompts. The choice of such prompts largely influences VICL performance, standing out as a key challenge. Prior work has made substantial progress on prompt retrieval and reranking strategies, but mainly focuses on prompt images while overlooking labels. We reveal these approaches sometimes get visually similar but label-inconsistent prompts, which potentially degrade VICL performance. On the other hand, higher label consistency between query and prompts preferably indicates stronger VICL results. Motivated by these findings, we develop a framework named LaPR (Label-aware Prompt Retrieval), which highlights the role of labels in prompt selection. Our framework first designs an image-label joint representation for prompts to incorporate label cues explicitly. Besides, to handle unavailable query labels at test time, we introduce a mixture-of-expert mechanism to the dual encoders with query-adaptive routing. Each expert is expected to capture a specific label mode, while the router infers query-adaptive mixture weights and helps to learn label-aware representation. We carefully design alternative optimization for experts and router, with a VICL performance-guided contrastive loss and a label-guided contrastive loss, respectively. Extensive experiments show promising and consistent improvement of LaPR on in-context segmentation, detection, and colorization tasks. Moreover, LaPR generalizes well across feature extractors and cross-fold scenarios, suggesting the importance of label utilization in prompt retrieval for VICL. Code is available at https://github.com/luotc-why/CVPR26-LaPR.

• The paper introduces a label-aware approach that fuses image and label cues to address inconsistencies in prompt retrieval for visual in-context learning.
• It employs a Mixture-of-Experts and query-adaptive routing, achieving up to a 15% mIoU gain across standard benchmarks like segmentation and detection.
• Ablation studies confirm that components such as prompt label fusion and decoupled optimization are critical for the enhanced performance of the framework.

Label-Aware Prompt Retrieval for Visual In-Context Learning

Introduction

Prompt selection is an essential component in Visual In-Context Learning (VICL), directly impacting the functional capacity of visual foundation models to perform diverse downstream tasks. Conventional prompt retrieval pipelines focus principally on encoding image-based similarity while discarding explicit label information, which can introduce substantial label inconsistency between prompts and queries. The paper "Love Me, Love My Label: Rethinking the Role of Labels in Prompt Retrieval for Visual In-Context Learning" (2604.03657) systematically re-examines this paradigm, empirically demonstrating the correlation between label consistency and VICL performance. The authors propose the first explicit, label-aware prompt retrieval framework, introducing a mechanism that fuses image and label cues for prompt representation and utilizes a Mixture-of-Experts (MoE) and query-adaptive routing for inference, even when the query label is unknown.

Figure 1: An empirical study on 100 query-prompt pairs reveals label matching consistency is positively correlated with VICL performance.

Motivation and Problem Analysis

The empirical findings underscore that mere image similarity in prompt retrieval is insufficient: visually similar but label-inconsistent prompts can degrade performance, especially in settings like segmentation, detection, or colorization where label semantics are not trivially entangled with visual similarity. As illustrated, image-only approaches tend to conflate prompt selection across queries with divergent label demands, which disrupts model conditioning and task-anchoring efficacy.

Figure 2: (a) Label-agnostic retrieval selects prompts by image similarity, susceptible to label mismatches. (b) Label-aware retrieval considers both visual and label consistency.

Label-Aware Prompt Retrieval Architecture

The proposed framework injects label information into prompt embeddings and utilizes a Mixture-of-Experts (MoE) mechanism on both query and prompt encoding branches. Each expert encodes a distinct mode, hypothesized to correspond to semantically meaningful patterns (beyond simple category-level distinctions). A query-conditioned router infers a weight distribution over these experts, which, in the absence of query-side label supervision, enables soft adaptation to implicit semantic cues within the query. At inference, label-aware prompt selection computes a cosine similarity between the routed, label-fused embeddings of the prompt and query, retrieving the most label-consistent candidate.

Figure 3: (a) Prompt labels are explicitly fused into prompt embeddings. Mode-specific experts produce distributed features, and a query-conditioned router extracts label-aware embeddings. (b) Training alternates performance-guided contrastive learning for experts and label-guided contrastive learning plus load balancing for the router.

Training Regimen

Training decouples the optimization trajectories of the experts and the router. For the experts, a contrastive loss is employed using VICL task performance for positive/negative sampling, encouraging mode specialization. The router is supervised using label-matching consistency between prompt/query pairs, alongside a KL-based load-balancing regularizer to enforce sufficient exploration and utilization of all expert paths.

Empirical Evaluation

Main Results

On standardized VICL benchmarks—foreground segmentation (Pascal-5$^i$), single-object detection (Pascal VOC 2012), and image colorization (ImageNet-1K)—the label-aware framework exhibits consistent empirical gains, surpassing prior SoTA retrieval (SupPR, Partial2Global) and rerank-based (RH-Partial2Global) pipelines. Noteworthy, segmentation mIoU improvements exceed 6% over the strongest retrieval baseline and 3.5% over reranking under canonical (no-voting) settings. For detection and colorization, complementary improvements are observed.

Qualitative Analysis

Qualitative case studies highlight the resolution of label inconsistency failure modes prevalent in label-agnostic selectors. The label-aware approach retrieves semantically aligned prompts, resulting in visibly improved task outputs (segmentation masks, colorizations).

Figure 4: Label-aware prompt retrieval robustly selects label-compatible examples, consistently outperforming label-agnostic retrieval in VICL prediction fidelity.

Cross-Fold and Cross-Extractor Generalization

The framework achieves strong transferability across cross-validation folds and generalizes effectively when substituting the backbone feature extractor (e.g., CLIP $\rightarrow$ DINOv2), indicating encoder-agnostic properties. Cross-fold transfer experiments show gains up to 15% mIoU over retrieval and rerank alternatives, supporting the argument that the expert-router architecture confers meaningful robustness across domains/tasks.

Interpretability and Mode Specialization

Expert activation analyses reveal that the query-conditioned router distributes focus differently across experts depending on the underlying category and task complexity, with some experts capturing fine-grained structures not reducible to simple class labels. Semantically similar classes exhibit analogous mode activations, suggesting the model’s ability to discover sub-structures in visual tasks.

Figure 5: Heatmap of expert mode selection ratios per category, revealing that expert activations align with semantic structure and category affinity.

Ablation Studies

Ablations confirm the criticality of each architectural component: removing prompt label fusion or query-side router reduces accuracy markedly. Optimizing the router and experts jointly (instead of alternately) leads to unstable optimization and degraded performance, underscoring the necessity for targeted, decoupled supervision. The performance-guided and label-guided contrastive objectives are synergistic; either’s removal adversely impacts results.

Implications and Future Directions

This work definitively substantiates the necessity and value of incorporating label cues in VICL prompt retrieval. The MoE with query-adaptive routing offers a principled strategy for handling label ambiguity and semantic diversity at prompt selection. The encoder-agnostic nature, empirical robustness, and substantial improvement argue for broader adoption of explicit label-aware conditioning mechanisms in multi-task, multimodal foundation models. Future research avenues include systematic exploration of mode interpretability, scalable expert/router architectures, and joint modeling of label hierarchies for complex structured prediction tasks.

Conclusion

The integration of label-aware prompt retrieval—via label-injected embeddings, expert-based mode decomposition, and query-adaptive routing—substantially improves visual in-context learning performance across a spectrum of vision tasks. Empirical analysis establishes the pivotal role of label consistency in prompt selection, augmenting both in-domain accuracy and cross-domain robustness. These results motivate a paradigm shift from image-focused to semantically balanced prompt retrieval strategies in VICL, with direct ramifications for generalist vision foundation models and prompt-based multi-task AI systems.

(2604.03657)