# Learning What Helps: Task-Aligned Context Selection for Vision Tasks

arXiv: https://arxiv.org/abs/2512.00489 · HTML: https://mlsanity.com/paper/2512.00489/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 29 Nov 2025.

## Builds on (most cited first, 30 of 31)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Vision Transformers for Dense Prediction](https://mlsanity.com/paper/2103.13413/) (2021)
- [Fine-Grained Visual Classification of Aircraft](https://mlsanity.com/paper/1306.5151/) (2013)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [Kvasir-SEG: A Segmented Polyp Dataset](https://mlsanity.com/paper/1911.07069/) (2019)
- [Skin Lesion Analysis Toward Melanoma Detection: A Challenge at the 2017 International Symposium on Biomedical Imaging (ISBI), Hosted by the International Skin Imaging Collaboration (ISIC)](https://mlsanity.com/paper/1710.05006/) (2017)
- [Perception Encoder: The best visual embeddings are not at the output of the network](https://mlsanity.com/paper/2504.13181/) (2025)
- [VisRAG: Vision-based Retrieval-augmented Generation on Multi-modality Documents](https://mlsanity.com/paper/2410.10594/) (2024)
- [Distilling Knowledge from Reader to Retriever for Question Answering](https://mlsanity.com/paper/2012.04584/) (2020)
- [BCN20000: Dermoscopic Lesions in the Wild](https://mlsanity.com/paper/1908.02288/) (2019)
- [End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering](https://mlsanity.com/paper/2106.05346/) (2021)
- [REVEAL: Retrieval-Augmented Visual-Language Pre-Training with Multi-Source Multimodal Knowledge Memory](https://mlsanity.com/paper/2212.05221/) (2022)
- [Retrieval Augmented Classification for Long-Tail Visual Recognition](https://mlsanity.com/paper/2202.11233/) (2022)
- [Describe Anything: Detailed Localized Image and Video Captioning](https://mlsanity.com/paper/2504.16072/) (2025)
- [Learning Customized Visual Models with Retrieval-Augmented Knowledge](https://mlsanity.com/paper/2301.07094/) (2023)
- [VILA-M3: Enhancing Vision-Language Models with Medical Expert Knowledge](https://mlsanity.com/paper/2411.12915/) (2024)
- [Are Natural Domain Foundation Models Useful for Medical Image Classification?](https://mlsanity.com/paper/2310.19522/) (2023)
- [RAG-DDR: Optimizing Retrieval-Augmented Generation Using Differentiable Data Rewards](https://mlsanity.com/paper/2410.13509/) (2024)
- [Retrieval Augmented Generation and Understanding in Vision: A Survey and New Outlook](https://mlsanity.com/paper/2503.18016/) (2025)
- [Retrieval-augmented Few-shot Medical Image Segmentation with Foundation Models](https://mlsanity.com/paper/2408.08813/) (2024)
- [Adaptive Multi-view and Temporal Fusing Transformer for 3D Human Pose Estimation](https://mlsanity.com/paper/2110.05092/) (2021)
- [SmartRAG: Jointly Learn RAG-Related Tasks From the Environment Feedback](https://mlsanity.com/paper/2410.18141/) (2024)
- [Few-Shot Recognition via Stage-Wise Retrieval-Augmented Finetuning](https://mlsanity.com/paper/2406.11148/) (2024)
- [Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models](https://mlsanity.com/paper/2505.03075/) (2025)

## Abstract (arXiv)

Humans often resolve visual uncertainty by comparing an image with relevant examples, but ViTs lack the ability to identify which examples would improve their predictions. We present Task-Aligned Context Selection (TACS), a framework that learns to select paired examples which truly improve task performance rather than those that merely appear similar. TACS jointly trains a selector network with the task model through a hybrid optimization scheme combining gradient-based supervision and reinforcement learning, making retrieval part of the learning objective. By aligning selection with task rewards, TACS enables discriminative models to discover which contextual examples genuinely help. Across 18 datasets covering fine-grained recognition, medical image classification, and medical image segmentation, TACS consistently outperforms similarity-based retrieval, particularly in challenging or data-limited settings.