Boosting the Potential of Large Language Models with an Intelligent Information Assistant

Authors: yujia  zhou
zheng  liu
zhicheng  dou
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:48:51

Abstract: The emergence of Large Language Models (LLMs) has significantly advanced natural language processing, but these models often generate factually incorrect information, known as "hallucination." Initial retrieval-augmented generation (RAG) methods like the "Retrieve-Read" framework was inadequate for complex reasoning tasks. Subsequent prompt-based RAG strategies and Supervised Fine-Tuning (SFT) methods improved performance but required frequent retraining and risked altering foundational LLM capabilities. To cope with these challenges, we propose Assistant-based Retrieval-Augmented Generation (AssistRAG), integrating an intelligent information assistant within LLMs. This assistant manages memory and knowledge through tool usage, action execution, memory building, and plan specification. Using a two-phase training approach—Curriculum Assistant Learning and Reinforced Preference Optimization—AssistRAG enhances information retrieval and decision-making. Experiments show AssistRAG significantly outperforms benchmarks, especially benefiting less advanced LLMs, by providing superior reasoning capabilities and accurate responses.

Main paper: 
 introduction: 
 The emergence of Large Language Models (LLMs) has significantly advanced the field of natural language processing, demonstrating an impressive ability to mimic human-like language patterns [1] . However, despite their extensive knowledge acquired during training, LLMs can occasionally generate factually incorrect information, a phenomenon referred to as "hallucination" [2, 3] . To address this, the integration of retrieval systems with LLMs has been suggested, allowing these models to tap into external databases to generate more reliable responses [4] .Initially, retrieval-augmented generation (RAG) relied on a simple "Retrieve-Read" framework [5] , which was adequate for basic question-answering but insufficient for complex, multi-step reasoning tasks. As language models advanced, various prompt-based RAG strategies emerged [6, 7] , incorporating pre-retrieval and post-retrieval prompts to refine the process. However, these strategies heavily relied on the foundational capabilities of the language models. Consequently, the focus shifted to Supervised Fine-Tuning (SFT)-based RAG methods [8] , which involve fine-tuning language models specifically for RAG tasks to enhance their performance.While SFT-based methods have improved the quality of generated responses, they face two limitations that hinder their practical application. Firstly, these fine-tuned models are not easily adaptable to emerging LLMs, requiring retraining for each new foundational LLM. Secondly, directly fine-tuning a foundational LLM in the RAG scenario may change its innate abilities, potentially leading to negative impacts on the model's performance on other tasks. To address these challenges, we propose Assistant-based Retrieval-Augmented Generation (ASSISTRAG), which integrates an intelligent information assistant as a plugin within LLMs. This approach comprises a trainable assistant for information management and a static main LLM dedicated to task execution, as depicted in Figure 1 .As an intelligent information assistant, ASSISTRAG operates in two primary categories to handle complex tasks: memory management and knowledge management. Memory management involves integrating and analyzing content from internal memory, while knowledge management focuses on leveraging external knowledge. These two main functions are supported by four core capabilities of ASSISTRAG: (1) Tool usage, which involves recalling relevant information from both internal memory and external knowledge bases through a retriever; (2) Action execution, which involves processing, analyzing, and extracting information; (3) Memory building, which involves recording essential knowledge and reasoning patterns from historical interactions; (4) Plan specification, which involves determining the necessity of each step in the process. These four capabilities work together to ensure that ASSISTRAG can provide accurate and comprehensive support to the main LLM.To implement ASSISTRAG, we adopt a two-phase training approach. The first phase, Curriculum Assistant Learning, enhances the assistant's capabilities in note-taking, question decomposition, and knowledge extraction through progressively complex tasks. The second phase, Reinforced Preference Optimization, uses reinforcement learning to tailor the assistant's feedback to the main LLM's specific needs, optimizing knowledge extraction based on feedback from the main LLM.During the inference stage, ASSISTRAG operates through a three-step process: (1) Information Retrieval and Integration: The assistant understands the main LLM's needs, retrieves relevant knowledge from internal and external sources, and extracts valuable information. (2) Decision Making: The assistant evaluates and decides whether to provide the retrieved memories and knowledge to the main LLM based on their relevance. (3) Answer Generation and Memory Updating: The main LLM generates an answer using its internal knowledge and the assistant's information, while the assistant updates its memory with crucial reasoning steps.Results from experiments across three complex question-answering datasets reveal that ASSISTRAG exhibits superior reasoning capabilities and markedly outperforms existing benchmarks. Notably, when applied to different foundational LLMs, ASSISTRAG appears to confer more pronounced benefits on less advanced LLMs. 
