見出し画像

【実装付きAI論文解説】LightRAG: Simple and Fast Retrieval-Augmented Generation｜TECH BLOG #04

AIVALIX株式会社




1. はじめに

はじめまして、AIVALIX株式会社 Chief AI Engineer を務めております、遠藤大貴と申します。私は機械学習を基盤にプログラミングスキルを磨き上げ、当社ではモバイルアプリの保守やデータSQLを活用したバックエンド開発、さらに生成AIの実装など、多岐にわたるプロジェクトに携わってまいりました。

弊社では、AIエンジニアによるテックブログの発信に力を注いでおり、最新のAI論文を実装例を交えながら分かりやすく解説することを目指しています。AIの技術や可能性をより多くの方に感じていただくために、こうした情報を継続的に発信していく予定です。少しでも皆さまの学習や研究の一助となれば幸いです。

さて、今回テックブログ第4回としてご紹介するのは、ナレッジグラフを活用したRetrieval-Augmented Generation（RAG）の一手法であるLightRAGです。

昨年、Microsoft社がナレッジグラフを組み込んだGraphRAGを発表し、従来のRAG手法の課題であった文脈理解の精度を大きく改善させたことによって注目されている手法です。本稿では、LightRAGの紹介に加えて、ナレッジグラフを適用した場合と従来のRAG手法との比較も紹介させていただきます。


2. LightRAG 論文解説

RAG（Retrieval-Augmented Generation）とは、大規模言語モデル（LLM）に外部データを付加して、モデルが学習していない領域の質問にも答えられるようにする方法です。専門的な知識や社内文書を追加することで、幅広いタスクに応用できます。しかし従来のRAG手法は、文章内で複雑に関連し合う情報同士の関係性を十分に考慮できないという課題がありました。

この課題を解決するためにナレッジグラフが導入されました。ナレッジグラフを使うことで、情報同士の関係性を整理し、文章間の文脈を理解しながら学習できるようになります。結果として、クエリ（質問）に対して、より文脈に即した的確な回答が可能になります。
また、LightRAGはナレッジグラフの導入以外にも、検索の効率化や柔軟な回答の生成を可能にする手法が使用されています。

2-1. Methods

ナレッジグラフの構築
テキストデータから動的にエンティティ（ノード）とリレーション（エッジ）を抽出し、重複を削除しつつグラフ構造を構築します。エンティティやリレーションは特定のプロンプトに基づいてLLMが抽出を行い、類似度の高いエンティティやリレーションとの関係をナレッジグラフに反映させます。また各要素は、元のソースとなる情報と紐づいているため、グラフを参照することで、高速化つ正確な検索、情報収集が可能です。LightRAGは、ナレッジグラフを再構築することなく新たなノードとエッジを追加できるため、コストを抑えつつ情報を更新・追加できます。

デュアルレベル検索
低レベル検索と高レベル検索を同時に行うことで、さまざまな形式のクエリ（質問）に柔軟に対応することが可能になり、クエリの文脈把握能力を高めます。
また、ナレッジグラフ内のエンティティとリレーションを検索するため、高速な検索が可能になります。

手法
まず、クエリからLLMを活用して具体的なキーワードと抽象的なキーワードを抽出します。
低レベル検索では、具体的なキーワードを、特定のエンティティやリレーションを対象にコサイン類似度によって検索しています。一方、高レベル検索では、抽象的なキーワードを、複数の関連するエンティティやリレーション（概念・分野）を対象にコサイン類似度によって検索します。
また、より広範囲に関連性を学習し、検索されたエンティティの近傍ノードも含めて（1hop以上離れている）情報も収集することでより文脈を正確に把握することが可能になります。

ナレッジグラフを活用するLLMの回答
ナレッジグラフからデュアルレベル検索を行い、LLMによってクエリへの回答を行います。
LightRAGのナレッジグラフは高速検索を想定して構築されているため、チャンクの検索ではなくエンティティと関係の検索を行い、検索オーバーヘッドを削減し、高速な検索、応答を行うことができます。


2-2. Results

データセット
データセットは以下のデータを使用して検証が行われています。

UltraDomainベンチマーク（Qian et al., 2024）
（農業、コンピュータサイエンス、法学、混合分野）

  • Agriculture: 12文書、約2,017,886トークン。

  • CS: 10文書、約2,306,535トークン。

  • Legal: 94文書、約5,081,069トークン（最大規模）。

  • Mixed: 61文書、約619,009トークン。

