Gradient-free Task-Conditioned Retrieval for On-Device In-Context Learning

2026-07-30

AI Summary

This work addresses the challenge of efficiently retrieving local examples for in-context learning on edge devices, where constraints on computation, memory, and data privacy limit conventional approaches. The authors propose CoRA, a novel framework that enables task-conditioned retrieval without fine-tuning, backpropagation, or queries to the target model. CoRA constructs a task-conditional representation space using a frozen encoder, aligns representations via closed-form ridge regression, and incorporates optimal low-rank compression theory with a two-stage streaming indexing algorithm to support efficient multimodal example retrieval. Experiments demonstrate that CoRA significantly improves retrieval efficiency across ten text and four multimodal benchmarks and is successfully deployed on a Raspberry Pi 5, showing compatibility with models such as Llama-3.2-1B.

Abstract

On-device in-context learning (ICL) relies on pre-inference retrieval to select demonstrations for useful context before downstream model inference. This retrieval must exploit task-specific information while operating over local memories under limited computation, memory, and data-exposure budgets. We propose Conditional Retrieval Alignment (CoRA), a gradient-free framework that converts a frozen encoder into a task-conditioned retriever using paired candidate inputs and outputs. CoRA selects complementary encoder layers, constructs an output-derived conditioning space from candidate memory, and aligns candidate input representations to this space through closed-form ridge regression. Low-rank factorization then produces a compact retrieval basis where candidate outputs are used only during offline index construction, whereas query-time retrieval requires only the query input and precomputed index. We show that CoRA's rank-constrained basis is the optimal low-rank compression of the output-conditioned fitted representation, and derive an exact two-pass streaming construction that avoids materializing the full fitted matrix. We further extend the framework to multimodal exemplar retrieval by incorporating visual representations into the conditioning and retrieval spaces. Experiments across ten textual datasets and four multimodal benchmarks with Llama-3.2-1B, MobileLLM-Pro, OpenFlamingo-3B, and Qwen3.5-2B, as well as end-to-end Raspberry Pi~5 deployment demonstrate that CoRA supports effective task-conditioned retrieval without retriever fine-tuning, backpropagation, or target-model calls.

Problem

• on-device in-context learning
• task-conditioned retrieval
• gradient-free
• limited resource constraints
• demonstration selection

Innovation

• gradient-free retrieval
• task-conditioned retrieval
• in-context learning
• low-rank compression
• on-device AI