retrieval-augmented generation: 
 RAG represents a significant advancement in the domain of LLMs, particularly for tasks demanding extensive knowledge. This paradigm begins with a retrieval step, where the LLM accesses an external database to gather relevant information before addressing queries. Traditionally, RAG follows a "Retrieve-Read" framework [9, 5, 10, 11] , with efforts focused on refining either the retriever or the generator through pre-training approaches to augment RAG's accuracy. Building on this foundation, new RAG strategies have emerged, including the use of prompt-based methods like Chain-of-Thought (CoT) reasoning [7, 12] , iterative retrieval processes [13, 14, 15, 16] , and leveraging LLM-generated content for dynamic retrieval [6, 17, 18] . These strategies underscore the LLMs' ability to select relevant information adaptively in response to specific contexts. Concurrently, research on finetuning LLMs for RAG applications is rapidly expanding [19, 20] , focusing on enhancing skills such as query reformulation [21] and knowledge integration [22, 23, 24] , as well as developing critical functions like determining the necessity of retrieval and appraising the value of retrieved data [8, 25] . Departing from these approaches, our paper introduces Assistant-based RAG, integrating an intelligent information assistant with the main LLM to boost its potential.Recent advancements in LLMs have facilitated the development of LLM-based autonomous agents such as AutoGPT [26] , Toolformer [27] , and MetaGPT [28] , which utilize LLMs for effective decision-making. Notably, ReAct [29] combines LLMs with external tools to manage knowledgeintensive tasks, allowing for dynamic responses to environmental changes. Additionally, models like WebGPT [30] integrate reinforcement learning with GPT-3, enabling the autonomous operation of search engines during text generation. Innovative methods used by Flare [17] and Self-Ask [6] determine optimal times for information retrieval, while Reflexion [31] endows LLMs with introspective mechanisms that continually refine their outputs. Our proposed Assistant-based RAG model further enhances LLM capabilities by combining memory management and knowledge management, thus providing robust support to the main LLM in tackling complex tasks.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: [TODO]Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methodology: 
 In this section, we first define the task of RAG and then introduce our proposed framework, AS-SISTRAG. ASSISTRAG enhances the capabilities of LLMs through the support of an intelligent information assistant. With abilities to use tools, execute actions, build memory, and plan, ASSIS-TRAG can provide precise memory and knowledge management services for LLMs.Given a question q and a collection of documents D = {d i } |D| i=1 , the main LLM aims to generate an answer y based on both the question and the relevant documents. This can be formalized as y = LLM main ([D q , q]), where D q represents the set of documents retrieved for the query q, and [•, •] denotes the concatenation of the retrieved documents with the query. Expanding this concept, ASSISTRAG employs an intelligent information assistant, LLM Assist , to enhance the main LLM's responses by providing relevant information, formalized as y = LLM main ([LLM Assist (q), q]). In the following sections, we will detail the capabilities of the ASSISTRAG framework, along with its training and inference procedures.By incorporating an intelligent information assistant, ASSISTRAG aims to boost the potential of LLMs in handling complex reasoning tasks. As illustrated in Figure 2 , this framework consists of two main components: a frozen main LLM tasked with generating answers based on the information provided, and a trainable assistant LLM responsible for information management. This assistant LLM is designed with two tasks: Memory Management involves storing interactions with the main LLM and retrieving relevant past memories to assist in addressing similar questions. Knowledge Management encompasses retrieving relevant information from external databases and processing it to support the main LLM in formulating responses to new questions.To effectively accomplish these tasks, we have endowed the assistant with four key capabilities:• Tool Usage: Retrieving relevant information from internal memory and external knowledge bases. • Plan Specification: Determining the necessity of assistance during answer generation.These four capabilities synergize to ensure that ASSISTRAG offers precise and comprehensive support to the main LLM. In the following sections, we will provide a detailed examination of the role and implementation of each capability.Effective memory management is crucial for enhancing the main LLM's performance by storing and retrieving historical interactions. This functionality comprises two key processes: capturing new insights and retrieving previously stored information. This stage activates the following three capabilities of AssistRAG:• Action I: Note Taking. This action F NT records critical information and the reasoning patterns behind each historical interaction. Given the historical interactions of the main LLM, which include question q, reference r, and answer y, the assistant is tasked with memorizing the key reasoning process behind the answer into the memory slot m q : m q ←-F NT (q, r, y). The accumulation of memory slots for all prior questions forms the assistant's memory M, which is utilized for subsequent memory retrieval.• Tool I: Memory Retriever. Given the question q and the assistant's memory M, the memory retriever retrieves historically relevant memories, represented as: M q ←-R memory (q, M).• Plan I: Assessing the Usefulness of Retrieved Memory. If the question is entirely new, the retrieved memories may not only be unhelpful but also negatively impact the main LLM's response. Therefore, we implement this plan to determine whether the retrieved memory slots should be provided to the main LLM. Using a prompt, the assistant evaluates whether the retrieved memories are beneficial for answering the current question. Only if the answer is affirmative will the retrieved memories be supplied to the main LLM. Effective knowledge management is essential for an intelligent information assistant, involving the efficient gathering of necessary knowledge to support the main LLM. This process includes analyzing the information needs of the main LLM, retrieving relevant knowledge, and integrating it. This process involves the following four capabilities of AssistRAG:• Action II: Question Decomposition. This action F QD aims to break down the current question into multiple sub-queries to facilitate the retrieval of knowledge across various aspects:Q ′ ←-F QD (q),where Q ′ represents a series of sub-queries derived from the question q.• Tool II: Knowledge Retriever. Utilizing a batch of sub-queries Q ′ , the knowledge retriever sources relevant documents from external knowledge bases D, denoted as:D Q ′ ←-R knowledge (Q ′ , D).• Action III: Knowledge Extraction. This action F KE involves extracting essential knowledge from a large number of retrieved documents. Given the question q and the retrieved documents D Q ′ , the assistant is responsible for extracting the relevant knowledge K q from the search results:K q ←-F KE (q, D Q ′ ).• Plan II: Evaluating the Relevance of Extracted Knowledge. To ensure the accuracy and relevance of the information provided to the main LLM, this plan determines whether the extracted knowledge should be included in the response generation process. Similarly, we prompt the assistant to assess whether the extracted knowledge is relevant to the current question.To summarize, we have endowed the assistant with memory capabilities and designed three actions, two retrieval tools, and two planning strategies integral to ASSISTRAG. Next, we will introduce the training strategies developed for ASSISTRAG, focusing on enhancing the accuracy of these actions and ensuring their compatibility with the main LLM.The training objectives of ASSISTRAG focus on two main goals: (1) enhancing the effectiveness of each action within the RAG process, and (2) ensuring that its outputs align with the main LLM's requirements. To achieve these two goals, as depicted in Figure 3 , we implement curriculum-based assistant learning and reinforced preference optimization to optimize the training of ASSISTRAG.Several studies have demonstrated that GPT-4 can achieve human-like annotation accuracy [32] .Based on this consideration, we leverage it to collect training data for the three actions. The supervised training samples for each specific action are cataloged as C QD , C KE , and C NT , preparing these for the assistant's subsequent training phase.Motivation. The tasks of question decomposition, knowledge extraction, and note-taking are interconnected, each contributing towards navigating the reasoning path from a question to its answer.To equip the assistant with a comprehensive understanding of the RAG process, we devise a step-wise curriculum assistant learning strategy designed to evolve from simpler to more complex tasks to foster a deepened mastery over time.Training Objective. The curriculum learning strategy integrates training samples across three sequential phases C QD → C KE → C NT . Each phase dedicates 60% of its focus to the task at hand, with the remaining 40% evenly divided between the other two tasks. The assistant's training employs the standard next token prediction target based on the training set D gen :EQUATIONwhere ϕ symbolizes the generator's adjustable parameters, and (x, y) is the pair of input and expected output. This methodical training strategy is designed to progressively refine the assistant's proficiency in each component of the RAG process, thereby boosting its effectiveness.Motivation. Although ASSISTRAG effectively handles RAG tasks after assistant learning, its output may sometimes not fully meet the downstream LLM's specific needs. To enhance integration, we implement reinforced preference optimization, a technique that adjusts the assistant's output based on feedback from the main LLM, ensuring tailored assistance that better meets its requirements.Training Objective. To optimize the assistant for better alignment with the main LLM, we adopt Direct Preference Optimization (DPO) [33] . This approach involves generating two sets of references, one from externally retrieved knowledge and the other generated by the assistant itself. The main LLM evaluates these sets, with a preference determined by comparing the F1 scores of its responses against correct answers. For reinforced preference optimization, we leverage the DPO algorithm's optimization objective, utilizing paired preference data D dpo :EQUATIONwherer θ (x, y i ) = β π θ (yi|x) πref(yi|x)is the reward implicitly defined by the language model π θ and the reference model π ref . This reinforced training stage enhances the assistant's capability to deliver assistance that aligns more closely with the main LLM's preferences, enhancing overall efficacy.Upon completing its training phase, ASSISTRAG initiates its inference process through three steps:Information Retrieval and Integration. At this initial stage, ASSISTRAG first activates Action II to understand the main LLM's information needs. It then uses Tool I and Tool II to retrieve relevant information from internal memory and external knowledge bases, respectively. Subsequently, it invokes Action III to extract essential knowledge from the retrieved documents.Decision Making. In this stage, ASSISTRAG decides whether to provide the retrieved memories and extracted knowledge to the main LLM. It activates Plan I and Plan II to evaluate the relevance and usefulness of the retrieved memories and knowledge for the current question. If the assistant deems them helpful, they are supplied to the main LLM to aid in answer generation.Answer Generation and Memory Updating. In the final phase, we prompt the main LLM to generate an answer based on the question, its internal knowledge, and the information provided by the assistant. Following this, ASSISTRAG activates Action I to utilize its note-taking feature, capturing crucial reasoning steps from the interaction and incorporating them into its memory. This ensures the assistant's knowledge base remains up-to-date.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
