Learning to Select Visual In-Context Demonstrations

Eugene Lee; Yu-Chi Lin; Jiajie Diao

cs.LGMar 24, 2026

Learning to Select Visual In-Context Demonstrations

Authors: Eugene Lee, Yu-Chi Lin, Jiajie Diao

Organizations: University of Cincinnati · University of California, Los Angeles

Abstract

Multimodal Large Language Models (MLLMs) adapt to visual tasks via in-context learning (ICL), which relies heavily on demonstration quality. The dominant demonstration selection strategy is unsupervised k-Nearest Neighbor (kNN) search. While simple, this similarity-first approach is sub-optimal for complex factual regression tasks; it selects redundant examples that fail to capture the task's full output range. We reframe selection as a sequential decision-making problem and introduce Learning to Select Demonstrations (LSD), training a Reinforcement Learning agent to construct optimal demonstration sets. Using a Dueling DQN with a query-centric Transformer Decoder, our agent learns a policy that maximizes MLLM downstream performance. Evaluating across five visual regression benchmarks, we uncover a crucial dichotomy: while kNN remains optimal for subjective preference tasks, LSD significantly outperforms baselines on objective, factual regression tasks. By balancing visual relevance with diversity, LSD better defines regression boundaries, illuminating when learned selection is strictly necessary for visual ICL.

Figures & tables

Figure 1: An overview of our LSD (Learning to Select Demonstrations) framework. The process is a training loop where the MLLM acts as the Environment. (1) The Agent (a Dueling DQN) receives the current state $s_{t}$ , which contains the query embedding $\mathbf{e}_{q}$ and the embeddings of all previously selected demonstrations $\{\mathbf{e}_{1},\dots,\mathbf{e}_{t-1}\}$ . (2) The agent’s query-centric decoder outputs an advantage query $\mathbf{a}_{s}$ , which is used to retrieve candidates $A_{\text{cand}}$ from the Task’s Data via FAISS. (3) The agent selects the next best demonstration, $d_{t}$ . (4) The full prompt (including the selected demos $d_{1}\dots d_{K}$ and the query) is sent to the MLLM (Environment), which makes a prediction. (5) A Reward $r_{t}$ is calculated based on the prediction’s accuracy (e.g., MAE). (6) This reward is used to update the agent’s policy.

			Random				kNN				LSD (Ours)
Dataset	Sup.	0-Shot	K=1	K=4	K=8	K=16	K=1	K=4	K=8	K=16	K=1	K=4	K=8	K=16
UTKFace	4.42 [ 27 ]	6.10	6.14	10.66	14.86	12.51	5.98	7.27	7.61	7.60	5.90	6.27	7.05	6.64
AVA	-	1.32	1.21	1.03	0.92	0.92	1.20	0.98	0.83	0.86	1.20	1.06	0.98	0.98
SCUT-FBP5500	0.26 [ 45 ]	0.59	0.58	0.64	0.64	0.68	0.53	0.39	0.40	0.44	0.55	0.62	0.67	0.75
KonIQ-10k	0.39 [ 34 ]	0.42	0.42	0.48	0.44	0.56	0.40	0.44	0.55	0.61	0.39	0.40	0.51	0.51
KADID-10k	-	0.94	0.89	1.07	1.05	1.07	0.87	0.87	0.91	0.92	0.76	0.79	0.82	0.84

Table 1: Main Performance (MAE

\downarrow

) Comparison vs. Number of Shots (

K

). We report the Mean Absolute Error (MAE) for all methods on the five benchmark datasets, evaluated with Gemma 3 4B-it for

K\in\{1,4,8,16\}

. Our proposed method, LSD , consistently outperforms all baselines, and the performance gap widens as

K

increases. The 0-shot and a fully Supervised (Sup.) baseline are also provided for reference. Best results are in bold .

Panel of Figure 2 — Figure 2: Performance vs. Number of Shots ( $K$ ) on four datasets. We plot the MAE as $K$ increases. The results are task-dependent: (a), (c), (d) Objective Tasks (UTKFace, KonIQ, KADID): Our LSD policy (blue) consistently outperforms the kNN baseline (orange). (b) Subjective Task (AVA): The kNN baseline, which is based on visual similarity, consistently outperforms LSD.

