Treehop Rag

🚀 TreeHop: マルチホップ質問応答のための次のクエリ埋め込みを効率的に生成およびフィルタリングする

TreeHopは、Retrieval-Augmented Generation (RAG)の分野における従来の再帰的検索パラダイムの計算効率の低さを解決するために設計された軽量の埋め込みレベルのフレームワークです。反復的なLLMベースのクエリ書き換えを不要にすることで、TreeHopは最先端の性能を維持しながら待ち時間を大幅に削減します。これは、動的なクエリ埋め込みの更新と枝刈り戦略を通じて、効率的な「Retrieve-Embed-Retrieve」ワークフローを可能にします。

🚀 クイックスタート

TreeHopを使用するには、まずシステム要件を満たし、必要なデータベースを用意する必要があります。以下のセクションで詳細を説明します。

✨ 主な機能

• 複雑なクエリの処理：現実世界の質問は関連情報を取得するために複数のホップを必要とすることが多く、従来の検索方法では対応が困難ですが、TreeHopはこれを効率的に処理できます。
• コスト効率：既存のクエリリライターが数十億のパラメータを持つのに対し、TreeHopは2500万のパラメータで済み、計算オーバーヘッドを大幅に削減します。
• 高速性：反復的なLLMアプローチと比較して99%高速な推論が可能で、応答速度が重要な産業アプリケーションに最適です。
• 高性能：取得するパッセージの数を制御しながら高いリコール率を維持し、システムを圧迫することなく関連性を確保します。

📦 インストール

システム要件

Ubuntu 18.06 LTS以上またはMacOS Big Sur以上。 Nvidia GPUまたは32GB以上のRAMを持つApple Metal。 再現には16GBのシステムRAM、トレーニングには64GBのシステムRAMが必要です。 ハードドライブに50GBの空き容量が必要です。

Python環境

Python環境については、requirements.txtを参照してください。

事前準備

このリポジトリには、再現目的で評価用の埋め込みデータベースが含まれています。リポジトリをクローンするには、git LFSを有効にして、git lfs clone [リポジトリのリンク]を実行するか、既存のローカルリポジトリの下で以下のコマンドを実行してデータを取得してください。

git lfs pull

埋め込みデータベース

埋め込み生成にはBGE-m3を採用し、この上でTreeHopモデルをマルチホップ検索用にトレーニングします。以下の2つのスクリプトを実行すると、すべての必要なトレーニング用および評価用の埋め込みデータベースが生成されます。TreeHopをトレーニングする予定がない場合は、すべての必要な評価用埋め込みデータベースがリポジトリに提供されているため、これらのスクリプトを実行する必要はありません。

python init_train_vectors.py
python init_multihop_rag.py

💻 使用例

基本的な使用法

以下の例では、MultiHop RAG評価データセットを使用します。リポジトリには、この例を実行するために必要なファイルが含まれています。詳細は事前準備を参照してください。

from tree_hop import TreeHopModel
from passage_retrieval import MultiHopRetriever


EVALUATE_DATASET = "multihop_rag"

# HuggingFaceからTreeHopモデルをロード
tree_hop_model = TreeHopModel.from_pretrained("allen-li1231/treehop-rag")

# リトリーバーをロード
retriever = MultiHopRetriever(
    "BAAI/bge-m3",
    passages=f"embedding_data/{EVALUATE_DATASET}/eval_passages.jsonl",
    passage_embeddings=f"embedding_data/{EVALUATE_DATASET}/eval_content_dense.npy",
    # faissインデックスが初期化されている場合はこの行のコメントを解除すると、高速にロードできます
    # faiss_index=f"embedding_data/{EVALUATE_DATASET}/index.faiss",
    tree_hop_model=tree_hop_model,
    projection_size=1024,
    save_or_load_index=True,
    indexing_batch_size=10240,
    index_device="cuda"     # またはApple Metalではcpu
)

⚠️ 重要提示

• パッセージのjsonlファイルには、検索データベース内の各パッセージのID、タイトル、テキストが含まれています。
• passage_embeddingsファイルとfaiss_indexファイルには、それぞれnumpy配列形式とfaissインデックス形式のパッセージ埋め込みが含まれています。独自のデータベースで置き換えるには、init_multihop_rag.pyのロジックを参照してください。
• パッセージファイルの詳細な構造については、MultiHop RAG評価パッセージファイルを参照してください。

