SeaKR: Self-aware Knowledge Retrieval for Adaptive Retrieval Augmented Generation

1 Introduction

• Retrieving for every query can be inefficient, and noisy retrieved knowledge may mislead when the LLM already knows the answer.The paper frames adaptive retrieval as a way to decide dynamically whether external knowledge is needed.Source: §1 · p. 1
• SEAKR uses self-aware uncertainty from LLM internal states to determine when retrieval is needed.The uncertainty is extracted from feed-forward network states at the last generated token and measured through consistency across multiple generations.Source: §1 · pp. 1–2
• For integration, SEAKR re-ranks snippets by uncertainty reduction and selects among reasoning strategies using generation uncertainty.Its strategies include direct generation and comprehensive reasoning over retrieved knowledge.Source: §1 · p. 2
• SEAKR substantially improves over existing adaptive RAG methods on complex QA benchmarks.The ablation study found that dynamically integrating retrieved knowledge brought more performance gain than self-aware retrieval.Source: §1 · p. 2

2 Related Work

• The related-work discussion covers retrieval-augmented generation and methods for analyzing LLM internal states.These are the two areas identified as the paper's related-work scope.Source: §2 · p. 2

2.1 Retrieval Augmented Generation

• Adaptive RAG methods decide whether to retrieve, but output-based decisions may be unreliable and retrieved knowledge is often neither re-ranked nor used to optimize reasoning paths.Prior methods trigger retrieval using token probabilities, prompted decisions, or a trained factuality classifier; retrieval-augmented reasoning includes iterative retrieval and decomposing questions into subquestions.Source: §2.1 · p. 2

2.2 Self-awareness in Internal States of LLMs

• LLM internal representations can reveal hallucinations and uncertainty, while decoding them into tokens loses information relevant to self-awareness.Prior work detects non-factual generations using internal-state classifiers or consistency across multiple generations.Source: §2.2 · p. 3

3 Self-Aware Knowledge Retrieval

• SEAKR combines a search engine, an LLM, and an uncertainty estimator in an iterative process that gathers knowledge and rationales before producing an answer.The estimator measures uncertainty at each stage of retrieval and generation.Source: §3 · p. 3

3.1 Self-aware Retrieval

• SEAKR triggers retrieval when uncertainty in the current context exceeds a threshold and uses uncertain tokens in a pseudo-generated rationale to form the search query.Retrieved snippets are added to the context for rationale generation; otherwise, generation continues without them.Source: §3.1 · p. 4

3.2 Self-aware Re-ranking

• SEAKR re-ranks multiple retrieved snippets by selecting the one that produces the least self-aware uncertainty.This prioritizes uncertainty reduction rather than query relevance alone, which may overlook conflicts with the LLM's intrinsic knowledge.Source: §3.2 · p. 4

3.3 Self-aware Reasoning

• For the final answer, SEAKR compares direct generation from gathered rationales with chain-of-thought reasoning over retrieved knowledge, choosing the lower-uncertainty strategy.Retrieval stops when the final-answer statement is generated or the maximum retrieval limit is reached.Source: §3.3 · p. 4

3.4 Self-aware Uncertainty Estimator

• SEAKR scores uncertainty using the determinant of a regularized Gram matrix computed from EOS-token hidden representations across multiple generations.The estimator uses internal-state consistency, avoiding dependence on natural-language expression of equivalent semantics.Source: §3.4 · pp. 4–5

4 Experiments

• The experiments compare SEAKR with commonly used baseline RAG methods on question answering, reporting F1 and exact match scores.The setup introduces the datasets and baselines and describes implementation details.Source: §4 · p. 5, §4.1 · p. 5, §4.2 · p. 6

4.1.1 Benchmark Datasets

• Complex QA uses 2WikiMultiHopQA, HotpotQA, and answerable IIRC for multi-hop questions needing multiple supporting facts; simple QA uses NaturalQuestions, TriviaQA, and SQuAD in an open-domain setting.SEAKR is tuning-free and uses a sampled subset of NaturalQuestions' training split for hyperparameter search; the paper follows IRCoT's official development set and DPR for simple QA.Source: §4.1.1 · p. 5

4.1.2 Baselines

• Baselines include non-adaptive CoT and IRCoT, plus adaptive Self-RAG, FLARE, and DRAGIN, which use different retrieval triggers and integration approaches.IRCoT retrieves at every reasoning step; Self-RAG is fine-tuned to signal retrieval and critique knowledge, while FLARE and DRAGIN use low-probability tokens to trigger retrieval.Source: §4.1.2 · p. 5