datasets and evaluation metrics: 
 In this study, we evaluate the effectiveness of our proposed method through experiments on three intricate question-answering datasets: HotpotQA [34] , 2WikiMultiHopQA [35] , and Bamboogle [6] . These datasets, all derived from Wikipedia documents, provide a uniform corpus and retrieval mechanisms to supply external references for LLMs. To manage costs, we follow a similar approach as previous studies [17] by selecting a maximum of 500 questions from each dataset's validation set for our experiments. To assess the performance, we employ Exact Match (EM), F1 score, and Precision (Prec.).We benchmark our model against three foundational models: LLaMA2-chat 7B [36] , ChatGLM3 6B [37] , and ChatGPT, assessing their performance in Closebook, Naive RAG, and ASSISTRAG settings. To compare our RAG framework with other RAG models, we include advanced prompt-based methods ReAct [29] , IRCoT [7] , Self-Ask [6] , an SFT-based RAG model Self-RAG [8] , and a knowledge extraction model LLMLingua [38] . We ensure a fair comparison by standardizing evaluation conditions across all models.Training Settings: In the assistant learning phase, we create a dataset comprising 50k training samples based on instruction-following input-output pairs across three distinct task types. The assistant LLM, which is based on ChatGLM3-6B [37] , is fully fine-tuned across all parameters. The training is conducted over 2 epochs with a batch size of 32 and a peak learning rate of 2e-5. For the preference optimization phase, we employ a DPO trainer and uses LoRA for fine-tuning, with a learning rate set to 1e-5 and the training duration extended to 2 epochs. The training code and data can be accessed at https://github.com/smallporridge/AssistRAG.Inference Settings: For employing ChatGPT, we opt for the gpt-35-turbo-16k model, accessed through its API at a temperature setting of 0. We use a Wikipedia dump as our document corpus, breaking down articles into 100-token passages. For both memory and knowledge retrieval, we deploy the off-the-shelf LLM Embedder [39] to fetch up to 5 documents per input.5 Results and Analysis 
