LightRAG: Simple and Fast Retrieval-Augmented Generation for Legal Doc Analysis

LightRAG是一个创新的方法，基于 GraphRAG，它将知识图谱的属性与基于嵌入的检索系统相结合，既快速又高性能，取得了最先进的研究成果。由郭子瑞等人开发。来自香港大学（HKU）的 LightRAG，通过采用双层检索系统，缓解了社区遍历中与计算开销和缓慢、耗时相关的挑战。因此，LightRAG 在各种基准测试中都优于GraphRAG和传统的 RAG 系统。

我们知道，特别是法律、金融和医疗保健等行业，可以从最先进的 RAG（检索增强生成）方法中受益，因为这些行业拥有大量、复杂且多领域的数据，需要频繁更新以确保获得相关见解。

如果你是投资银行或专业服务公司，你在一个"零错误"的环境中运营。在涉及数百万美元的情况下，你不能容忍LLM出现幻觉并给出不准确的答案。

– RUI BAI，Invisible Technologies 的产品经理：来源

像 GraphRAG 和 LightRAG 这样的高级 RAG 技术是这些领域的改变者，相比朴素 RAG 提供了更有意义的改进和卓越的精确度。

根据论文，LightRAG 在法律数据集上表现尤为出色，这是其优势之一。为了验证这一点，我们将使用"2013 年公司法"和"印度宪法" PDF 作为法律文档分析的样本。

文章讨论的主题包括：

• 
  传统 RAG 的不足
• 
  基于 GraphRAG 方法的必要性
•  GraphRAG 中的痛点
• 
  LightRAG 如何解决这些问题？
• 
  LightRAG 演示的代码演练与 Ollama
• 
  使用 LightRAG 和 Neo4j 进行法律文档分析

1. 
   传统 RAG 系统的局限性
2. 
   微软 GraphRAG 的温和介绍
3. 
   LightRAG 架构：一个简单、快速且高效的图谱 x RAG
   1. 
      基于图的文本索引
   2. 
      双层检索范式
4. 
   LightRAG 框架的计算成本
5. 
   论文中的基准测试讨论
6. 
   使用 LightRAG 进行法律文档分析
7.  代码演练
   1.  LightRAG 索引
   2. LightRAG Querying
8. Response Comparison 1: Targeted at Local Entity
9. Response Comparison 2: Targeted at Global Relations
10. Takeways
11. Conclusion
12. References

This is the fifth article in our series of blogs on LLMs and RAG.

1. Deciphering LLMs: From Transformers to Quantization
2. Fine Tuning LLMs with PEFT
3. RAG with LLMs
4. Multimodal RAG with ColPali and Gemini
5. LightRAG:  Legal Doc Analysis

Limitations of Traditional RAG Systems

Problem Statement:

Existing RAG systems have significant limitations, including reliance on flat data    representations and inadequate contextual awareness, which can lead to fragmented answers that fail to capture complex inter-dependencies.

– LightRAG 2024 Paper

Baseline RAGs are found to be inefficient because,

• They fail to capture interconnectedness between fragmented information, spread across multiple documents making it challenging to outline comprehensive insights.
• Limited contextual understanding arises from a lack of a holistic overview of the retrieved chunks.
• Scalability inefficiencies occur when the corpus of data grows, resulting in poor retrieval quality.

To address these, Microsoft introduced a novel approach called GraphRAG which incorporates graph structures into the text indexing and retrieval process.

Note: We will explore nuances in GraphRAG and spend some time here, as it shares attributes and overlapping concepts with LightRAG.

Gentle Intro to GraphRAG by Microsoft

Knowledge Graphs are data structures consisting of a set of nodes that preserve relationships between different entities present across various data points. Structured knowledge graphs enable GraphRAG (Edge et al.) to excel in multi-hop reasoning by connecting dots or contrasting pieces of information.

For example, from contemporary scenarios a stock analyst might want to ask LLMs: “How will the U.S. election results impact Indian stock market trends?” or How did ChatGPT’s success fuel NVIDIA’s stock surge?

Questions like these are inherently a Query-Focused Summarization (QFS) where query context matters rather than simple retrieval. Baseline RAG often struggles to provide a comprehensive answer, as it lacks full context and fails to aggregate relationships between different nodes across the entire dataset. However GraphRAG shines at QFS to effectively draw conclusions with global understanding (like humans) even when the information needed is not explicitly stated but implied in the basket of sources.

Similar to typical RAG, GraphRAG pipeline also involves Indexing and Querying.

However the indexing process involves 2 stages as follows,

