PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning When to Think: Utility-Guided Adaptive Reasoning for Efficient Text Retrieval

Yingtao Luo, Kangyi Zhao, Shiyi Du, Yubo Li

arXiv:2610.21030v1cs.IRcs.CL

TL;DR

Reasoning can refine retrieval representations, but additional steps may introduce semantic drift, and existing methods offer limited guidance over intermediate refinements. AURA supervises intermediate states by retrieval specificity and estimates their utility to adapt reasoning depth. It achieves the strongest average BRIGHT performance among compared methods while reducing reasoning cost, and MS MARCO results suggest limited in-domain degradation.

  • Problem

    Existing retrieval methods generally supervise reasoning through final retrieval outcomes, providing little guidance over intermediate refinements, while further reasoning can introduce semantic drift that weakens retrieval discrimination.

    Source: §1 · p. 1

  • Method

    AURA uses relevance labels and retrieval candidates to supervise intermediate states by retrieval specificity, then predicts state utility to stop reasoning when utility ceases to improve.

    Source: §1 · p. 2

  • Results

    AURA achieves the strongest average BRIGHT performance among compared methods while substantially reducing reasoning cost; MS MARCO results suggest limited in-domain degradation.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    Utility-aware supervision, rather than denser supervision or trajectory filtering, accounts for the gains, and retrieval specificity closely tracks nDCG@10.

    Source: §5 · p. 9

    Evaluation focuses primarily on reasoning-intensive English retrieval, leaving broader validation on web search, multilingual retrieval, and generation tasks for future work.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Reasoning-conditioned embeddings can improve retrieval by progressively refining semantic constraints. However, additional reasoning may introduce over-generalized or semantically noisy abstractions that weaken retrieval specificity. Existing methods largely supervise reasoning through final retrieval outcomes, providing limited guidance over the refinement process itself. We propose AURA, a framework that directly supervises intermediate reasoning states according to how well they separate relevant documents from hard distractors. AURA further trains a lightweight utility estimator over prefix-conditioned embeddings to adapt reasoning depth at inference time. On the reasoning-intensive BRIGHT benchmark, AURA-4B achieves 34.7 average nDCG@10 while using 2.8 reasoning steps and 96 generated tokens per query. Under matched deployment conditions, adaptive stopping reduces latency from 4.6 to 0.76 seconds relative to full-trajectory reasoning. Results on MS MARCO further suggest that this specialization need not substantially degrade conventional in-domain retrieval.

1 Introduction

AURA treats retrieval reasoning as query-dependent refinement: added steps can sharpen discrimination, but can also cause semantic drift, making fixed-depth reasoning unreliable. It supervises intermediate states and estimates their utility to allocate reasoning at inference time.

1 Introduction

  • Existing systems generally apply more reasoning as if it were beneficial, although semantic drift and overly abstract reformulations can weaken discrimination, increase latency, and add computation cost.Some refinements sharpen specificity by identifying useful constraints or disambiguating entities.Source: §1 · p. 1
  • Retrieval specificity improves on average in early BRIGHT reasoning steps, but query-level variability rises with depth and over 50% of queries have passed their optimum by Step 10.The figure tracks mean specificity, query-level variability, and the fraction past each query's optimal step.Source: §1 · p. 1, Figure 1 · p. 1
  • AURA frames retrieval-oriented reasoning as utility-guided computation allocation, measuring usefulness by its effect on retrieval discrimination rather than trajectory length or textual sophistication.The approach targets heterogeneous query needs by adapting reasoning depth.Source: §1 · p. 2
  • AURA supervises intermediate reasoning states by estimating whether each refinement improves discrimination between relevant documents and retrieval candidates, avoiding semantically noisy refinements.Training uses relevance labels and retrieval candidates to assess each step's retrieval specificity.Source: §1 · p. 2
  • A lightweight utility estimator predicts each current reasoning state's retrieval utility, allowing online stopping when utility ceases to improve without exhaustively expanding the trajectory.The estimator addresses the fact that retrieval specificity is unavailable at inference time.Source: §1 · p. 2
  • AURA outperforms compared reasoning-conditioned methods across 12 BRIGHT domains, remains comparable to a strong reasoning-trained retriever on MS MARCO, and offers a better performance–cost trade-off under matched inference conditions.The results suggest no substantial in-domain degradation when extensive reasoning is unnecessary.Source: §1 · p. 2