main results: 
 The main results are presented in Table 1 . Several key findings can be observed as follows:Comparison among Different Reasoning Types. Applying our ASSISTRAG framework to Chat-GPT demonstrates a significant performance advantage over other models across all datasets. Specifi-cally, ASSISTRAG showcases its adaptability with different base LLMs, consistently outperforming them in both Closebook and Naive RAG settings. This result highlights the advantage of ASSISTRAG in effectively assisting a variety of downstream LLMs. Additionally, our method surpasses contemporary approaches employing prompt engineering or supervised fine-tuning, validating the efficacy of our curriculum assistant learning and reinforced preference optimization training strategies.Comparison among Different Base LLMs. By comparing the performance of ASSISTRAG across various base LLMs, it is observed that stronger base LLMs yield higher quality responses across all reasoning types. Notably, compared to Naive RAG settings, ASSISTRAG achieves performance improvements of 78%, 51%, and 40% for LLaMA, ChatGLM, and ChatGPT, respectively. This indicates that ASSISTRAG brings more substantial benefits to weaker base LLMs. A likely reason is that weaker models inherently have less robust noise resistance. Benefiting from the assistant's knowledge extraction capability, the main LLM only receives relevant knowledge to generate answers, leading to improved responses. 2 illustrates that removing or freezing any of the ASSISTRAG's actions results in decreased performance, underscoring the value of the assistant learning in the RAG context. Notably, maintaining these actions in a frozen state still outperforms completely removing them, highlighting their critical role in the RAG process. Concerning training strategies, the absence of planning, curriculum learning, and preference optimization slightly diminishes performance, indicating that a structured progression from simple to complex tasks and aligning with downstream LLM preferences contribute to the assistant providing more accurate information to the downstream LLM, thereby enhancing the accuracy of LLM responses. Token Usage of Different Methods. A notable benefit of our ASSISTRAG framework is its efficiency in preprocessing extensive information prior to engaging the main LLM. This not only enhances the inference speed of the main LLM but also minimizes token usage, which is particularly valuable when utilizing online API services like ChatGPT. We select a representative model for each reasoning type to compare their token consumption in terms of online API and SFT model, alongside their performance metrics F1 on the 2Wiki dataset. The results, as outlined in Table 3 , reveal significant differences in token usage among the methods. Prompt-based RAG methods tend to consume a large number of tokens due to their dependency on multiple API calls. On the other hand, SFT-based methods are more economical in terms of API calls but require retraining for new LLM adaptations. In contrast, our ASSISTRAG demonstrates a balanced approach by reducing API token costs while maintaining adaptability across different LLMs without the need for retraining. This method not only lowers the overall costs associated with API usage but also achieves superior performance. Training data size (k) Accuracy, Efficiency, and Cost Analysis. When evaluating an algorithm's value, we consider three dimensions: accuracy, efficiency, and cost. To compare different RAG methods, we calculate each method's F1 accuracy, inference speed, and cost. We then illustrate the relationships between these variables using three separate plots. From Figure 4 , we observe that ASSISTRAG stands out as the most balanced method, achieving the highest F1 accuracy of 45.6, while maintaining a comparable inference time of 5.73 seconds and a low cost of 0.009 cents per question. Although methods like IR-CoT show higher costs and longer inference times, they do not surpass ASSISTRAG in accuracy. These results demonstrate that ASSISTRAG is advantageous for applications requiring high accuracy without incurring significant costs.We examine the effect of training dataset size on model performance by creating subsets of 5k, 10k, and 20k instances from our original 50k training samples. These subsets fine-tune three separate model versions, evaluated on three datasets, and compared to the model trained on the full 50k dataset. We also compare the impact of curriculum learning and DPO training by evaluating performance with each strategy omitted. Figure 5 shows a clear performance improvement for ASSISTRAG as the training dataset size increases from 5k to 50k across both datasets, indicating potential further gains with larger datasets. The curriculum learning strategy performs better than random mixed training, especially with smaller datasets, showing its advantage when data is limited. In contrast, DPO training benefits more from larger datasets, likely because more data enables better training for high-quality data generation.Case Study. Table 4 is a case study that highlights the capabilities of AssistRAG in processing and answering complex comparative questions. In this case study, the main question "Who is older, Danny Green or James Worthy?" is systematically broken down by ASSISTRAG into simpler sub-questions regarding the birth dates of both individuals. This decomposition enables targeted information retrieval, allowing the system to accurately locate and extract relevant birth date information from the corpus. ASSISTRAG effectively retrieves multiple pieces of information, including relevant and irrelevant entries, and filters through them to extract the necessary facts. For instance, it identifies the birth dates of both Danny Green and James Worthy and ignores unrelated entries, such as those concerning another individual named Danny Green who is a boxer. The memory retrieval capability is then utilized to access previous similar questions and their answers, which aids in reinforcing Final Output: Based on the analysis, Main LLM outputs:-James Worthy was born on February 27, 1961, while Danny Green was born on June 27, 1987. Since James Worthy was born earlier, he is older. So the answer is James Worthy.Ground-truth: James Worthy the current decision-making process. This step provides context and supports consistency in the reasoning pattern applied by the system. Combining the extracted knowledge and retrieved memory, ASSISTRAG plans the final response by confirming the usefulness of both sources of information. The Main LLM then generates a comprehensive answer, stating that James Worthy, born on February 27, 1961 , is older than Danny Green, born on June 22, 1987 . This case study showcases ASSISTRAG's ability to manage complex tasks by leveraging its multi-step reasoning process, ensuring the delivery of accurate and reliable answers. The superior performance in accurately answering comparative questions is a testament to the system's robust architecture and its effective integration of question decomposition, information retrieval, knowledge extraction, planning, and memory capabilities. 