Panel of Figure 3 — Figure 3: Demonstration Set Analysis on UTKFace, plotted against $K$ shots. (a) MAE of Demo Labels vs. Query: The MAE between selected demo labels and the query’s true label. LSD finds demos with closer labels. (b) Pairwise Label MAE: The MAE computed over all pairwise label differences among the selected demos. (c) Demo-Query Feature Similarity: The cosine similarity between demo embeddings and the query embedding. LSD balances similarity with other factors. (d) Pairwise Feature Similarity: The cosine similarity between every pair of selected demonstrations. LSD actively seeks diverse (low-similarity) demos.

Figure 4: Qualitative Comparison of Selected Demonstrations ( $K=12$ ). (a) UTKFace: For an 8-year-old query, kNN selects only images with highly similar features (e.g., other young children). LSD selects a diverse spectrum of visual features (e.g., varied ages, genders, and lighting conditions) to build a richer context. (b) KADID-10k: For a motion-blurred query, kNN selects only other distorted versions of the same source image . LSD selects a varied set, including the pristine original and images with different distortion types from different source images , defining the quality boundaries.

Panel of Figure 5 — Figure 5: Cross-MLLM Generalization (MAE $\downarrow$ ) on UTKFace vs. Number of Shots ( $K$ ). We use the single LSD policy (trained on Gemma 3 4B-it) to select demos for two unseen MLLMs. The plots show our policy (blue line) versus the kNN (orange line) and Random (green line) baselines. (a) On Qwen 2.5 7B, our policy consistently outperforms kNN. (b) On Phi-3.5-vision, our policy performs on par with kNN. Both LSD and kNN significantly outperform the Random baseline.

Dataset	LSD (Learned Order)	LSD (Shuffled Set)
UTKFace	7.05	6.51
AVA	0.98	1.04
SCUT-FBP5500	0.67	0.53
KonIQ-10k	0.51	0.59
KADID-10k	0.82	0.82

Table 2: Analysis of Selection Order (MAE

\downarrow

) at

K=8

. We compare the performance of the agent’s learned demonstration sequence against the exact same set of demonstrations in a random order.

	MAE $\downarrow$
Decoder Input Strategy	K=4	K=8	K=16	Policy Behavior
Query-Centric	6.27	7.05	6.64	Query-specific demos
Concat Input	7.01	6.42	7.74	Non-query-specific

Table 3: Ablation Study on Decoder Input Strategy (MAE

\downarrow

). We compare our query-centric model against a standard decoder-only model (Concat Input) on UTKFace for

K\in\{4,8,16\}

, both using

L=2

layers. We also note the qualitative policy behavior.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Task / Dataset	Query Instruction ( $Q_{task}$ )
Age Prediction	“What is the age of the person in the last image? Only output the estimated age as a number.”
(UTKFace)
Aesthetic Assessment	“What is the aesthetic score of the last image on a scale from 0 to 10? Only output the score as a floating number.”
(AVA)
Facial Beauty	“What is the facial beauty score of the last image on a scale from 0 to 5? Only output the score as a floating number.”
(FBP5500)

Appendix

Table 4: List of task-specific instructions used to prompt the VLM. The model is provided with

K

interleaved image-label pairs prior to these instructions.

Feature	Discrete RL (e.g., PPO)	Proposed Method
Output Space	$N$ logits (One per item)	Vector $\in\mathbb{R}^{D}$
Selection	Softmax over $N$	ANN Retrieval ( $k$ ) $\to$ Argmax
Semantics	Orthogonal Actions	Semantic Neighbors
Complexity	Linear $O(N)$	Logarithmic $O(\log N+k)$
Scalability	Limited ( $N<10^{4}$ )	Unlimited (via FAISS)

Appendix

Table 5: Comparison between Standard Discrete RL and the proposed Continuous Embedding (Wolpertinger) approach for large-scale selection.

