@inproceedings{koul-etal-2026-robust,
    title = "Robust In-Context Selection via Online Learned Position-Corrected Attention",
    author = "Koul, Deeksha  and
      Kumar, Gaurav  and
      Sabale, Yash  and
      Sarawagi, Sunita",
    editor = "Liakata, Maria  and
      Moreira, Viviane P.  and
      Zhang, Jiajun  and
      Jurgens, David",
    booktitle = "Findings of the {A}ssociation for {C}omputational {L}inguistics: {ACL} 2026",
    month = jul,
    year = "2026",
    address = "San Diego, California, United States",
    publisher = "Association for Computational Linguistics",
    url = "https://preview.aclanthology.org/ingest-acl/2026.findings-acl.1743/",
    pages = "34919--34930",
    ISBN = "979-8-89176-395-1",
    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."
}