2 Related Work

Dense retrieval has moved from single-pass embeddings toward reasoning-based retrieval, including reasoning-conditioned representations. Yet prior approaches provide limited intermediate-state supervision and adaptive allocation, gaps AURA addresses.

2.1 Dense Retrieval and Embedding Models

  • Dense retrievers map queries and documents into a shared embedding space for vector-similarity retrieval, but most remain single-pass encoders without explicit semantic refinement or intermediate reasoning.The subsection traces embedding models from contrastive approaches through instruction-tuned and LLM-based systems.Source: §2.1 · p. 3

2.2 Reasoning for Retrieval

  • Reasoning-based retrieval spans query rewriting, decomposition, iterative retrieval, reranking, and reasoning-conditioned embeddings, with some methods scaling inference by sampling multiple trajectories.ReasonIR trains on reasoning trajectories and hard negatives for reasoning-intensive benchmarks.Source: §2.2 · p. 3
  • Reasoning can hurt retrieval through semantic drift, while existing methods either supervise only final outcomes or select complete trajectories with best-of-k, leaving intermediate refinement and adaptive allocation unsupervised.AURA is presented as addressing both gaps.Source: §2.2 · p. 3

3 Methodology

AURA builds retrieval embeddings at intermediate reasoning steps, supervises them according to how well they separate relevant documents from distractors, and estimates utility to adapt reasoning depth at inference. Its training combines retrieval and language-modeling objectives while limiting gradient interference between them.

3.1 Unified Reasoning Embedding

  • At each sentence boundary, AURA forms an embedding from the final-layer hidden state at the last reasoning-token position, projected into the embedding space.The reasoning trajectory is parsed at sentence boundaries, and the prefix through each sentence conditions its embedding.Source: §3.1 · p. 3
  • AURA scores each intermediate state by its separation of relevant documents from hard distractors and applies weighted contrastive supervision across reasoning prefixes.The state weight emphasizes refinements with larger specificity gains and reduces supervision on saturated or retrieval-harmful states.Source: §3.1 · pp. 3–4
  • Training combines contrastive retrieval and autoregressive language modeling, with a warm-up and staged utility-aware optimization.Contrastive gradients are stopped before the current embedding extraction point, reducing interference with generation while preserving language-model gradient flow across the reasoning sequence.Source: Figure 3 · p. 4, §3.1 · pp. 4–5

3.2 Inference-Time Utility Estimation

  • Because retrieval specificity depends on relevance labels and distractors unavailable at inference, AURA trains a lightweight estimator to predict each state's utility.A pairwise ranking loss targets the ordering of state utilities rather than absolute score calibration.Source: §3.2 · p. 5
  • At inference, reasoning stops after predicted utility fails to improve for a set number of consecutive steps, and the highest-utility state supplies the retrieval embedding.The estimator operates on prefix-conditioned embeddings and the stopping rule relies on relative state ordering within a query trajectory.Source: §3.2 · p. 5

4 Experiments

AURA is evaluated on BRIGHT using reasoning trajectories constructed from ReasonIR data, with comparisons spanning external query rewriting and endogenous reasoning. The experiments also specify the training schedule and implementation settings.

4.1 Training Data Construction

  • ReasonIR supplies 101K query-positive-negative samples; eight DeepSeek-R1 trajectories per query are segmented into states, leaving about 545K samples after filtering.Filtering retains 67.5% of trajectories, or about 5.4 per query on average.Source: §4.1 · p. 5

4.2 Training Procedure

  • AURA warms up the unified model before constructing specificity and utility labels from frozen embeddings, then enables utility-aware weighting during joint optimization.The utility estimator is a 3-layer, 10M-parameter MLP trained on the frozen embeddings.Source: §4.2 · p. 5