Panel of Figure 6 — Figure 6: Extended Feature-Space Analysis (Relevance and Similarity). The plots in the right column demonstrate that on all five datasets, LSD (blue line) actively seeks low redundancy, maintaining the trend $\text{LSD}\ll\text{kNN}$ in pairwise similarity, which is the key behavioral difference.

Panel of Figure 7 — Figure 7: Extended Label-Space Analysis (Emergent Relevance and Consistency). The results show a critical, task-dependent pattern in minimizing the label difference ( $\Delta$ Label) between the query and selected demos. For Objective Tasks (UTKFace, KonIQ-10k, KADID-10k), the LSD policy (blue) is the most effective implicit label retriever. Conversely, for Subjective Tasks (AVA, SCUT-FBP5500), the kNN Baseline (orange) consistently selects demos that minimize the label difference. This confirms that the optimal policy for minimizing label-space MAE aligns with the task’s underlying human perception structure.

	Target MLLM (Evaluation)
Source Policy (Training)	Gemma 3 4B	Qwen 2.5 7B	Phi-3.5-vision	InternVL2-8B
Baseline: kNN	7.27	6.54	5.96	7.38
LSD (Gemma)	6.27	5.58	6.05	10.62
LSD (Qwen)	6.81	5.95	6.17	9.69
LSD (InternVL)	6.01	5.77	5.06	8.87

Appendix

Table 6: Cross-Model Transfer Matrix (MAE

\downarrow

) on UTKFace at

K=4

. Rows represent the Source Policy (training model). Columns represent the Target MLLM (evaluation model). Diagonal entries (gray) represent Intra-Model performance (Specialization), where LSD typically achieves its peak performance. Off-diagonal entries represent Inter-Model performance (Generalization). While LSD consistently outperforms Random selection (not shown), it performs comparably to the kNN baseline in cross-model settings, suggesting that model-specific nuances play a significant role in optimal retrieval.

Panel of Figure 8 — Figure 8: Transfer Scaling for Source Policy: Gemma 3 4B-it. Performance of the Gemma-trained LSD policy evaluated across all four target models. The policy generalizes well, consistently beating kNN on Qwen and InternVL, and matching it on Phi.

Panel of Figure 9 — Figure 9: Transfer Scaling for Source Policy: Qwen 2.5 7B. Performance of the Qwen-trained LSD policy evaluated across all targets.

Panel of Figure 10 — Figure 10: Transfer Scaling for Source Policy: InternVL2-8B. Performance of the InternVL-trained LSD policy evaluated across all targets.

Panel of Figure 11 — Figure 11: Cross-Dataset Generalization Results. We compare LSD-Self (Blue, trained on target), LSD-Cross (Red, trained on Age), kNN (Orange), and Random (Green). (d) Successful Transfer: On KADID-10k, the Age policy (Red) transfers remarkably well, matching the Self-trained agent and beating kNN/Random. (b) Negative Transfer: On SCUT-FBP5500, the Age policy hurts performance, performing worse than random. (a) Generic Policy: On AVA, the Cross and Self policies perform identically, suggesting the learned retrieval strategy is task-agnostic but inferior to semantic matching (kNN) for aesthetics.

		K=1			K=4			K=8
Dataset	0-Shot	kNN	Self	Cross	kNN	Self	Cross	kNN	Self	Cross
AVA	1.38	1.58	1.70	1.61	1.22	1.26	1.39	0.98	0.99	1.18
SCUT-FBP5500	1.07	0.91	1.13	1.01	0.36	0.53	1.03	0.38	0.52	0.98
KonIQ-10k	0.78	0.73	0.74	0.72	0.45	0.48	0.57	0.37	0.38	0.44
KADID-10k	1.13	1.16	1.07	1.12	1.06	0.98	0.94	1.05	0.94	1.02

Appendix

Table 7: Cross-Dataset Generalization Analysis. We report the MAE (

\downarrow

) for the kNN baseline, the Domain-Specific Agent ( Self ), and the Cross-Trained Agent ( Cross , trained on UTKFace) across

K\in\{1,4,8\}

. LSD-Self typically sets the upper bound. Comparing LSD-Cross to kNN reveals where the learned “diversity” heuristic transfers effectively (e.g., KADID-10k) versus where domain-specific visual matching is superior (e.g., AVA).