4.1.3 Implementation and Variable Control

• Implementation uses LLaMA-2-chat (7B), BM25 via Elastic Search, and Wikipedia dumped on December 20, 2018; simple QA limits search time to one for all methods.The study recalls N = 3 snippets, samples the EOS hidden representation 20 times, and searches for the uncertainty threshold on the development set.Source: §4.1.3 · pp. 5–6

4.2.1 Results on Complex QA

• SEAKR scores 36.0%, 39.7%, and 23.5% F1 on 2WikiMultiHopQA, HotpotQA, and IIRC, exceeding the best baselines by 6.0%, 5.5%, and 0.6%, respectively.The paper notes that the gain is smaller on IIRC, which requires numerical reasoning that SEAKR does not optimize.Source: §4.2.1 · p. 6, Table 1 · p. 6
• CoT reaches 22.3%, 27.5%, and 17.3% F1 on the three complex-QA datasets, while IRCoT also lags SEAKR, supporting adaptive retrieval over answering without retrieval or retrieving at every step.The paper attributes CoT's results to questions within LLMs' knowledge boundaries and reports that SEAKR outperforms Self-RAG, whose simple-QA fine-tuning data creates a distribution shift on complex QA.Source: §4.2.1 · p. 6

4.2.2 Results on Simple QA

• SEAKR achieves 63.1% F1 on TriviaQA and 36.5% F1 on SQuAD, leads the baselines on both, and is comparable to FLARE but behind Self-RAG on NaturalQuestions.The paper reports that the performance gap is less pronounced on simple QA, especially on NaturalQuestions and TriviaQA, because these tasks demand less knowledge integration.Source: §4.2.2 · pp. 6–7, Table 2 · p. 6
• Simple-QA performance gaps are less pronounced, especially on NaturalQuestions and TriviaQA, because single-step machine reading comprehension demands less knowledge integration.Source: §4.2.2 · pp. 6–7

5.1 Ablation Study

• Removing dynamic retrieval degrades performance on both complex and simple QA, consistent with retrieved information sometimes misleading the LLM.The ablation retrieves knowledge at every generation step without dynamically deciding when to retrieve.Source: §5.1 · p. 7
• Replacing self-aware re-ranking with the search engine’s first result undermines performance; re-ranking de-noises retrieved knowledge.The ablation selects the first search result instead of using self-aware uncertainty to choose knowledge.Source: §5.1 · p. 7
• Ablating self-aware re-ranking reduces performance more than removing self-aware retrieval, underscoring the contribution of knowledge integration.The ablation table notes that self-aware reasoning applies only to complex QA because simple QA does not require multiple retrievals.Source: §5.1 · p. 7, Table 3 · p. 7
• Fixed rationale-only and knowledge-only reasoning strategies both underperform SEAKR’s adaptive reasoning, which combines retrieved knowledge and acts as an ensemble.Self-aware reasoning integrates previously retrieved knowledge more effectively; the table’s reasoning ablation applies only to complex QA.Source: §5.1 · pp. 7–8, Table 3 · p. 7

5.4 Case Study

• In the second iteration, U(c) = −4.4 triggers retrieval; the second-ranked snippet reduces uncertainty and supplies critical information, unlike the top search result.The third retrieved snippet overlaps with the second and also has a relatively low uncertainty score.Source: Table 5 · p. 8, §5.4 · p. 8

5.2 Backbone LLMs

• SEAKR benefits from the stronger LLaMA-3 backbone, while alignment-tuned backbones achieve higher performance than their non-alignment-tuned counterparts.The comparison examines LLaMA-2-chat and LLaMA-3 instruct and chat variants.Source: Table 4 · p. 7, §5.2 · p. 8

5.3 Hyper-parameter Search

• The best Gram-determinant generation count k falls in [10–25]; the most indicative internal state is extracted from middle layer l=16; δ > −6 triggers retrieval, under which fewer than 80% of questions cannot be answered correctly.Figure 3 presents the hyper-parameter search results.Source: §5.3 · p. 8, Figure 3 · p. 8

• SEAKR extracts LLM self-aware uncertainty from internal states to invoke retrieval and dynamically integrate retrieved knowledge, outperforming existing adaptive baselines on complex and simple QA.Source: §6 · p. 9