Indexing – GraphRAG

STAGE 1: Knowledge Graph (KG) Creation
To construct a strIncremental Indexinguctured KG, GraphRAG uses LLM’s to extract entities and relationships from the source documents. This process relies on a meticulously crafted prompt and multi-part gleaning checks.

Step 1: Entities are represented as individual nodes in the KG having information about name, organization or category and a brief description about the entity.

Step 2: Relationships are represented as edges between different nodes with a source and target entity ID. This defines the relationship between entities and rates them with connection weights aka strength scores ranging between 1 to 10.

STAGE 2 : Semantic Clustering

In the next stage, based on the density of node connections and for scalability, the Leiden algorithm is applied to discover modular communities by grouping closely related nodes into hierarchical clusters

 Community detection is used to partition the graph index into groups of elements (nodes, edges, covariates) that the LLM can summarize in parallel at both indexing time and query time.

– GraphRAG, Microsoft 2024

This helps to reduce the search space efficiently navigating the KG by focusing on highly correlated communities rather than the entire graph.

Using LLMs, these communities are summarized with a bottom-up approach  which serves as descriptors that have complete coverage of graph index.

Querying – GraphRAG

During the querying phase, when a user asks a question, entities and relationships within the query are identified for QFS. By comparing these elements between questions and the graph index, the most relevant communities are identified. 

Then these community summaries are randomly shuffled and intermediate responses at different community levels (local or global level) are generated by LLMs along with a helpfulness score ranging between 0 to 100 . This score indicates the relevance of the generated answer to the user query. The final global answer is generated using a multi stage map reduce approach, aggregating intermediate partial responses ranked descendingly based on helpfulness score as the LLMs context. 

Graph RAG supports both local and global query workflows. We will learn more about this in the LightRAG methodology which follows a similar approach with a  small difference.

Advantages of GraphRAG

• GraphRAG handles global or high-level queries exceptionally well and consistently performs better than naive or baseline RAG.
• It doesn’t suffer from scalability issues when the dataset is large like naive RAG.
• It offers enhanced contextual understanding and greater interpretability by grounding information within the graph index enabling entity or relation source citations.

Disadvantages of GraphRAG

• GraphRAG is often painfully slow to run as it requires multiple LLM API calls potentially hitting rate limits.
• It is extremely costly. The internet community based on testing suggests with GPT4o, it could cost around $6- $7 to index a simple book like ‘A Christmas Carol by Dickens’ having 32k words.
• To incorporate new data into existing graph indexes, we will need to reconstruct the entire KG for previous data as well, which is an inefficient approach.
• There are no explicit deduplication steps performed on repeated elements which results in a noisy graph index.

Even though GraphRAG seems promising, it’s not an efficient solution due to its operational cost and computation complexity. So what’s next?

KG-RAG is still in its early stages, and there is no consensus on how to effectively retrieve entities and relationships from KGs or how to integrate vector similarity search with graph structures. 

 –   Milvus DB : [Source]

LightRAG appears to be a compelling solution for the above statement.

How does LightRAG differ from NaiveRAG and GraphRAG?

Compared to Naive RAG, LightRAG offers:

• Comprehensive Information Retrieval with diverse answers.
• Efficient and Low-Cost Retrieval
• Fast Adaptation to Data and updates them with minimal reindexing.

LightRAG address two main pain points of GraphRAG by

• Reducing indexing and response time through a better approach than community traversal.
• Easily adapting to new data with an incremental update algorithm that updates only specific element instances.

LightRAG Architecture:  A Simple, Fast, and Efficient Graph x RAG

LightRAG enhances the retrieval process by segmenting documents into smaller, more manageable chunks Di. This chunking strategy enables quick identification of relevant content without the need to go through the entire document.

Graph Based Text Indexing

To identify the  key information needed, an LLM is used to identify entities (e.g., names, dates, locations, and events) and their relations to build an overarching knowledge graph.

This phase can be mathematically formulated as,

Where, 

represents the resulting knowledge graph.

Let’s take the following example for better understanding and consider it uses a sentence level chunking mechanism.

Sample text 1: 

Chunk 1:

Elon Musk envisions SpaceX‘s Mechazilla tower enhancing booster reusability to support sustainable space exploration. 

Chunk 2:

Operated by SpaceX, Mechazilla is designed to catch the Super Heavy booster, reducing disposable components, lowering costs, and minimizing environmental impact. 

Sample text 2:

Chunk 1:

“Musk’s commitment to sustainability is also reflected in Tesla, where he drives green energy innovations through electric vehicles and solar solutions. Tesla’s focus on renewable energy technologies complements Musk’s sustainability goals for SpaceX.”

