GEM-RAG: Graphical Eigen Memories For Retrieval Augmented Generation

TL;DR

Inspired by human memory, Graphical Eigen Memories for Retrieval Augmented Generation (GEM-RAG) is introduced, which tags information with LLMgenerated “utility” questions, links information together in a graph by similarity, and uses eigendecomposition to form higherlevel summary information.

Abstract

The ability to form, retrieve, and reason about memories in response to stimuli serves as the cornerstone for general intelligence - shaping entities capable of learning, adaptation, and intuitive insight. Large Language Models (LLMs) have proven their ability, given the proper memories or context, to reason and respond meaningfully to stimuli. However, they are still unable to optimally encode, store, and retrieve memories - the ability to do this would unlock their full ability to operate as AI agents, and to specialize to niche domains. To remedy this, one promising area of research is Retrieval Augmented Generation (RAG), which aims to augment LLMs by providing them with rich in-context examples and information. In question-answering (QA) applications, RAG methods embed the text of interest in chunks, and retrieve the most relevant chunks for a prompt using text embeddings. Motivated by human memory encoding and retrieval, we aim to improve over standard RAG methods by generating and encoding higher-level information and tagging the chunks by their utility to answer questions. We introduce Graphical Eigen Memories For Retrieval Augmented Generation (GEM-RAG). GEM-RAG works by tagging each chunk of text in a given text corpus with LLM generated ``utility'' questions, connecting chunks in a graph based on the similarity of both their text and utility questions, and then using the eigendecomposition of the memory graph to build higher level summary nodes that capture the main themes of the text. We evaluate GEM-RAG, using both UnifiedQA and GPT-3.5 Turbo as the LLMs, with SBERT, and OpenAI's text encoders on two standard QA tasks, showing that GEM-RAG outperforms other state-of-the-art RAG methods on these tasks. We also discuss the implications of having a robust RAG system and future directions.

Figures (2)

• Figure 1: An overview of the graph construction and retrieval process for GEM-RAG. Given a corpus of text, the text is first grouped into chunks of text. GEM-RAG then generates utility questions for each chunk of text using an LLM, where a utility question asks something that could be answered given the text chunk as context. Next GEM-RAG uses these utility questions and respective embeddings to build a weighted graph. Summary nodes are then generated using the graph's spectral decomposition, using the eigenvectors to represent different orthogonal modes or "eigenthemes" of the text. For retrieval, GEM-RAG embeds the question or prompt, and searches the graph for the optimal nodes or context to return.
• Figure 2: Ablation study on the effect of number of utility questions and eigencomponents for a 50 passage subset of the QuALITY data set. The accuracy from more utility questions quickly becomes saturated, whereas the best number of eigencomponents is two. Intuitively the number of utility questions capture the amount of information per chunk, whereas the number of components covers the importance of summaries in the domain of questions.

Theorems & Definitions (3)

• Remark 1
• Remark 2
• Remark 3