• SEAKR requires access to LLM internal states, limiting its usability to open-sourced models; estimating uncertainty from outputs remains an area to explore.Source: Limitations · p. 9
• Evaluation focuses mainly on short-form QA, leaving tasks such as long-form QA and creative writing outside the reported coverage.Source: Limitations · p. 9
• Computing the Gram determinant requires 20 pseudo-generations and is costly, although vLLM parallel inference makes their latency roughly that of one generation.Source: Limitations · p. 9
• Limited compute prevented experiments with backbones larger than 8 billion parameters; the scaling comparison instead uses LLaMA-2 (7B) and LLaMA-3 (8B).Source: Limitations · p. 9
• The authors note that advanced information-retrieval methods could surpass SEAKR’s self-aware re-ranking in the future.Source: Limitations · p. 9

Ethical Considerations

• The work’s stated intention is to improve LLM trustworthiness through retrieval-augmented generation aimed at increasing factual correctness.SEAKR is developed on open-source LLMs, and the authors plan to release its code.Source: Ethical Considerations · p. 9
• SEAKR’s self-aware uncertainty estimator could be used as an adversarial training signal, and integrating LLMs with information-retrieval systems could enable automated cyber manhunt.The authors hope open-source development and code release help monitor and prevent misuse.Source: Ethical Considerations · p. 9
• The paper uses six QA benchmark datasets and credits the creators of its datasets and reproduced methods; all cited artifact licenses allow academic use.The listed datasets are 2WikiMultiHopQA, HotpotQA, IIRC, NaturalQuestion, TriviaQA, and SQuAD.Source: Ethical Considerations · p. 9

A.1 Case Study for Self-aware Retrieval

• Case #1 triggers retrieval when generation misses the predefined threshold; Cases #2 and #3 bypass potentially redundant retrieval after confident, correct generation.Case #3 also shows self-aware retrieval during multi-step reasoning with nonempty knowledge and rationale buffers.Source: §A.1 · p. 13

A.2 Case Study for Self-aware Re-ranking

• After retrieval, SEAKR performs pairwise re-ranking to identify the passage used for subsequent reasoning.The example contrasts three retrieved passages, only one of which contains the relevant parent-company information.Source: §A.2 · p. 13
• SEAKR’s self-aware uncertainty prioritizes Knowledge #2, which says FastJet Tanzania was founded as a subsidiary of a Kenyan company.Knowledge #1 and #3 instead list its current headquarters and major shareholder, respectively.Source: §A.2 · p. 13

A.3 Case Study for Self-aware Reasoning

• The knowledge-based strategy yields a better uncertainty score (−6.20) than the average rationale score (−4.73) in the example.The initial rationale falsely identifies Stormbreaker as a fantasy film; reasoning from retrieved evidence instead uses information about The Spiderwick Chronicles.Source: §A.3 · p. 13

B.1 Self-aware Retrieval

• At each reasoning iteration, SEAKR evaluates a pseudo-generation’s uncertainty and retrieves when it fails the predefined threshold.The pseudo-generation is limited to a single next step using a period stop token.Source: §B.1 · p. 13

B.2 Self-aware Re-ranking

• When direct-generation uncertainty fails the threshold, SEAKR retrieves and pairwise re-ranks a pseudo-generation.The appendix specifies a period stop token for this generation.Source: §B.2 · p. 13

B.3 Self-aware Reasoning

• The final stage selects either the answer extracted from the last rationale or a full chain-of-thought response using all retrieved passages.Both response groups use a newline stop token.Source: §B.3 · p. 13

B.4 In context learning examples

• The paper uses the same in-context examples for simple QA and different examples for each multihop QA dataset following IRCoT.The cited multihop datasets are 2WikiMultiHopQA, HotpotQA, and IIRC.Source: §B.4 · p. 14

C Datasets And Settings

• Hyperparameter search uses 3,000 Natural Questions training samples.Dataset statistics are summarized in Table 6.Source: C Datasets And Settings · p. 14
• The implementation modifies vLLM 0.4.2 to compute uncertainty from internal states and uses PyTorch 2.3.0 and Elasticsearch 7.17.9.Source: C Datasets And Settings · p. 14
• Tables 7–9 provide additional examples for self-aware retrieval, re-ranking, and reasoning, respectively.Source: Table 7 · p. 15, Table 8 · p. 16, Table 9 · p. 17
• Figures 4–7 show examples for simple QA, 2WikiMultiHopQA, HotpotQA, and IIRC, respectively.Source: Figure 4 · p. 18, Figure 5 · p. 19, Figure 6 · p. 20