To synthesize KG for the above chunks from two samples, three preprocessing steps are performed, 

Step 1: Extracting Entities and Relationships, R(.): As we saw in GraphRAG, similarly here key word extraction prompt is sent to the LLM to obtain the nodes and edges. From the above sample sentence, “Elon Musk envisions SpaceX’s Mechazilla tower enhancing booster reusability”

LLM analyzes the text and extracts the following,

Entities (Nodes): Elon Musk, SpaceX, Mechazilla tower, Booster reusability, Sustainable space exploration

Relationships (Edges):
Elon Musk → envisions → Mechazilla tower
Mechazilla tower → enhances → Booster reusability
Booster reusability → supports → Sustainable space exploration
SpaceX → operates → Mechazilla tower

Step2: LLM Profiling for Key Value Pair Generation P(.): Once the elements are extracted, an LLM empowered profiling function P(.) is employed to create structured entity key-value pairs and relationship KV pair for each entity node in 𝓥 and relation in Ɛ respectively. 

A Key (K) is a word or short phrase (e.g., Elon Musk) while Value (V) is a paragraph summarizing relevant chunks. Typically entities will have their name as sole index whereas relations may have multiple index keys to capture broader themes from connected entities. Using these KV pairs an initial knowledge graph is constructed.

This profiling enables efficient retrieval with relevant summaries that support text generation. 

An entity KV pair will look like,

• Key (K): “Elon Musk”

• Value (V): “CEO of SpaceX and Tesla, Elon Musk envisions using the Mechazilla tower to enhance rocket reusability, aiming to reduce space travel costs and advance sustainable space exploration.”

As discussed in the introduction, LightRAG combines graph indexing and standard embedding based approach. To be specific, the value(V) of entities and relations in key-value pairs after LLM profiling are passed to an intermediate step to generate vector embeddings using an embedding model. These KV data structures offer a more precise retrieval than less accurate standard embedding only RAG or inefficient chunk traversal techniques in GraphRAG.

Step 3: Deduplication to Optimize Graph Operation D(.) : For efficient data processing, repetitive entities and relations are merged thereby reducing both overhead and overall graph size.

For e.g., repeated instances of entities like SpaceX and Mechazilla will be merged into the same node, allowing for different edges to represent distinct relationships.

After deduplication we get a final optimized version of the initial KG.

Dual – Level Retrieval Paradigm

Aided by Graph Based Text Indexing, LightRAG pipeline employs a dual level retrieval strategy. This approach identifies low level and high level keys from multi hop subgraphs within the KG to answer a diverse set of queries.

Low level Retrieval:  Targets specific, granular information from individual nodes and edges to handle local queries, such as “What is Mechazilla?” At this level it provides detailed node-level insights.

High Level Retrieval:  Aggregates information from multiple entities from different documents to answer global queries that require broader themes or abstractions such as “How does Elon Musk’s vision promote sustainability across his ventures?” 

How does LightRAG utilize knowledge graphs?

For a given query, LightRAG’s retrieval algorithm extracts both local k^(l) and global query keywords k^(g). Then using vector similarity relevant entities are matched to local query with low level keys and global query keywords with high level concepts.

By gathering one-hop neighboring nodes in local subgraphs, LightRAG integrates an additional layer of context, improving the relevance of results with edges within the graph index. This dual-level retrieval structure combines keyword matching with relevant structural information induced from the constructed KG.

The retrieved content is the output of the LLM profiling stage which contains names, descriptions of entities and relations and short summaries from the original text.

Computation Cost of LightRAG framework

Here, an LLM is called both during indexing and querying phases similar to GraphRAG. 

During querying, LLM extracts keywords from the query to retrieve relevant entities and relations rather than entire chunks which reduces retrieval input tokens overhead compared to standard RAG. 

In the retrieval phase, LightRAG requires less than 100 tokens and a single API call, whereas in GraphRAG requires `no.of communities x avg no. tokens per community` tokens, demanding excessive API calls that could potentially hit rate limits.

For incremental data updates, GraphRAG incurs significantly higher costs due to the need for community restructuring within the KG, whereas LightRAG efficiently manages updates with minimal API calls.

Benchmarks Discussion from the Paper

The benchmark was conducted on the UltraDomain dataset, a collection of textbooks encompassing 18 distinct domains.

The authors selected the Agriculture, CS, Legal and Mix data domains alone for their experiments and they report that LightRAG is the new SOTA across various evaluation metrics like comprehensiveness, diversity, empowerment and overall response quality.