Figure 12: Extended Qualitative Comparison of Selected Demonstrations ( $K=11$ ) across benchmark datasets. Row 1: UTKFace (Age). For a query of a young child , the kNN baseline retrieves a homogeneous set of other children and babies. In contrast, LSD (Ours) retrieves a diverse timeline of faces, ranging from toddlers to adults and the elderly, providing the VLM with a complete regression scale. Row 2: KADID-10k (Quality). For a query of a beach scene, kNN fails by retrieving near-duplicate versions of the same source image , adding zero new information. LSD selects visually distinct scenes (sports, cityscapes) with varying distortion types. Row 3: KonIQ-10k (Quality). For a query of abstract metal spheres, LSD retrieves a broad semantic range (animals, food, landscapes), whereas kNN gets stuck in a narrow cluster of indoor/portrait shots. Row 4: AVA (Aesthetics). For a query of red berries, kNN relies on color and content matching, retrieving flamingos and other birds. LSD ignores the specific content, selecting architecture and objects to illustrate broad aesthetic principles. Row 5: SCUT-FBP5500 (Beauty). For a query of an Asian male, kNN exhibits high demographic bias, retrieving only other Asian males. LSD retrieves a diverse set of demographics (varying gender and race), reducing bias and helping the model score attractiveness independent of demographic features.

In-Context Example Selection

In-context example selection chooses demonstrations that help language and multimodal models perform a query-specific task without changing model weights.

See all 29 papers

MAG: MAnifold Guided Semi-Supervised Multi-modal In-Context Learning
Aug 13, 2026Zirui Cheng, Xun Xu, Tiankai Chen +7In-Context Example Selection Multimodal ICL
College of Design and Engineering, National University of Singapore · Institute for Infocomm Research, ASTAR, Singapore
Retrieving Counterfactuals Improves Visual In-Context Learning
Mar 17, 2026Guangzhi Xiong, Sanchit Sinha, Zhenghao He +1Multimodal ICL Retrieval-Augmented ICL
University of Virginia
Watch and Learn: Teaching Visual In-Context Learning by Removing Textual Shortcuts
Oct 22, 2026Nimrod Shabtay, Asaf Yehudai, Avihu Dekel +2Multimodal ICL In-Context Learning
IBM Research · Tel-Aviv University

Aug 13, 2026J/K move · Enter open · S save

MAG: MAnifold Guided Semi-Supervised Multi-modal In-Context Learning

Zirui Cheng, Xun Xu, Tiankai Chen, Fady Rezk+6

College of Design and Engineering, National University of Singapore · Institute for Infocomm Research, ASTAR, Singapore

In-Context Example Selection Multimodal ICL Multimodal Large Language Models Semi-Supervised Learning

Open paper AskPDF arXiv

Few-shot in-context learning (ICL) with multi-modal large language models (MLLMs) enables task adaptation without parameter updates, but its performance is highly sensitive to the quality and coverage of the selected demonstrations. While unlabeled multi-modal data is abundant, it remains elusive how to exploit them for ICL. We propose MAG (MAnifold-Guided semi-supervised in-context demonstra- tion selection), an efficient framework that leverages unlabeled data to improve multi-modal ICL. MAG formulates demonstration selection as a semi-supervised propagation problem on a multi-modal graph and adopts a two-stage strategy: (i) relevance score propagation identifies a compact set of high-impact unlabeled samples for pseudo-labeling, reducing MLLM inference cost; (ii) multi-modal relevance is used to select the final demonstrations. We show that textual represen- tations are more effective for relevance propagation, while both visual and textual modalities are crucial for high-quality demonstration selection. Experiments on eight multi-modal benchmarks demonstrate that MAG consistently outperforms strong baselines in label-scarce regimes, achieving significant gains with a limited pseudo-labeling budget.

Abstract

In-Context Example Selection

Explore similar work

MAG: MAnifold Guided Semi-Supervised Multi-modal In-Context Learning

Retrieving Counterfactuals Improves Visual In-Context Learning

Watch and Learn: Teaching Visual In-Context Learning by Removing Textual Shortcuts