conclusion: 
 In this study, we introduce ASSISTRAG to augment LLMs with an intelligent information assistant, significantly improving their ability to tackle tasks requiring complex reasoning. By implementing a two-stage training methodology that integrates curriculum assistant learning with reinforced preference optimization, we enhance the assistant's skills in memory and knowledge management. Experiments demonstrate that ASSISTRAG surpasses existing baselines with a notable margin.Looking ahead, we plan to further expand the assistant's skills to include long-text processing [40] and personalized support [41] , thereby providing more effective assistance to the main LLM. 2. Preference Optimization Phase:• Optimization Technique: We employed a DPO (Distributed Preference Optimization) trainer to refine the assistant's feedback mechanisms. • Fine-tuning with LoRA: Low-Rank Adaptation (LoRA) was utilized for fine-tuning, which helps in adjusting a subset of model parameters efficiently, reducing the computational load. • Learning Rate and Duration: The learning rate was set to 1e-5, with the training duration extending to 2 epochs.Training Resources: The entire training process was conducted using 8 A800 GPUs, providing substantial computational power to handle the intensive training tasks efficiently. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer:[Yes] Justification: [TODO] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: [TODO] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? 
 
 Appendix: During the inference phase, the ASSISTRAG system was fine-tuned to operate seamlessly with the main LLM, ensuring effective retrieval and processing of information.• Model Selection: For the main LLM, we opted for the gpt-35-turbo-16k model, accessed through its API. This model was chosen for its extended context window and advanced capabilities in handling complex queries. • Temperature Setting: The temperature was set to 0, ensuring deterministic outputs which are crucial for consistency in inference tasks. • Document Corpus: A Wikipedia dump was used as the primary document corpus. Articles were segmented into 100-token passages to facilitate efficient retrieval. • Retrieval tool: Both memory and knowledge retrieval processes utilized the off-the-shelf LLM Embedder [39] , capable of fetching up to 5 documents per input, ensuring comprehensive coverage of relevant information.Inference Resources: The inference tasks were efficiently handled using a single A800 GPU, highlighting the system's capability to deliver high performance with minimal resource usage.The training data annotation process was critical in ensuring the quality and relevance of the dataset. We utilized GPT4-turbo for annotating the dataset, ensuring high accuracy and consistency across the training samples. The dataset included approximately 50k training samples, labeled to cover a wide range of instruction-following tasks.The following tables ( Secondly, the effectiveness of ASSISTRAG depends on the quality and comprehensiveness of the external knowledge bases and internal memory it accesses. In scenarios where the available data is sparse, outdated, or biased, the assistant's ability to retrieve and integrate relevant information may be compromised, leading to suboptimal or erroneous outputs from the main LLM. This dependence on data quality underscores the need for continuous updates and maintenance of the knowledge sources.Lastly, the decision-making process during inference, which involves the assistant evaluating the relevance of retrieved information, is inherently difficult. The assistant's ability to accurately determine the necessity and applicability of specific knowledge is crucial for effective support. However, this process is susceptible to errors, particularly in scenarios involving ambiguous or multifaceted queries.Enhancing the precision and reliability of this decision-making mechanism is a key area for further research.NeurIPS Paper Checklist Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The reviewers agree that this paper addresses an important problem and makes a significant contribution to the existing literature on RAG-based systems. No major weaknesses were identified.

