LLMs Get a Memory Boost with HippoRAG

Large Language Models (LLMs) have quickly proven themselves to be invaluable tools for thinking. Trained on massive datasets of text, code, and other media, they can generate human-quality writing, translate languages, generate images, answer your questions in an informative way, and even write different kinds of creative content.  But for all their brilliance, even the most advanced LLMs have a fundamental constraint: their knowledge is frozen in time.  Everything they "know" is determined by the data they were trained on, leaving them unable to adapt to new information or learn about your specific needs and preferences.

To address this limitation, researchers developed Retrieval-Augmented Generation (RAG). RAG allows LLMs access to access datastores that can be updated in real-time. This access to dynamic external knowledge bases, allows them to retrieve relevant information on the fly and incorporate it into their responses. Since they tend to rely on keyword matching, however, standard RAG implementations struggle when a question requires connecting information across multiple sources —  a challenge known as "multi-hop" reasoning. 

Inspired by how the brain stores and retrieves memories, researchers developed HippoRAG, a new approach to RAG that retrieves and incorporates more meaningful sources in generated responses. In this post, we'll delve into how HippoRAG works, explore its advantages over traditional RAG, and glimpse its potential to unlock new levels of reasoning and understanding in AI systems.

When RAG Falls Short: The Need for Deeper Connections

In a typical RAG system, you have two key components: a retriever and a generator. The retriever's job is to search a massive database of text (the knowledge base) – imagine Wikipedia, a company's internal documents, or even your personal files – and find documents relevant to a given question. This often involves converting the question and documents into numerical representations (embeddings) and using clever algorithms to quickly find the documents with the most similar embeddings to the question. The generator, typically a powerful LLM, then takes these retrieved documents as context and crafts a comprehensive, well-informed answer.

For example, if you ask a RAG system "What are the main tourist attractions in Paris?", the retriever would search its knowledge base for documents containing information about Paris and tourist attractions. It might find articles from Wikipedia, travel blogs, or even tourist guides. The LLM would then use these retrieved documents to generate a response, perhaps listing popular sites like the Eiffel Tower, the Louvre Museum, and the Arc de Triomphe.

While this is a powerful tool, traditional RAG systems often struggle when a question requires connecting the dots across multiple pieces of information – a challenge known as "multi-hop" reasoning. Think about asking your AI assistant: "Should I pack an umbrella for my trip to London next week?" To answer this, the assistant needs to pull information from your calendar to confirm the travel dates, check the weather forecast for London during those dates, and consider your personal packing preferences (do you always pack an umbrella, or only when absolutely necessary?). Standard RAG systems, which often rely on simple keyword matches, might find documents mentioning "London" and "umbrella," but wouldn't necessarily understand the temporal connection to your trip or your packing preferences.

Similarly, a question like "What is the capital of the country where the current CEO of Google was born?" requires linking information about the CEO's birthplace to the capital of that country – a connection that might not be explicit in any single document. Traditional RAG would struggle to make these connections effectively.

HippoRAG: Mimicking the Brain's Memory Index

Enter HippoRAG, a new RAG framework that takes inspiration from how our own brains store and retrieve memories. The human brain doesn't just store information in isolated chunks; it creates a rich web of associations between different concepts. This ability to link related ideas is what allows us to reason, make inferences, and answer complex questions that require piecing together information from multiple sources.

The hippocampal indexing theory, which inspired HippoRAG, provides a model for how this works in the brain:

• Neocortex: The "thinking" part of your brain processes sensory information and stores complex knowledge. This is analogous to the LLM in HippoRAG.
• Hippocampus: A region deep in your brain acts like an "index" for your memories. It doesn't store the complete memories themselves, but it creates links (associations) between different pieces of information stored across the neocortex. Think of it like a mental map of how concepts are related. This is the role of the knowledge graph in HippoRAG.

When you experience something new, your neocortex processes it, and the hippocampus creates links between the relevant concepts, forming a memory trace. Later, when you're reminded of a part of that experience, your hippocampus activates the associated links, triggering the retrieval of the complete memory from the neocortex.

