AI Innovations and Insights 16: ASSISTRAG

The Dual Engine of Thinking and Memory

This video contains a mind map:

This article is the 16th in this promising series. Today, we will explore an advancement of RAG.

ASSISTRAG: The Dual Engine of Thinking and Memory

Training code and data: https://github.com/smallporridge/AssistRAG.

Vivid Description

ASSISTRAG works like a helpful assistant that helps organize your work materials and keeps track of what's important. When you have a difficult problem to solve, ASSISTRAG helps by breaking it into smaller, easier parts and finding the information you need.

Now, let's dive into the detailed introduction.

Background

Early RAG methods like the "Retrieve-Read" framework were inadequate for complex reasoning tasks. While subsequent prompt-based RAG strategies and Supervised Fine-Tuning (SFT) methods improved performance, they required frequent retraining and risked altering foundational LLM capabilities.

Figure 1: Comparisons of Naive, Prompt-based, SFT-based and our Assistant-based RAG frameworks. Source: ASSISTRAG.

By introducing an intelligent information assistant that integrates memory and knowledge management, ASSISTRAG compensates for LLMs' shortcomings in information accuracy and reasoning depth. As shown in Figure 1, this approach consists of a trainable assistant that manages information and a static main LLM that handles task execution.

Overview

Figure 2: Overview of ASSISTRAG. ASSISTRAG enhances LLMs by providing an intelligent information assistant. Endowed with the ability of tool usage, action execution, memory building and plan specification, it can achieve effective memory and knowledge management. Source: ASSISTRAG.

As shown in Figure 2, ASSISTRAG handles two key tasks: Memory Management for storing and retrieving past interactions, and Knowledge Management for accessing external information to support responses. It has four capabilities:

• Tool Usage: Retrieving information from memory and knowledge bases

• Action Execution: Reasoning and knowledge analysis

• Memory Building: Recording key knowledge from past interactions

• Plan Specification: Determining when assistance is needed

See Figure 3 for details.

Figure 3: The details of ASSISTRAG’s tasks. Image by author.

Training

Figure 4: Training framework of ASSISTRAG. It undergoes a two-stage training pipeline through curriculum assistant learning and reinforced preference optimization. Source: ASSISTRAG.

ASSISTRAG uses a two-phase training strategy to enhance retrieval effectiveness and improve generation quality without modifying the core LLM.

• Curriculum Assistant Learning progressively trains the system through a structured curriculum that evolves from simpler to more complex tasks, emphasizing deep mastery over time.

• Reinforced Preference Optimization adjusts ASSISTRAG’s outputs based on feedback from the main LLM using Direct Preference Optimization (DPO), focusing on alignment with the LLM’s specific needs.

Implementation and Case Study

For implementation, the details is shown in Figure 5.

Figure 5: The implementation details of ASSISTRAG. Image by author.

Figure 6 demonstrates AssistRAG's handling of comparative questions through a case study asking "Who is older, Danny Green or James Worthy?"

Figure 6: Case Study of AssistRAG. Source: ASSISTRAG.

The system breaks this into sub-questions about birth dates, retrieves and filters relevant information, and accesses similar past questions and their answers from memory. By combining extracted knowledge with memory retrieval, AssistRAG determines that James Worthy (born February 27, 1961) is older than Danny Green (born June 22, 1987).

Commentary

ASSISTRAG takes a new approach by keeping the main LLM separate from its helper system. This separation, combined with better ways to find information and solve problems, makes it an important new development.

In my opinion, the assistant introduces additional computation during training and inference, which may pose a challenge. Investigating lightweight implementations or hybrid models that reduce latency without significant accuracy trade-offs would enhance its appeal.

In addition, the decision-making process of the current system is not transparent to users, which may lead to trust issues in some high-risk domains (such as healthcare, finance, etc.). To address this, we can provide more transparent processes for memory retrieval and knowledge integration. For example, display retrieved information and its sources through an interactive interface, enabling users to intervene or adjust decisions.

Moreover, how can we make sure the memory system only keeps the most useful and up-to-date information? If old or wrong information stays in the system's memory, it can lead to poor decisions and faulty reasoning. Potential improvements include implementing memory decay strategies, leveraging user feedback to determine which memories to retain, and creating hierarchical memory structures that distinguish between short-term, long-term, or critical memories.

---

Finally, if you’re interested in the series, feel free to check out my other articles.