Optimizing Inference-Time Reasoning in LLMs via Retrieval-Augmented Reflection

Zihao Wang, Haowei Lin, Ruilin Yan, Xiangyu Wang, Jiaqi Li, Weiye Shi, weipeng chen, Xiaojian Ma, Anji Liu, Yitao Liang

Keywords: Retrieval-augmented Generation, Reasoning, Large Language Models

Abstract: Empowering LLMs to improve their performance through increased inference-time computation is a crucial step in developing self-improving agents capable of operating in open-ended natural language contexts. In this paper, we explore how iteratively revising a chain of thoughts guided by information retrieval significantly improves large language models' reasoning ability in challenging tasks, while hugely mitigating hallucination. In particular, the proposed method --- \emph{retrieval-augmented reflection} (RaR) --- revises the generation tokens step by step, leveraging multiple pieces of retrieved information relevant to the intermediate reasoning steps and the instruction. Applying RaR during inference-time to a various set of language models substantially improves their performances on various reasoning tasks; on relatively increasing scores by up to +16.4\% on code generation, +11.6\% on mathematical reasoning, and 29.1\% on embodied task planning. Moreover, we find that with more inference-time computation given to the LLM for multi-times retrieval-augmented reflection, the LLM can continuously improve on various reasoning benchmarks. A small LM can surpass the performance of the LM with more than 10 times the parameters, when giving more computation cost.

Primary Area: foundation or frontier models, including LLMs

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 10331