All the experiments were conducted using GPT-4o-mini with a chunk size of 1200 for consistency. This is similar to the benchmarking approach discussed in the GraphRAG paper where the LLM (GPT4o-mini) evaluates responses of LightRAG against others like HyDE, NaiveRAG and GraphRAG. The winner is chosen based on the quality of responses specific to local or global queries.

Legal Document Analysis with LightRAG

Now it’s time to do some hands-on testing. For this, we’ll use two legal documents as the legal domain presents an ideal use case for RAG. Individuals or corporate people need an effective and straightforward system to track the legal aspects of their organization or rights. Legal contexts are highly complex, and even trained legal professionals often struggle to navigate between case files to identify patterns or connect the dots between related case verdicts or judgements. In these scenarios GraphRAG principles become handy providing overarching insights reducing the cumbersome manual search which provides contextually grounded citations as nodes and edges in KG.

If you’re working in integrating GenAI in legal space, you may find valuable in going through this [Link] insightful post discussing the effectiveness of graph based approaches in legal contexts. 

Our Legal Document Analysis with LightRAG workflow will look as follows,

Click here to download the source code to this post

Code Walkthrough

To set up LightRAG, you can either clone the repository or install it via PyPI.

Installing Dependencies:

or

The following set of code is mostly adapted from LightRAG repo,

 examples/lightrag_ollama_demo.py

Setup Ollama locally in your system and download the model from the Ollama  model hub.

Before running the script it’s recommended to run ollama serve in a separate terminal so that we can monitor real-time logs and interactions with local LLM API endpoints during indexing and querying.

Import Dependencies

logging – to enable structured logging

pdfplumber – to easily extract text and data from PDFs

All the indexing files will be saved under the specified WORKING_DIR, including vector embeddings in the Vector DB storage and a JSON KV Storage for retrieval of key-value pairs.

This is the main step where the LightRAG instance is configured with necessary parameters. We are using the llama3.1 8b Q4 model as the LLM, with a context length of 32K tokens. As discussed in the benchmark section, we are sticking to a chunk size of 1200 to divide the pdf_text as manageable chunks for processing through the LLM. LightRAG’s function definitions are designed to natively integrate asynchronous methods within its pipeline enabling efficient resource utilization with minimal idle time.

• By default, Ollama models connect to a  local server at localhost:11434 which we will use to access model endpoints to get responses. The  llm_model_max_async=4 setting allows handling 4 concurrent async requests.
• For embedding the raw text chunks, nomic-embed-text is used which is a robust embedding model at its param size with an embedding dimension of 768.

Indexing

The indexing flowchart of LightRAG is as follows,

Using pdfplumber the PDF is read and its contents are saved as a single string in the pdf_text variable.

1. The kv_store_text_chunks.json stores doc text chunks with associated metadata such as chunk_size (tokens), actual content, chunk index, parent doc id etc.

2. The kv_store_llm_response.json holds the LLM generated summaries about the entity and relations. It uses a caching mechanism to prevent redundant indexing of identical ids.

3. Next, vdb_entities.json contains vector embeddings of entities extracted from the text chunks with unique id, entity name etc.

4. Similarly, the vdb_relationship.json file stores relationships between entities linking source and target entities by their IDs to represent connections.

5. Finally the graph index constructed is saved as graph_chunk_entity_relation.graphml, which we will utilize later to visualize the KG in Neo4j.

Incremental Indexing

To index new files, simply change the pdf_path. This will automatically update the LLM responses in kv_store_llm_response.json to include any new entities and relationships without any conflicts or redundancies with existing data. 

The following is an example of an updated response.  

On an RTX 3080 GPU with 12GB VRAM and an intel i7 13th Gen CPU,  it took about an hour to index the 288-page Companies Act PDF which had 190 total chunks and an additional 70 minutes for updating the Constitution of India PDF. Based on our testing, the llama3.1:8b model works well without many issues; however when the qwen2.5:7b or llama3.2:1b family of models is used, we encounter several issues in the indexing process like slowdowns or halts. While the reason remains unclear we appreciate sharing any insights you may have in the comments.

Querying

Depending on the query mode such as naive, local, global and hybrid, relevant keywords are extracted from the query and are compared with the KV store and Nano vector DB to retrieve the candidate entities and relationships based on cosine similarity.

The querying phase of LightRAG is outlined in the following flowchart.

We will save the LLM output responses for different query approaches in an output_response.txt file for referencing outputs anytime.

Naive RAG

This is a standard RAG setup retrieving top_k = 60 chunks.

