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 st , which contains the query embedding eq and the embeddings of all previously selected demonstrations {e1,…,et−1} . (2) The agent’s query-centric decoder outputs an advantage query as , which is used to retrieve candidates Acand from the Task’s Data via FAISS. (3) The agent selects the next best demonstration, dt . (4) The full prompt (including the selected demos d1…dK and the query) is sent to the MLLM (Environment), which makes a prediction. (5) A Reward rt 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 ↓ ) 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∈{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 .
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.
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.
Figure 5: Cross-MLLM Generalization (MAE ↓ ) 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 ↓ ) 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 ↓
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 ↓ ). We compare our query-centric model against a standard decoder-only model (Concat Input) on UTKFace for K∈{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 ( Qtask )
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 ∈RD
Selection
Softmax over N
ANN Retrieval ( k ) → Argmax
Semantics
Orthogonal Actions
Semantic Neighbors
Complexity
Linear O(N)
Logarithmic O(logN+k)
Scalability
Limited ( N<104 )
Unlimited (via FAISS)
Appendix
Table 5: Comparison between Standard Discrete RL and the proposed Continuous Embedding (Wolpertinger) approach for large-scale selection.
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 LSD≪kNN in pairwise similarity, which is the key behavioral difference.
Figure 7: Extended Label-Space Analysis (Emergent Relevance and Consistency). The results show a critical, task-dependent pattern in minimizing the label difference ( Δ 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 ↓ ) 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.
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.
Figure 9: Transfer Scaling for Source Policy: Qwen 2.5 7B. Performance of the Qwen-trained LSD policy evaluated across all targets.
Figure 10: Transfer Scaling for Source Policy: InternVL2-8B. Performance of the InternVL-trained LSD policy evaluated across all targets.
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 ( ↓ ) for the kNN baseline, the Domain-Specific Agent ( Self ), and the Cross-Trained Agent ( Cross , trained on UTKFace) across K∈{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 chooses demonstrations that help language and multimodal models perform a query-specific task without changing model weights.
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.
Zirui Cheng, Xun Xu, Tiankai Chen +7
College of Design and Engineering, National University of Singapore · Institute for Infocomm Research, ASTAR, Singapore
Vision-language models (VLMs) have achieved impressive performance across a wide range of multimodal reasoning tasks, but they often struggle to disentangle fine-grained visual attributes and reason about underlying causal relationships. In-context learning (ICL) offers a promising avenue for VLMs to adapt to new tasks, but its effectiveness critically depends on the selection of demonstration examples. Existing retrieval-augmented approaches typically rely on passive similarity-based retrieval, which tends to select correlated but non-causal examples, amplifying spurious associations and limiting model robustness. We introduce CIRCLES (Composed Image Retrieval for Causal Learning Example Selection), a novel framework that actively constructs demonstration sets by retrieving counterfactual-style examples through targeted, attribute-guided composed image retrieval. By incorporating counterfactual-style examples, CIRCLES enables VLMs to implicitly reason about the causal relations between attributes and outcomes, moving beyond superficial correlations and fostering more robust and grounded reasoning. Comprehensive experiments on four diverse datasets demonstrate that CIRCLES consistently outperforms existing methods across multiple architectures, especially on small-scale models, with pronounced gains under information scarcity. Furthermore, CIRCLES retrieves more diverse and causally informative examples, providing qualitative insights into how models leverage in-context demonstrations for improved reasoning. Our code is available at https://github.com/gzxiong/CIRCLES.
Vision-language models (VLMs) struggle with in context learning compared to their textual counterparts. Our preliminary study shows that VLMs rely on textual shortcuts and underutilize the visual demonstrations. We therefore hypothesize that VLMs are insufficiently trained to perform visual in-context learning (vICL). To this end, we construct MIME, a task-agnostic vICL training set without textual shortcuts. We ensure the instruction does not state the rule linking the demonstration images to their answers, so the model must infer it from the visual demonstrations. Fine-tuning on MIME improves vICL accuracy by 8.1 to 18.0 points across three VLMs, preserves single-image performance, and improves multi-image and text only in-context learning. In contrast, stating the rule in the training data sharply reduces this gain. We further propose DTAD (Demonstration Text Attention Dropout), which randomly withholds demonstration labels during training so that the VLM must infer the rule from the demonstration images. While achieving comparable vICL accuracy to fine-tuning on MIME, DTAD yields more consistent attention to the demonstration images. On TrueMICL, Qwen3-VL-8B trained on MIME matches frontier models such as Claude Opus 5, and surpasses them with DTAD. Finally, adding MIME to the instruction-tuning mix of LLaVA 1.5 13B more than doubles its vICL accuracy without hurting single-image abilities, whereas a same-size multi-image set improves it only slightly. These results suggest that vICL data should be a standard component of VLM training.