Beyond Semantic Similarity: Learning Behavior-Aligned Retrieval for Tool-Augmented LLMs

Yixin Chen, Ying Xiong, Shangyu Wu, Yufei Cui, Xue Liu, Nan Guan, Chun Jason Xue


Abstract
Tool-augmented LLMs invoke external functions to extend their capabilities, but errors in the invocation decision, such as calling a tool when none is needed or omitting a needed call, can produce unreliable outputs and unnecessary cost. A lightweight remedy is to prepend retrieved examples so LLMs decide tool use in context. However, existing retrievers rank examples by semantic similarity alone. Lexically close or semantically close queries can require opposite behavior, so the retrieved examples may be behaviorally inconsistent and silently mislead the model. We propose Behavior-Aligned Retrieval (BAR), a backbone-agnostic training recipe that teaches a dense retriever a behavior-aware similarity, keeping semantically related candidates close only when their tool-use behavior is compatible. BAR does not predict invocation labels; instead, it ranks demonstrations while leaving the final tool-use decision to the LLM. Applied to multiple retrieval backbones, including BERT, Contriever, and Qwen3-Reranker, BAR consistently improves invocation reliability and reduces unnecessary API calls across 14 LLMs and 3 benchmarks.
Anthology ID:
2026.findings-emnlp.694
Volume:
Findings of the Association for Computational Linguistics: EMNLP 2026
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
Findings
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
14549–14571
Language:
URL:
https://aclanthology.org/2026.findings-emnlp.694/
DOI:
Bibkey:
Cite (ACL):
Yixin Chen, Ying Xiong, Shangyu Wu, Yufei Cui, Xue Liu, Nan Guan, and Chun Jason Xue. 2026. Beyond Semantic Similarity: Learning Behavior-Aligned Retrieval for Tool-Augmented LLMs. In Findings of the Association for Computational Linguistics: EMNLP 2026, pages 14549–14571, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
Beyond Semantic Similarity: Learning Behavior-Aligned Retrieval for Tool-Augmented LLMs (Chen et al., Findings 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.findings-emnlp.694.pdf