Think Before You Retrieve: Learning Test-Time Adaptive Search with Small Language Models

TL;DR

This work tackles the mismatch between complex user information needs and static, one-shot retrieval by introducing Orion, a training framework that enables compact models (350M–1.2B) to learn test-time adaptive search. Through synthetic trajectory generation, turn-level reinforcement learning (GRPO), and structured inference-time beam search, Orion learns diverse, backtrack-capable search strategies rather than relying on ever-larger models. Empirical results show that Orion achieves competitive or superior performance across multiple benchmarks (e.g., 77.6% SciFact, 25.2% BRIGHT with 1.2B parameters) despite using only a fraction of the training data, and can outperform retrievers that are orders of magnitude larger on several tasks. The findings suggest that retrieval intelligence can emerge from learned search strategies and adaptive querying, offering a cost-effective path to high-performance information retrieval in production systems.

Abstract

Effective information retrieval requires reasoning over partial evidence and refining strategies as information emerges. Yet current approaches fall short: neural retrievers lack reasoning capabilities, large language models (LLMs) provide semantic depth but at prohibitive cost, and query rewriting or decomposition limits improvement to static transformations. As a result, existing methods fail to capture the iterative dynamics of exploration, feedback, and revision that complex user queries demand. We introduce Orion, a training framework that enables compact models (350M-1.2B parameters) to perform iterative retrieval through learned search strategies. Orion combines: (1) synthetic trajectory generation and supervised fine-tuning to encourage diverse exploration patterns in models, (2) reinforcement learning (RL) that rewards effective query refinement and backtracking behaviors, and (3) inference-time beam search algorithms that exploit the self-reflection capabilities learned during RL. Despite using only 3% of the training data available, our 1.2B model achieves 77.6% success on SciFact (vs. 72.6% for prior retrievers), 25.2% on BRIGHT (vs. 22.1%), 63.2% on NFCorpus (vs. 57.8%), and remains competitive on FEVER, HotpotQA, and MSMarco. It outperforms retrievers up to 200-400x larger on five of six benchmarks. These findings suggest that retrieval performance can emerge from learned strategies, not just model scale, when models are trained to search, reflect, and revise.

Figures (7)

• Figure 1: Overview of Orion. We illustrate two established query reformulation baselines alongside our proposed Orion framework. While query decomposition fails without corpus feedback and query rewriting yields static reformulations that ignore retrieval results, Orion performs tree-based exploration with structured reasoning spans, revising its strategy as it encounters conflicting evidence -- backtracking, refining, and generating new hypotheses to recover relevant information.
• Figure 2: We present further behavioral analysis of Orion-Large on BRIGHT: (a) illustrates the proportion of queries with unchanged rankings across turns, indicating repetitive search patterns and the inability to overcome search stagnation, while (b) shows search query length distribution, demonstrating our models generate relatively succinct search queries, while (c) measures backtracking behavior by counting queries where rankings ($r$) deteriorate then recover ($r_{i-1} > r_i < r_{i+1}$). Results for Orion-Small and Orion-Medium variants are provided in Appendix \ref{['app:search-behaviour-graphics']}.
• Figure 3: We demonstrate how successful queries distribute across search turns for different models of Orion-Large on BRIGHT. Results for Orion-Small and Orion-Medium variants are provided in Appendix \ref{['app:search-behaviour-graphics']}.
• Figure 4: Search behavior analysis across models: demonstrates how successful queries distribute across search turns for different models.
• Figure 5: Search behavior analysis across models: illustrates the proportion of queries with unchanged rankings across turns, indicating repetitive search patterns and the inability to overcome search stagnation.