実装と評価
論文では以下のツール、評価法が採用されています。

  • ベクトルデータ管理にはNanoベクトルデータベースを使用。

  • すべてのLLM操作にはGPT-4o-miniを使用。

  • チャンクサイズはすべてのデータセットで1200に統一。

  • 各分野に関し、クエリが5つ生成され、回答を生成する。

  • LLMによって、包括性、多様性、有用性の観点から2つのRAGを比較し、二者の勝率を算出し評価。


大規模コーパスおよび複雑なクエリを扱っているため、LightRAGやGraphRAGのようなナレッジグラフを使用しているRAGシステムは、単純なチャンクベースの検索手法（従来のRAG手法であるNaiveRAG、HyDE、RQ-RAG）を一貫して上回る結果が出ました。
評価項目のうちの多様性において顕著な優位性が見られることから、柔軟な回答がデュアルレベル検索によって実現されていると考察できます。


アブレーションスタディ
低レベル検索、高レベル検索、スニペット（参照した元データ）をそれぞれ使用しなかった場合と比較した検証が行われています。
全体的に、どの評価項目も全て使用しているLightRAGが上回っているものの、スニペットを除いたほうが評価されている分野が見られました。これはスニペットに含まれるノイズが除去されたことが起因していると考えられます。


2-3. Conclusion & Comment

LightRAGの注目する特徴としては、多様なクエリに対応できる柔軟性をもち、ナレッジグラフを参照しているものの検索速度が速い点が挙げられます。評価として、LLMによる勝率算出を使用している点が少し疑問に感じつつも出力結果の精度や推論の一貫性については高い印象を受けました。


3. LightRAGの実装と比較

今回はLightRAGの実装、そして従来のRAG手法との回答比較を行いました。

コードはGitHubに公開されており、実行環境はGoogleColabで行いました。

生成モデル、埋め込みモデルは、

  • meta-llama/Llama-3.1-8B-Instruct

  • sentence-transformers/all-MiniLM-L6-v2

を使用しています。

どちらもHugging Faceから使用できるモデルになります。

モデルのロード

!git clone <https://github.com/HKUDS/LightRAG>
%cd LightRAG
!pip install -e .

クエリは「ニューラルネットワークと機械学習の関係性を教えて」とし、テキストデータとしては、WikiPediaから「ニューラルネットワーク、機械学習、大規模言語モデル、生成的人工知能」の記事を引用して学習させています。

学習させたいテキストデータを学習させ、回答を出力。

import asyncio
import os
import inspect
import logging
from lightrag import LightRAG, QueryParam
from lightrag.llm import hf_model_complete, hf_embedding
from transformers import AutoModel, AutoTokenizer
from lightrag.utils import EmbeddingFunc

nest_asyncio.apply()

logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)

WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)

# Initialize LightRAG with Hugging Face model
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=hf_model_complete,  # Use Hugging Face model for text generation
    llm_model_name='meta-llama/Llama-3.1-8B-Instruct',  # Model name from Hugging Face
    # Use Hugging Face embedding function
    embedding_func=EmbeddingFunc(
        embedding_dim=384,
        max_token_size=5000,
        func=lambda texts: hf_embedding(
            texts,
            tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
            embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        )
    ),
)

# your_text_data を自分の使用したいテキストデータの名前に変更
combined_file_path = "./your_text_data"
if not os.path.exists(combined_file_path):
    raise FileNotFoundError(f"{combined_file_path} が見つかりません。ファイルが存在するか確認してください。")

with open(combined_file_path, "r", encoding="utf-8") as f:
    rag.insert(f.read())

# クエリを実行して結果を出力
def execute_queries():
    queries = [
        "ニューラルネットワークと機械学習の関係性を教えて"
    ]
    modes = ["hybrid"]

    for mode in modes:
        print(f"\n--- Query Mode: {mode} ---")
        print(
            rag.query(queries[0], param=QueryParam(mode=mode))
        )

    print("\n--- Stream Response (Hybrid) ---")
    resp = rag.query(
        queries[0],
        param=QueryParam(mode="hybrid", stream=True),
    )

    async def print_stream(stream):
        async for chunk in stream:
            print(chunk, end="", flush=True)

    if inspect.isasyncgen(resp):
        asyncio.run(print_stream(resp))
    else:
        print(resp)