Local Query – LightRAG

By default the pipeline retrieves top_k = 60 candidate entities for local query mode.

Global Query – LightRAG

For the global query model, LightRAG pipeline retrieves top_k = 60 candidate relations.

Hybrid Query – LightRAG

The hybrid mode combines the high level and low level context to generate better responses targeted at both local entities and global relations.

To evaluate the win rate of responses to different queries targeting local entities and global relationships, we can follow a similar process discussed in both the GraphRAG and LightRAG papers, where GPT-4o was used as a judge. Additionally, we will manually review the responses to derive further insights.

Response Comparison 1: Targeted at Local Entity

To prevent OOM error the following query was inferred with {"num_ctx": 16000}

Query:  “What are the requirements for an independent director in Section 149?”

Note: The original page from the text was added from the PDF for reference and is independent of LightRAG pipeline

Naive RAG

Local Query – LightRAG Response

Global Query – LightRAG Response

Hybrid (Local + Global) Query – LightRAG Response

GPT-4o review about the response and ratings is as follows:

Overall Rating Summary by GPT-4o:

• Naive: 4/10
• Local: 9/10
• Global: 6/10
• Hybrid: 8.5/10

The Local Response is the most effective, with the Hybrid Response following closely. Both responses address the query well, with Local Response being particularly concise and informative.

Response Comparison 2: Targeted at Global Relations

{"num_ctx": 16000}

Query:  “How are workers’ rights protected across both constitutional and corporate legal frameworks in India?”

Naive RAG Response

Local Query – LightRAG Response

Global Query – LightRAG Response

Hybrid (Local + Global) Query – LightRAG Response

Overall Ratings Summary by GPT-4o:

• Naive: 7/10
• Local: 8.5/10
• Global: 8/10
• Hybrid: 9/10

The Hybrid Response is the most effective, combining detailed information on protections, challenges, and reform efforts, followed closely by the Local Response for its in-depth and structured approach. Both responses effectively address the query from a global perspective.

Even upon manual review we see that hybrid responses in both sections were better in most cases. Hybrid retrieval is one of the key strengths of the LightRAG framework.

Knowledge Graph Visualization with Neo4j

Uploading the .graphml to Neo4j cloud is straightforward, you can find the code for this in the repository. Simply personalize the credentials with your specific details by signing up for a free Neo4J AuraDB account. After upserting into the Neo4j instance, for those two PDF contents we had Nodes: 3,352 and Relationships: 4,752.

Sample Cypher Query in Neo4J Console:

where, () represents nodes (or entities), while [] represents relationships between them.

Reproduce Benchmarks

To reproduce the benchmarks results, the team has also released evaluation scripts. If you are interested, simply download the UltraDomain dataset and follow the instructions outlined in the README file.

Takeways

• As expected, LightRAG’s Graph Indexing and dual-level retrieval yielded promising results as rightly outlined in the paper. In our comparison, when querying a specific entity, the local query response was the most desirable, followed by the hybrid approach. In a second comparison focused on global relationships, the global and hybrid modes produced the most preferred summaries.
  
• At times, queries targeted at global and local contexts didn’t work as expected, and the naive approach performed better, while other responses simply stated, “Sorry, I can’t answer this question as I don’t have enough information.” . However, subject to further experiments optimal querying strategy is to be found.

• While LightRAG is a promising project, its high level of abstraction can make it challenging to follow through logs alone. The community has suggested adding more intermediate logging, and the authors have acknowledged this feedback. If you are getting started with LightRAG, it is highly recommended to review the internal codebase, the workflow image in the README, and both the GraphRAG and LightRAG papers for a better understanding.

GenAI Powered Legal Services

• Harvey AI
• SALI Alliance
• Relevance AI
• Lyzr AI
• A&O Shearman
• DocDraft
• Sarvam A1 Legal

Conclusion

Based on our experiments, we find LightRAG to be a promising and effective RAG framework. The legal domain, by nature, poses challenging queries that demand accurate insights without hallucinations, and we observed that LightRAG demonstrates strong potential in meeting these demands.

Although still in its early stages, there is still room for further enhancements. The authors are actively maintaining the repository, continuously implementing updates and addressing issues. Kudos to the team at HKU for releasing this as an open-source project under the MIT license.

References

• LightRAG Github
• LightRAG Arxiv
• GraphRAG Arxiv
• Prompt Engineer– Youtube
• Zilliz – Milvus DB
• Companies Act ,2013 PDF 
• Constitution of India PDF
• GraphRAG Blog: Cohesion Force