retrieverには、単一クエリとバッチクエリの両方を取得できるmultihop_search_passagesメソッドがあります。

単一クエリの場合

retrieve_result = retriever.multihop_search_passages(
    "Did Engadget report a discount on the 13.6-inch MacBook Air \
        before The Verge reported a discount on Samsung Galaxy Buds 2?",
    n_hop=2,
    top_n=5
)

バッチクエリの場合

LIST_OF_QUESTIONS = [
    "Did Engadget report a discount on the 13.6-inch MacBook Air \
        before The Verge reported a discount on Samsung Galaxy Buds 2?",
    "Did 'The Independent - Travel' report on Tremblant Ski Resort \
        before 'Essentially Sports' mentioned Jeff Shiffrin's skiing habits?"
]

retrieve_result = retriever.multihop_search_passages(
    LIST_OF_QUESTIONS,
    n_hop=2,
    top_n=5,
    # デバイスに合わせてバッチサイズを変更してパフォーマンスを最適化します
    index_batch_size=2048,
    generate_batch_size=1024
)

取得したパッセージと対応するマルチホップ検索パスにアクセスするには、以下のコードを使用します。

# 質問に対する取得パッセージを表示
print(retrieve_result.passage)

# networkxグラフを使用してマルチホップ検索を描画
retrieve_result = retriever.multihop_search_passages(
    LIST_OF_QUESTIONS,
    n_hop=2,
    top_n=5,
    index_batch_size=2048,
    generate_batch_size=1024,
    return_tree=True        # この引数を追加するだけです
)
# `retrieve_result.tree_hop_graph` は、LIST_OF_QUESTIONS内のクエリの検索パスに対応するnetworkxオブジェクトのリストです。
# 最初のクエリを例にとって、対応するパスを描画するには、以下のコードを実行します。
retrieval_tree = retrieve_result.tree_hop_graph[0]
retrieval_tree.plot_tree()

# ノードは検索グラフ内のパッセージを表します
# 元のパッセージのメタデータを保存します。
print(retrieval_tree.nodes(data=True))

📚 ドキュメント

論文再現

TreeHopのマルチホップ検索性能を評価するには、以下のコードを実行します。ここでは、2WikiMultihopデータセットと3ホップでのrecall@5を例に挙げています。このスクリプトは、各ホップでのリコール率と取得パッセージの平均数、および質問のタイプ別の統計を出力します。

python evaluation.py \
    --dataset_name multihop_rag \
    --revision paper-reproduction \
    --n_hop 3 \
    --top_n 5 \
    --redundant_pruning True \
    --layerwise_top_pruning True

⚠️ 重要提示

• 評価データセットを変更するには、2wikiをmusiqueまたはmultihop_ragに置き換えてください。
• n_hopとtop_nを変更して、ホップ数と上位検索設定を変更できます。
• redundant_pruningとlayerwise_top_pruningを切り替えて、停止基準に関するアブレーション研究を再現できます。

TreeHopのトレーニング

以下のコードを実行して、グラフを生成し、TreeHopをトレーニングします。このスクリプトの引数については、training.pyのparse_args関数を参照してください。トレーニング用の埋め込み生成については、init_train_vectors.pyのコードを参照してください。

python training.py --graph_cache_dir ./train_data/

📄 ライセンス

このプロジェクトはMITライセンスの下で公開されています。詳細については、LICENSEを参照してください。

🔧 技術詳細

TreeHopは、Retrieval-Augmented Generation (RAG)の分野における従来の再帰的検索パラダイムの計算効率の低さを解決するために設計された軽量の埋め込みレベルのフレームワークです。反復的なLLMベースのクエリ書き換えを不要にすることで、TreeHopは最先端の性能を維持しながら待ち時間を大幅に削減します。これは、動的なクエリ埋め込みの更新と枝刈り戦略を通じて、効率的な「Retrieve-Embed-Retrieve」ワークフローを可能にします。

📚 引用

@misc{li2025treehopgeneratefilterquery,
      title={TreeHop: Generate and Filter Next Query Embeddings Efficiently for Multi-hop Question Answering}, 
      author={Zhonghao Li and Kunpeng Zhang and Jinghuai Ou and Shuliang Liu and Xuming Hu},
      year={2025},
      eprint={2504.20114},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2504.20114}, 
}

📄 情報テーブル