execute_queries()


3-1. LightRAGによる出力結果


取得された関連文書

High-level keywords:

  • ニューラルネットワークと機械学習 (Artificial Intelligence,Machine Learning)

  • 関係性 (Relationship, Connection)

Low-level keywords:

  • ニューラルネットワーク (Neural Network)

  • 機械学習 (Machine Learning)

  • 人工知能 (Artificial Intelligence)

  • ディープラーニング (Deep Learning)

  • AI (Artificial Intelligence)

生成された回答

ニューラルネットワーク
ニューラルネットワークは、人工知能の分野で広く知られているものであり、生物の学習メカニズムを模倣した機械学習手法です。ニューラルネットワークは、人工知能の分野で広く使用されているものであり、分類・回帰・生成など様々なクラスのタスクに教師あり/教師なし問わず利用されます。

機械学習
機械学習は、ニューラルネットワークを含む人工的な知能を指す人工知能の分野です。機械学習は、ニューラルネットワークを含む人工的な知能を指すため、ニューラルネットワークは機械学習の重要な部分です。機械学習は、機械がデータから学習し、予測や決定を行う技術です。

ニューラルネットワークと機械学習の関係性
ニューラルネットワークは、機械学習の分野で広く使用されている技術です。ニューラルネットワークは、機械学習のモデルとして利用され、分類・回帰・生成など様々なクラスのタスクに教師あり/教師なし問わず利用されます。ニューラルネットワークは、機械学習の分野で広く使用されているため、機械学習の重要な部分です。

生成されたナレッジグラフ（一部）


3-2. 従来のRAG手法による出力結果

従来のRAG手法を用いて同様のクエリに対する回答を比較します。

埋め込み、生成モデルはLightRAGと同様のモデルを使用し、関連度の高い上位5文書を利用してLLMが回答を生成しています。

関連文書

  • 機械学習の評価指標

  • 機械学習における変数の種類

  • それに加えて、学習の時間複雑性と実現可能性についても研究している。計算論的学習理論では、多項式時間で終了する計算を実現可能とみなす。

  • 半教師あり学習

  • 統計的機械学習

生成された回答

ニューラルネットワークは、機械学習の重要な手法の1つであり、特定の入力から出力を予測する能力を持つ人工知能のモデルです。ニューラルネットワークは、ニューロン（ニューロン）を模した構造を持っており、入力データを処理して出力データを生成します。これにより、ニューラルネットワークは複雑なデータを分析し、パターンを認識し、予測を行うことができます。機械学習では、ニューラルネットワークを使用して、データからパターンを学び、モデルをトレーニングし、精度を向上させることができます。ニューラルネットワークは、画像認識、自然言語処理、時間系列予測などの多くの分野で活用されています。

3-3. 考察

LightRAGの回答は、ナレッジグラフを活用して「ニューラルネットワーク」と「機械学習」の関係性にフォーカスした簡潔な回答を提示しており、冗長性が低い反面、横方向の情報展開は少ない回答形式でした。

一方、従来のRAGの回答は、類似度検索をベースとして関連文書を集めるため、回答に含まれる情報の広がりが大きい回答形式でした。（ニューラルネットワークに関する情報が回答の多くを占めています。）


4. 最後に

ナレッジグラフの組み込みは、従来のRAGでは困難だった文脈理解を改善し、LLMによる推論能力をより強力にすることを可能にしました。今後も、LLMの推論能力を向上させる様々な技術に注目したいです。

今回のテックブログは以上になります。最後までお読みいただき、ありがとうございました。この記事が、少しでも皆様の参考になり、新たな気づきや挑戦のきっかけとなれば幸いです。今後も引き続き、有益な情報をお届けしてまいりますので、どうぞご期待ください。


AIVALIX株式会社では、AIに関する開発やR&Dの案件を常時承っております。また、共同開発・共同研究にご協力いただける企業様も広く募集しております。

AI技術を通じた新たな価値創造に向け、一緒に挑戦していただけるパートナーを心よりお待ちしております。ご興味をお持ちの際は、ぜひお気軽にお問い合わせください。

▼ コーポレートサイトはこちら ▼


いいなと思ったら応援しよう！