4.3 Implementation Details

  • Experiments use Qwen3-4B-Instruct on 8 NVIDIA H100 GPUs and report nDCG@10, generated tokens, reasoning steps, and matched-setting latency.Convergence consumes approximately 100 GPU hours.Source: §4.3 · p. 5

4.4 Baselines

  • The baselines distinguish retrieval without inference-time reasoning, external query rewriting, and AURA’s endogenous reasoning with utility-guided depth allocation.Qwen3-4B-Instruct used as an external query rewriter does not improve any evaluated retriever and lowers average BRIGHT nDCG@10 by 0.1–2.7 points.Source: §4.4 · pp. 5–6

4.5 Main Results

On BRIGHT, AURA achieves the highest average nDCG@10 among the compared systems, outperforming both external-reasoning baselines and the strongest baseline without inference-time reasoning. External query rewriting also improves Setting II over Setting I across retrievers.

4.5 Main Results

  • Setting II consistently outperforms Setting I across retrievers, with external GPT-4 query expansion producing particularly large improvements for lexical retrievers such as BM25.Source: §4.5 · p. 6
  • 34.7 average nDCG@10: AURA beats DIVER-Retriever-4B with GPT-4 query expansion at 32.1, a +2.6-point gain.AURA generates reasoning within the retrieval model, without an external reasoning model.Source: §4.5 · p. 6
  • 34.7 versus 28.9: AURA exceeds the strongest baseline without inference-time reasoning by 5.8 points, or 20.1% relative.Source: §4.5 · p. 6

4.6 Alignment with Corpus-Level Retrieval

The paper checks whether retrieval specificity aligns with full-corpus ranking quality and evaluates AURA on conventional MS MARCO retrieval. Specificity closely tracks corpus-level nDCG@10, while AURA remains comparable to a larger reasoning-trained retriever across the reported splits.

4.6 Alignment with Corpus-Level Retrieval

  • 0.97 average query-level correlation links retrieval specificity across intermediate states with full-corpus nDCG@10; specificity-based oracle selection is nDCG@10-optimal for 98.1% of queries.The evaluation uses full-corpus retrieval for every available intermediate state in naturally varying-length trajectories.Source: §4.6 · pp. 6–7

4.7 Retrieval on Traditional IR Benchmarks

  • AURA-4B is comparable to RaDeR-7B on MS MARCO DEV, TREC-DL’19, and TREC-DL’20, despite using a smaller backbone.These held-out-query splits provide an in-domain check; broader general-purpose evaluation is left to future work.Source: §4.7 · p. 7

4.8 Ablation Study

The ablations test utility supervision, stopping, negative aggregation, trajectory filtering, gradient isolation, utility estimation, and confidence-based stopping. They indicate that utility-aware state differentiation and adaptive termination matter, while filtering provides only a small benefit.

4.8 Ablation Study

  • Negative aggregation has minimal impact: mean top-3, max, and mean over all negatives score 34.7, 34.6, and 34.4, respectively.Mean top-3 is the default aggregation.Source: §4.8 · p. 7
  • 34.7 for AURA versus 33.6 for uniform intermediate-state weighting: utility-aware differentiation adds 1.1 nDCG@10 beyond denser contrastive supervision.Uniform weighting itself improves 2.4 points over final-state-only supervision at 31.2; monotonic labels score 33.2 and random labels 31.9.Source: §4.8 · p. 7
  • 34.7 with gradient isolation falls to 29.8 when isolation is removed, whereas retaining trajectories without improving intermediate states changes performance only from 34.7 to 34.6.The ablation supports gradient isolation; filtering contributes only a small benefit.Source: §4.8 · p. 7
  • 34.7 from the utility estimator is 0.3 points below oracle selection at 35.0; random state selection reaches 32.3.The oracle uses ground-truth relevance labels, unavailable to the estimator at inference.Source: §4.8 · p. 7
  • 32.8 nDCG@10 at 3.2 average steps for the strongest LLM-confidence rule trails AURA’s 34.7 at 2.8 steps.The comparison tests whether mean token entropy can identify the most retrieval-discriminative state.Source: §4.8 · p. 8
  • 34.7 with patience p=1 outperforms full-trajectory reasoning at 33.0; patience p=3 falls to 33.7, while greedy peak stopping reaches 34.6.The p=1 result is 1.7 points above no early stopping.Source: §4.8 · p. 8

