Robust In-Context Selection via Online Learned Position-Corrected Attention

ACL ARR 2026 January Submission10594 Authors

Keywords: in-context retrieval, tool selection

Abstract:

Large Language Models (LLMs) are often deployed in tasks that require selecting an item from a long list provided in the model’s context. LLMs’ native selection behavior is brittle: predictions are sensitive to the surface form of the identifiers, their placement within the context, and the ordering of candidate items. We present OLR-Heads, a robust method for list selection that harnesses attention patterns available from a single forward call on the LLM. OLR-Heads learns the logic for item selection using a few in-context examples, and a simple online position-debiasing mechanism to correct attention distortion. Across multiple database and tool selection benchmarks, OLR-Heads consistently improves selection performance over direct generation and prior attention-based methods, while remaining robust to prompt variations and item ordering.
The LLM's KV cache states are unaffected, and can be reused for subsequent response generation. In contrast, existing approaches either entail additional LLM calls, or task-specific offline learning, or position debiasing methods that modify the attention or encoding rendering the KV states unusable for subsequent generation.

Paper Type: Long

Research Area: Language Models

Research Area Keywords: in-context retrieval, tool selection

Contribution Types: NLP engineering experiment

Languages Studied: english

Submission Number: 10594