If ultimately accepted, I expect the authors to integrate the reviewer feedback into the final version.

Reviews: review: paper_summary: paper_summary: This paper proposes AssistRAG, an architecture for augmenting LLMs with a separate, trainable agent that helps with information retrieval and memory/knowledge management. The authors motivate such an architecture (as opposed to, say, fine-tuning the main LLM for RAG) and describe how to build and train it. The conduct experiments with 3 information heavy benchmarks and show consistent performance gains compared to a number of recent baselines.
strength_weakness: strength_weakness: I don't see any notable weaknesses in this paper. It's possible that I am not aware of the latest in this line of work and my assessment of novelty may not be accurate. I will defer to the other reviewers on the novelty aspect.

Some of the datasets used in this paper are rather old -- they are from the era of the train-test paradigm with 10s of thousands of training instances (e.g., HotpotQA was published in 2018!). I realize they are being used here in the model context of prompted LLMs. Nevertheless, I wonder if there are newer datasets that might be more timely.
questions: questions: Please see the weakness section.
limitations: limitations: Yes
paper_summary: paper_summary: To address the limitation of LLMs generating factually incorrect information, the authors have introduced AssistRAG, an intelligent information assistant with LLMs, building upon existing retrieval-augmented generation (RAG) strategies. The system operates in two main categories: memory management and knowledge management. AssistRAG employs a two-phase training approach: Curriculum Assistant Learning and Reinforced Preference Optimization. During inference, AssistRAG follows a three-step process: information retrieval and integration, decision-making, and answer generation with memory updating.
strength_weakness: strength_weakness: 1. The paper might benefit from addressing potential scalability issues and providing more extensive comparisons with a broader range of state-of-the-art models. 
2. Including real-world application scenarios and discussing potential limitations or challenges in deploying AssistRAG in practical settings would also enhance the paper's impact and applicability.
questions: questions: 1. How does AssistRAG perform in terms of computational efficiency and scalability when applied to very large datasets or in real-time applications? Can the authors provide any benchmarks or comparisons?
2. The ablation studies are informative, but could the authors include additional analysis on the sensitivity of AssistRAG to different hyperparameters or training configurations?
3. Incorporate user studies or qualitative evaluations to assess the practical usability and effectiveness of the assistant in aiding human users.
limitations: limitations: 1. Llama 2 Chat (7B parameters) and ChatGLM3 (6B parameters) have substantially fewer parameters than GPT-3.5, which is reported to have a much higher parameter count. Moreover, GPT4-turbo is used for annotating the dataset (as mentioned in Appendix section A.1.3), however, the model isn’t used, or mentioned, in the experiments. This disparity in model size should be taken into account when comparing the performance of these models, particularly in relation to GPT-3.5.
2. The analysis done on token usage (section 5.2) is only conducted on the 2WikiMultiHopQA dataset, however other experiments in the study are done on the HotpotQA, 2WikiMultiHopQA, and Bamboogle datasets. The paper could benefit from a consistent dataset used for all evaluations.
paper_summary: paper_summary: The paper introduces AssistRAG, a framework that integrates an intelligent information assistant within LLMs. AssistRAG employs a two-phase training approach involving Curriculum Assistant Learning and Reinforced Preference Optimization, focusing on memory management and knowledge management. Experiments on three QA datasets demonstrate the effectiveness.
strength_weakness: strength_weakness: 1. Questionable innovation, the key ideas of the proposed method (memory and knowledge management) become a standard practice in terms of Agent design.
2. The proposed method involves model training and serving, which could be costly, and also make it harder to adapt to different foundation LLMs.
3. A.2 shows that 58% of the errors stem from insufficient knowledge retrieval, but no further details are provided. Given such a scenario, the retrieval strategy may need to be improved. 
4. Using GPT-4 to generate fine-tuned data potentially limits the proposed method.
5. Notations are unclear, e.g., line 97, what is d_i? Eq. 2 is not well explained.
6. Table 1 shows that AssistRAG is combined with LLaMA2, ChatGLM, and ChatGPT, but the manuscript only mentioned using ChatGLM3 for training and ChatGPT in inference, which is inconsistent.
questions: questions: 1. Please include more details about the (ChatGPT) model versions employed in this paper.
2. The articles were segmented into 100-token passages - this is a relatively short length for LLM, and A.2 also shows that insufficient knowledge retrieval caused more than half of the errors, why don't consider increasing the chunk size?
3. Experiment data and questions are focused on Wikipedia, how about other domains?
limitations: limitations: 1. Given the current development of the Agent framework, it is necessary to compare the proposed method with the Agent solution.