4.9 Inference-Time Compute Efficiency

AURA’s utility-guided stopping reduces inference cost while improving retrieval performance, and utility-guided search can further improve the performance–compute trade-off on BRIGHT.

4.9 Inference-Time Compute Efficiency

  • 96 generated tokens (2.8 steps) per query reduce latency to 0.76 seconds and improve nDCG@10 to 34.7, versus 332 tokens (9.5 steps), 4.6 seconds, and 33.0 without early stopping.Under matched Qwen3-4B backbone, hardware, batch size, and software, this is a 3.5× token reduction and approximately 6× speedup; the utility estimator adds approximately 1 millisecond per query.Source: §4.9 · p. 8, Table 5 · p. 9
  • Best-of-k samples complete trajectories and selects the highest-utility one, whereas greedy expansion continues from the highest-utility candidate refinement at each step.AURA evaluates both strategies under increasing inference budgets.Source: §4.9 · pp. 8–9, Figure 4 · p. 9
  • 34.7 (k = 1) rises to 37.2 (k = 64) with greedy expansion, compared with 36.7 under Best-of-k on the same search scale.Increasing search budget improves performance for both strategies, while greedy expansion achieves better compute efficiency.Source: §4.9 · p. 9, Figure 4 · p. 9

4.10 Qualitative Analysis

The insects-and-LEDs example shows why AURA can stop before a reasoning trajectory ends: later, valid details may weaken document discrimination.

4.10 Qualitative Analysis

  • The example’s early reasoning distinguishes proximate from ultimate causation, improving separation from hard negatives.Later reasoning adds evolutionary and sensory-ecology background that repeats encoded constraints or generalizes toward hard negatives.Source: §4.10 · p. 9
  • AURA stops once the discriminative constraint is encoded, even though the discarded later steps can be semantically valid.Textual informativeness alone does not determine retrieval utility.Source: §4.10 · p. 9
  • The example illustrates that stopping should depend on document discrimination rather than textual detail.The added background is useful but can repeat constraints or broaden the query toward hard negatives.Source: §4.10 · p. 9

5 Conclusion

AURA supervises intermediate reasoning states and stops when predicted retrieval utility ceases to improve. The conclusion reports stronger BRIGHT performance at lower reasoning cost, limited MS MARCO degradation, and further gains from utility-guided greedy expansion.

5 Conclusion

  • AURA supervises intermediate reasoning states and stops generation when predicted retrieval utility ceases to improve.The framework’s conclusion reports strongest average BRIGHT performance among compared methods while substantially reducing reasoning cost, without an external inference-time reasoner.Source: §5 · p. 9
  • Controlled ablations and full-corpus analysis attribute gains to utility-aware supervision rather than denser supervision or trajectory filtering, and report that retrieval specificity closely tracks nDCG@10.MS MARCO results suggest limited in-domain degradation.Source: §5 · p. 9
  • Utility-guided greedy expansion further improves the performance–compute trade-off.Source: §5 · p. 9

Limitations

AURA remains slower than single-pass retrieval, its utility labels may become less accurate as the embedding space evolves, and its evaluation scope is limited.

Limitations

  • AURA remains slower than single-pass retrieval because reasoning is generated autoregressively.Source: Limitations · p. 9
  • Utility labels are constructed once from a warmed-up checkpoint and may become less accurate during subsequent optimization as the embedding space evolves.Source: Limitations · p. 9
  • Evaluation focuses primarily on reasoning-intensive English retrieval; broader validation is needed for web search, multilingual retrieval, and generation tasks.The paper identifies this validation as necessary to establish generality and characterize the remaining retrieval–generation trade-off.Source: Limitations · p. 9

Related papers

Loading 2610.21030v1…