Building a Better Memory for LLMs

HippoRAG mimics this brain-inspired model to give LLMs a more sophisticated memory system. Let’s take a look at how it works:

1. Building the Hippocampal Index

HippoRAG uses an LLM to extract key concepts and relationships from the knowledge base, building a knowledge graph where nodes represent concepts and edges represent relationships between them. This knowledge graph is like the hippocampus, storing the connections between ideas.

Here’s a simplified representation of what a knowledge graph might look like for our example:

• Nodes: Represent entities and concepts: "London", "England", "Weather", "Unpredictable", "Umbrella", "Protection", "Rain", "Trip", "Next Week", "Monday", "Friday", "Pack Light", "I", etc.
• Edges: Represent relationships between nodes: "London" -[is the capital of]-> "England", "London" -[is known for]-> "Unpredictable Weather", "Umbrella" -[provides]-> "Protection", "Protection" -[from]-> "Rain", "Trip" -[destination]-> "London", "Trip" -[time]-> "Next Week", "Next Week" -[includes]-> "Monday", "Next Week" -[includes]-> "Friday", "I" -[preference]-> "Pack Light", etc.

2. Query Time Reasoning With Personalized PageRank

Given a new question, the LLM identifies the key entities and maps them to nodes in the knowledge graph. Then, HippoRAG uses an algorithm called Personalized PageRank (PPR) to explore the knowledge graph, spreading activation across related nodes. This is like the hippocampus activating the relevant memory traces. PPR allows HippoRAG to efficiently gather information from multiple "hops" away from the original entities, capturing multi-hop relationships in a single step. In our Example:

• Entity recognition: As before, the LLM identifies key entities in the question: "umbrella", "trip", and "London".
• PPR on knowledge graph: Starting from the nodes representing these entities, PPR explores the knowledge graph, spreading activation across related nodes. It considers the strength and direction of the edges to determine the relevance of different paths.

In our example: PPR might highly activate paths leading to nodes like "Rain", "Unpredictable Weather", and "Protection" due to their connections to "London" and "Umbrella".

3. Single-Step Retrieval

The most highly activated nodes (and their associated text chunks from the knowledge base) are then retrieved. This provides the LLM with the necessary information to answer the question, including the crucial connections between concepts.

In our example, this will likely include chunks 1, 2, and 3 from our original example.

4. Answer Generation With LLM

The LLM now has all the pieces of the puzzle – the original question, the retrieved knowledge (enriched with graph-based connections), and any additional real-time information. It can leverage this richer knowledge to provide a more nuanced and accurate answer.

In our Example:

• Combined input: The LLM receives the original question, the retrieved knowledge (now enriched with graph-based connections), and the real-time weather forecast for London for the trip dates.
• Enhanced reasoning: The LLM can now leverage the richer knowledge to provide a more nuanced and accurate answer. It understands not only that London has unpredictable weather but also that umbrellas offer protection from rain and that the trip is scheduled during a time with a potential for rain.

From Multi-Hop to Path-Finding: The Future of AI Memory

The researchers behind HippoRAG demonstrated that it significantly outperforms standard RAG methods on multi-hop reasoning tasks. But the implications of this approach go far beyond simple question answering.

The concept of "path-finding" retrieval, enabled by HippoRAG, is particularly exciting. Imagine an AI system that can not only retrieve information but also discover new connections between concepts, even if those connections are not explicitly stated in the data. This would be a game-changer for fields like scientific discovery, legal reasoning, and personalized recommendations, where the ability to make novel connections is essential.

While HippoRAG faces challenges like scaling to massive knowledge graphs and managing the concept-context tradeoff, it represents a significant leap toward building LLMs with more human-like memory capabilities. As we continue to explore the intersection of neuroscience and artificial intelligence, we are moving closer to creating AI systems that can learn, remember, and reason with the same depth and flexibility as the human brain.