Comments: title: Re: Rebuttal by Authors
Response to authors
comment: Thank you for your positive feedback and valuable suggestions. We will explore integrating the Musique dataset and expand our comparison part to provide a clearer analysis of similar work in the updated version. We appreciate your support and look forward to enhancing our paper based on your recommendations.
Thank you for the response. Given your focus on Multihop QA datasets, `Musique` might have been a good one to try as it has some harder questions requiring 3 or 4 hops. Nevertheless, after looking at the other reviews and your responses, I remain in favor of accepting this paper.  The only other suggestion I have is to include a clearer comparison with other similar work in the updated version of the paper.  Thank you for the work!
Thank you for your thorough evaluation and for considering the content of our rebuttal. We sincerely appreciate your willingness to reevaluate our work and the constructive feedback you’ve provided.
***
Regarding your concern about insufficient knowledge retrieval, this is indeed an interesting issue. We conducted experiments by increasing the chunk size to 512 tokens while keeping the number of retrieved documents constant. The results show the following changes in error proportions:

| Error Type | Original Proportion | New Proportion |
|------------|---------------------|----------------|
| Insufficient Knowledge Retrieval | 58% | 48% |
| Knowledge Extraction Errors      | 12% | 20% |
| Answer Reasoning Mistakes        | 20% | 22% |
| Other                             | 10% | 10% |

Increasing the chunk size did reduce the proportion of insufficient knowledge retrieval errors but also led to a higher likelihood of knowledge extraction errors due to the introduction of more irrelevant context. This trade-off is worth further exploration. Our opinion is that for models with strong knowledge extraction capabilities, increasing the chunk size or the number of retrieved documents can be an effective strategy.
***
Once again, thank you for your invaluable insights and your support of our work. Your recognition is our greatest encouragement.
Thank you for your response and the additional information provided. After carefully considering the paper and the rebuttal content, I have decided to slightly increase my evaluation, based on a comprehensive assessment of the work's contribution to the field.

I would like to further clarify my concern regarding the insufficient knowledge retrieval. The current experimental results indicate that insufficient knowledge retrieval accounts for 58% of the errors, suggesting potential issues at the retrieval stage. This problem typically does not stem from the retrieval algorithm, package, or vector database, but rather from the data processing workflow. This raises questions about the use of 100-token passages, as current best practices typically involve chunk sizes in the range of 500-1000 tokens to capture more semantic content.

Thank you for your efforts, and good luck.