Source-linked AI summary
RAT: Retrieval Augmented Thoughts Elicit Context-Aware Reasoning in Long-Horizon Generation
Zihao Wang, Anji Liu, Haowei Lin, Jiaqi Li, Xiaojian Ma, Yitao Liang
TL;DR
Long-horizon generation requires context-aware reasoning, but intermediate chain-of-thought steps can contain factual errors and task prompts may not reveal which information later steps need. RAT retrieves information using the task and thought context to revise each step progressively; it reports improvements across code, mathematics, embodied planning, and creative writing. The approach remains dependent on the base model's reasoning and in-context learning or retrieval capabilities.
Problem
Long-horizon tasks need multi-step, context-aware reasoning, yet intermediate thoughts can hallucinate and a task prompt may not reveal which facts later steps require.
Source: §1 · p. 1, §2.1 · p. 3Method
RAT generates zero-shot thoughts, retrieves information relevant to the task and current and past thoughts, and revises each thought step progressively.
Source: §1 · p. 1, §2.2 · p. 4Results
RAT improves results across code, mathematical reasoning, embodied planning, and creative writing, including 2.96 times on Minecraft executability and +51.94% on plausibility.
Source: §1 · p. 2Takeaways & Limitations
The paper reports that zero-shot retrieval-guided revision offers advantages over vanilla CoT, RAG, and other baselines on the evaluated long-horizon tasks.
Source: §5 · p. 10RAT relies on the base LLM's chain-of-thought reasoning and in-context learning or RAG capabilities, so its effect on smaller, weaker models such as GPT-2 is uncertain.
Source: Limitations · p. 10
Abstract
from arXiv · showhide
We explore how iterative revising a chain of thoughts with the help of information retrieval significantly improves large language models' reasoning and generation ability in long-horizon generation tasks, while hugely mitigating hallucination. In particular, the proposed method -- *retrieval-augmented thoughts* (RAT) -- revises each thought step one by one with retrieved information relevant to the task query, the current and the past thought steps, after the initial zero-shot CoT is generated. Applying RAT to GPT-3.5, GPT-4, and CodeLLaMA-7b substantially improves their performances on various long-horizon generation tasks; on average of relatively increasing rating scores by 13.63% on code generation, 16.96% on mathematical reasoning, 19.2% on creative writing, and 42.78% on embodied task planning. The demo page can be found at https://craftjarvis.github.io/RAT
1. Introduction
Long-horizon tasks make factual errors in intermediate reasoning especially consequential, while ordinary retrieval and chain-of-thought prompting each leave challenges. RAT retrieves information using the task and thought steps, then revises thoughts progressively; it improves results across code, mathematics, planning, and creative writing.
1. Introduction
- Hallucinations in reasoning paths are a particular concern for zero-shot CoT and long-horizon tasks requiring multi-step, context-aware reasoning.Factually valid intermediate thoughts can be critical to completing these tasks successfully.Source: §1 · p. 1
- RAT uses the task prompt and initial zero-shot chain of thought to retrieve relevant information, then revises the current thought step using retrieved information and past revised steps.The response is generated progressively from the revised thoughts.Source: §1 · p. 1
- Code-generation scores improve +20.94% on HumanEval, +18.89% on HumanEval+, +14.83% on MBPP, and +1.86% on MBPP+.The introduction reports advantages over vanilla CoT and RAG approaches.Source: §1 · p. 2
- Minecraft planning improves 2.96 times on executability and +51.94% on plausibility, while creative-writing human score improves +19.19%.The reported improvements are over the paper's baselines.Source: §1 · p. 2
- Mathematical reasoning improves +8.36% on GSM8K and +31.37% on GSMHard.The results are reported for RAT combined with the evaluated LLMs.Source: §1 · p. 2
2. Retrieval Augmented Thoughts
Long-horizon generation needs both factual information and intermediate reasoning, but a task prompt alone may not reveal which facts later steps will require. RAT first produces zero-shot thoughts, retrieves with the task and thought context, and revises each step before generating the response progressively.
2. Retrieval Augmented Thoughts
- A task prompt alone may not provide enough clues to retrieve facts needed at later reasoning and generation steps, limiting conventional one-shot RAG for complex tasks.The section frames relevant retrieval and effective correction of reasoning steps as two open questions.Source: §2 · p. 2, §2.1 · p. 3
2.1. Preliminary
- RAG supplies relevant texts, while chain-of-thought prompting supplies intermediate reasoning steps that act as a task scratch space.Without direct supervision, intermediate thoughts can still err because of missing domain knowledge or hallucinations.Source: §2.1 · pp. 3–4
2.2. Our Approach
- RAT begins with zero-shot step-by-step thoughts, then uses the task prompt and current and past thoughts to retrieve documents and revise the current step.The retrieved documents are prepended to the prompt; revised thoughts can form the response directly or guide step-by-step generation of a complete response.Source: §2.2 · p. 4
- The authors hypothesize that thought context improves retrieval clues, intermediate questions can be easier to retrieve for than final answers, and targeted stepwise revision can avoid disturbing correct steps.They point to ablation studies as quantitative evidence for these design choices.Source: §2.2 · p. 5
3. Experiments
RAT is evaluated on long-horizon code, math, creative-writing, and Minecraft planning tasks, with comparisons across retrieval strategies and reasoning modes. The reported results show gains across most benchmarks, while case analyses and ablations examine how retrieval and progressive reasoning contribute.
3.1. Experimental Setups
- The experiments cover code, multi-step math, creative writing, and 100 Minecraft planning tasks, using pass rates, accuracy, and human ratings.The settings include zero-shot prompts, multiple baselines, and external search resources for creative writing, code, math, and Minecraft.Source: §3.1 · pp. 5–6
3.2. Results
- RAT outperforms other methods on most benchmarks; HumanEval pass@1 improves by up to 20.94% and pass@5 by up to 25.68% relative to base models.The evaluation includes code generation, mathematical reasoning, embodied planning, and creative writing.Source: §3.2 · pp. 6–7, Table 1 · p. 6, Table 2 · p. 6
3.3. Case Analysis
- In Minecraft, RAT progressively retrieves knowledge about plan items, addressing procedural errors and RAG’s tendency to miss intermediate items.The case analysis contrasts a CoT error identifying crafting-table materials as four wood blocks rather than four planks with question-only retrieval focused on the final target.Source: §3.3 · p. 7
3.4. Ablation Study
- RAT’s iterative retrieval-query refinement outperforms question-only RAG-1 and full-thought CoT+RAG on HumanEval pass@1 and pass@5.On HumanEval, the reported gains over the baseline are 8.7 percentage points for pass@1 and 7.9 percentage points for pass@5.Source: §3.4 · p. 8
- Causal reasoning and retrieval outperform CoT and non-causal retrieval on HumanEval pass@1 and pass@5.The causal approach exceeds CoT by 11.9 percentage points in pass@1 and 4.6 percentage points in pass@5 on HumanEval.Source: §3.4 · p. 8
3.5. Robustness of RAT
- RAT’s gains extend across model scales and zero-shot tasks, but CodeLlama-7b RAT degrades on MBPP+.The authors attribute the largest GPT-4 improvement to its stronger in-context learning from retrieved text; math improvements are reported as 18.44% on average with GPT-3.5 and 10.26% with GPT-4.Source: §3.5 · p. 9
4. Related Works
Prior work uses retrieval to ground generation and reasoning methods to improve multi-step problem solving. RAT combines these directions by retrieving from evolving reasoning and draft answers rather than relying on a single question query.
Retrieval-augmented Generation
- RAG guides generation with retrieved knowledge across tasks including code generation, question answering, and creative writing.It can interact with external knowledge without updating model parameters.Source: §4 · p. 9
Reasoning-enhanced RAG
- Reasoning-enhanced RAG methods use CoT for retrieval, iteratively search for premises, or retrieve from a final-step generation, while RAT retrieves autoregressively from reasoning and draft answers.The passage contrasts these prior methods’ single-query retrieval with RAT’s approach.Source: §4 · p. 9
Language Model for Reasoning
- Prompting and feedback methods improve reasoning through structured steps or expanded reasoning paths, whereas RAT accesses external knowledge through RAG.Examples include Selection-Inference, STaR, Scratchpad, Tree-of-Thought, and Graph-of-Thought.Source: §4 · pp. 9–10
5. Conclusion
The conclusion presents RAT as a zero-shot approach that combines CoT with retrieval to revise thoughts progressively. It reports advantages over CoT, RAG, and other baselines across several long-horizon tasks.
5. Conclusion
- RAT combines chain-of-thought prompting and retrieval-augmented generation for long-horizon reasoning and generation.The method revises zero-shot thoughts using retrieval and progressively revises thoughts while generating the response.Source: §5 · p. 10
- RAT demonstrates advantages over vanilla CoT, RAG, and other baselines on code, math, embodied planning, and creative writing tasks.Source: §5 · p. 10
- RAT uses thoughts as retrieval queries and progressively revises them while generating the response.Source: §5 · p. 10
Limitations
RAT has three stated limitations: it depends on base-model reasoning and retrieval capabilities, on the quality of retrieved knowledge, and on a step-by-step CoT structure that may not suit every question.
Limitations
- RAT's performance depends on the base LLM's chain-of-thought reasoning and in-context learning or RAG capability; its effect on weaker models such as GPT-2 is uncertain.The work does not train the model, so applying RAT does not change the base LLM's capability.Source: Limitations · p. 10
- Irrelevant retrieved knowledge may be unhelpful, while maintaining and searching a large knowledge base can be expensive and reduce retrieval precision.The authors identify building and evaluating knowledge bases for efficient, effective retrieval as future work.Source: Limitations · p. 10
- RAT follows step-by-step CoT, which may be redundant for straightforward questions and may not match questions needing more complex structures such as tree-of-thoughts.The authors identify exploring better reasoning methods for LLMs as future work.Source: Limitations · p. 10
Ethics Statement
The authors report no significant ethical considerations associated with the work and state that most datasets and models are publicly accessible, with exceptions.
Ethics Statement
- All datasets and models are publicly accessible except OpenAI's GPT series and text embedding APIs.Source: Ethics Statement · p. 10
- The authors report no significant ethical considerations associated with this work.They state that RAT can improve LLM generation and reduce hallucinations.Source: Ethics Statement · p. 11
A. Task Details
The appendix describes benchmark tasks and evaluation setups for code generation, mathematical reasoning, and Minecraft embodied planning, including their metrics, baselines, and retrieval sources.
A.1. Code Generation
- Code generation is evaluated on HumanEval, HumanEval+, MBPP, and MBPP+, using pass@k to measure code that executes and passes tests.The study uses DIRECT, one- and five-shot RAG, and zero-shot CoT baselines; code outputs are post-processed for sandbox execution, and codeparrot/github-jupyter supplies retrieval data.Source: §A.1 · pp. 15–16
A.2. Mathematical Reasoning
- Mathematical reasoning uses GSM8K and GSM-HARD, with the first 200 samples from each test set evaluated using PAL and accuracy.Each question receives three execution attempts, whose average score is recorded; baselines are DIRECT, CoT, and one- and five-shot RAG.Source: §A.2 · p. 16
A.3. Embodied Planning
- Minecraft planning tests plans for obtaining items in survival mode, from simple goals to complex crafting objectives, on 100 randomly selected tasks.Prompts follow DEPS and evaluation uses MC-TextWorld executability plus human plausibility ratings, which assess crafting details that executability does not capture.Source: §A.3 · pp. 16–17
A.4. Creative Writing
Creative writing evaluates open-ended generation tasks with web retrieval and human ratings. The prompt details describe initial-answer, query-generation, and answer-revision prompts.
A.4. Creative Writing
- Creative-writing tasks include a survey of RAG methods, a description of Jin-Yong's life, and a timeline-based summary of the American Civil War.Three variants were created for each task, and Google search results were converted to markdown to provide retrieval context.Source: §A.4 · p. 17
- Creative-writing outputs are compared with DIRECT and one- and five-shot RAG using expert human evaluation and TrueSkill ratings.The framework assesses quality, accuracy, relevance, and innovativeness.Source: §A.4 · pp. 17–18
B. Prompt Details
- The prompt sequence generates an initial answer, creates an open-search query where needed, then revises the answer against retrieved material while preserving its structure.The query prompt seeks a short, specific search query, and the revision prompt asks for corrections or missing details supported by retrieved text.Source: B. Prompt Details · pp. 18–19
C. TrueSkill Evaluation Framework
Human evaluations compare two methods’ responses to the same instruction, with annotators rating relevance, accuracy, completeness, readability, and executability. Labels are aggregated using TrueSkill scores.
C. TrueSkill Evaluation Framework
- More than 10 professional annotators label tasks in each group; “Tie” and “Both are bad” count as ties.The four choices are “A is better”, “B is better”, “Tie” and “Both are bad”.Source: C. TrueSkill Evaluation Framework · p. 19
- The TrueSkill package calculates win rates and scores, with each method starting at a default score of 25.Ties and “Both are bad” responses are counted as ties.Source: C. TrueSkill Evaluation Framework · p. 19
- For each round, the interface provides a question with corresponding answers from two anonymous models, asks users to vote for the better answer (or mark both bad), then offers a new round.Source: C. TrueSkill Evaluation Framework · p. 19
- Evaluation criteria cover relevance, accuracy, completeness, readability, and executability.The guidance illustrates accuracy and completeness with Minecraft crafting examples.Source: C. TrueSkill Evaluation Framework · p. 19
D. Disscussions on Benchmark Contamination
The benchmark-contamination discussion describes preprocessing to remove matching or overly similar code from the search library and additional CodeLlama fine-tuning tests. The reported comparison says direct fine-tuning with the code corpus damages benchmark performance, whereas RAT in-context generation performs better.
D. Disscussions on Benchmark Contamination
- Preprocessing removes direct matches and overly similar code snippets from the search vector library to keep code-generation evaluation uncontaminated by existing solutions.The procedure is attributed to Guo et al. (2024).Source: D. Disscussions on Benchmark… · pp. 19–20
- Direct CodeLlama-7b-Python fine-tuning on the code corpus damages code-generation benchmark performance, while RAT in-context generation performs better.The comparison is described as a benchmark-contamination study.Source: Table 5 · p. 20
- To further examine potential benchmark contamination, the authors finetuned CodeLLaMA-7B-Python using Table 5’s code corpus; direct finetuning damaged code-generation benchmark performance, whereas RAT-based in-context generation yielded better performance.Source: D. Disscussions on Benchmark… · p. 20, Table 5 · p. 20
E.1. Emboddied Planning
The Minecraft examples contrast a zero-shot plan containing factual errors with a RAT plan whose error rate is described as significantly reduced. RAT retrieves information relevant to intermediate items, although the example still lacks a furnace prerequisite.
E.1. Emboddied Planning
- The planning task depends on world and causal knowledge because earlier steps commonly serve as preconditions for later actions.Examples include Minecraft recipes, action tools, and item-quantity restrictions.Source: §E.1 · p. 21
- The zero-shot plan contains errors in the crafting-table recipe, missing pickaxe materials, iron-ore mining tools, and the golden-apple recipe.The cited example identifies wood instead of planks, omits sticks, skips use of a stone pickaxe, and substitutes a water bucket for gold ingots and an apple.Source: §E.1 · p. 21
- RAT’s example still omits mining cobblestone to craft a furnace before smelting iron, but its plan’s error rate is significantly reduced relative to ChatGPT’s.The passage describes the missing furnace as a remaining error despite the overall reduction.Source: §E.1 · p. 22
- RAT retrieves links related to intermediate items at each step, whereas traditional RAG retrieves from the instruction and can find only the final step.The authors say this retrieval pattern improves accuracy and plausibility.Source: §E.1 · p. 22
E.2. Creative Writing
The creative-writing examples show different generated histories of the United States, from concise chronologies to longer accounts spanning colonial history, national founding, and later developments. The RAT example includes pre-colonial diversity, the Revolution, constitutional formation, and more recent history.
E.2. Creative Writing
- The five-shot RAG response also spans colonization through contemporary challenges, including expansion, the Civil War, world wars, and civil-rights movements.It characterizes the nation as diverse and influential, with a complex political landscape.Source: §E.2 · p. 24
- The original ChatGPT response presents a chronology from European colonization through independence, westward expansion, twentieth-century wars, and the present.It describes the official history as beginning in 1776, while also noting Indigenous presence before European settlement.Source: §E.2 · p. 25
- The RAT response covers pre-colonial Native American populations and colonial settlement by Spanish, French, Dutch, and British powers.It describes colonies with distinct economic and political frameworks.Source: §E.2 · p. 25
- The RAT response frames the Revolution and Declaration of Independence as a founding period, then discusses the Constitution and Bill of Rights.It also describes global reverberations and rights in connection with the Revolution and founding documents.Source: §E.2 · p. 25
- The RAT response continues through westward expansion, industrialization, the Civil War, immigration, urbanization, and overseas conflicts.It also describes the United States’ emergence as a global superpower after World War II.Source: §E.2 · p. 25
- The RAT response includes recent social movements, technological advances, and challenges involving healthcare and international relations.Source: §E.2 · p. 25