🤖 machine learning

TreeHop: Efficient Embedding-Level Query Rewriter

TreeHop is a novel, LLM-free embedding-level framework for multi-hop question answering that dynamically fuses semantic information from prior queries and retrieved documents to streamline the retrieval process, achieving performance comparable to state-of-the-art methods while drastically reducing computational costs and latency.

Original authors: Zhonghao Li, Kunpeng Zhang, Jinghuai Ou, Shuliang Liu, Xuming Hu

Published 2026-08-11

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhonghao Li, Kunpeng Zhang, Jinghuai Ou, Shuliang Liu, Xuming Hu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, multi-layered mystery, like figuring out who the grandfather of a famous celebrity is. You can't just ask one question and get the answer; you have to follow a trail of clues. First, you find out who the celebrity's father is. Then, you have to take that new name and ask, "Who is his father?" This is called "multi-hop" reasoning. In the world of artificial intelligence, computers use a system called Retrieval-Augmented Generation (RAG) to do this. Think of RAG as a super-smart librarian who doesn't just memorize facts but can run to the shelves to find the right books to help answer your questions.

However, there's a catch. When the computer needs to follow a long trail of clues, the old way of doing things is incredibly slow and expensive. It's like asking a human detective to read a whole book, write a new question on a piece of paper, hand that paper to a different detective, and then have that detective run to the library again. Every time the detective has to "rewrite" the question using a powerful brain (a Large Language Model), it takes a lot of time and energy. If you have to do this ten times to solve a mystery, you're waiting a long time and burning a lot of electricity. Scientists have been looking for a way to make this detective work faster without losing the ability to solve the hardest puzzles.

This is where a new method called TreeHop comes in. The researchers behind TreeHop realized that you don't need a full-blown detective to rewrite every single question. Instead, you can just tweak the "map" the computer is using to find the books. Imagine the computer's search query as a glowing, magical compass needle. In the old system, every time you found a clue, you had to stop, melt the compass, and forge a brand new one using a heavy, slow furnace. TreeHop is like a magical, instant-adjustment knob. It takes the old compass needle and the new clue you just found, and it instantly fuses them together to point in the right direction for the next step. It does this by mixing the "vibes" (mathematical embeddings) of the old question and the new clue, skipping the slow furnace entirely.

The paper shows that this new "knob" works incredibly well. By using this embedding-level update, TreeHop can solve these multi-step mysteries just as accurately as the slow, heavy methods, but it is vastly faster. In fact, the researchers found that TreeHop is about 92.8% to 97.8% faster than the current top methods that rely on rewriting questions with big AI models. It also uses a tiny fraction of the computer power—only 2.2% to 29.4% of the size of other solutions. To make sure the computer doesn't get lost in an endless loop of checking the same clues, TreeHop also has a smart "stop rule." It prunes away paths that lead to dead ends or places it has already visited, keeping the search tree neat and efficient.

In short, TreeHop suggests that we don't need to call in the heavy artillery (massive AI models) for every little step of a search. Instead, we can use a lightweight, clever mechanism to nudge the search direction in real-time. The results suggest this is a game-changer for making AI systems that are not only smart enough to solve complex puzzles but also fast and cheap enough to use in real-world applications, like on a phone or in a busy customer service chatbot.

Technical Summary: TreeHop: Efficient Embedding-Level Query Rewriter

Problem Statement

Retrieval-Augmented Generation (RAG) systems struggle with Multi-Hop Question Answering (MHQA), where answers require synthesizing information across multiple document chunks. Conventional approaches rely on iterative loops involving Large Language Models (LLMs) to rewrite queries based on retrieved context (e.g., "Retrieve-Rewrite-Vectorize-Retrieve"). This process incurs significant computational costs and latency due to repeated LLM invocations and multi-stage processing. Furthermore, existing efficient methods often fail to balance the trade-off between computational overhead and retrieval effectiveness, particularly in low-resource or latency-sensitive deployment scenarios.

Methodology: TreeHop

TreeHop proposes a novel framework that performs query refinement entirely at the embedding level, eliminating the need for LLM-based query rewriting. The core mechanism operates through the following components:

1. Embedding-Level Query Update

Instead of generating new text queries, TreeHop dynamically updates the query embedding by fusing semantic information from the previous query and retrieved documents. The update rule is defined as:
$q_{r+1}^i = q_r - \text{ForgetGate}(q_r, c_r^i) + \text{UpdateGate}(q_r, c_r^i)$

ForgetGate: Uses a cross-attention mechanism to identify and suppress query semantics that have already been resolved by the retrieved chunk $c_r^i$ . This prevents redundant retrieval of information already captured.
UpdateGate: Selectively incorporates new, relevant information from the retrieved chunk to form the next-step query embedding.
Architecture: Both gates utilize a modified cross-attention mechanism that operates along the embedding dimensions rather than token positions, as the input consists of fixed-dimensional vectors from a frozen encoder (BGE-m3). This allows the model to scale coordinates of the embedding vector to discard or inject specific semantic features.

2. Pruning Strategies

To prevent the exponential growth of retrieval branches inherent in multi-hop systems, TreeHop employs two rule-based stopping criteria:

Redundancy Pruning: Terminates paths where the retrieved document chunk has already been seen in previous iterations.
Layer-wise Top-K Pruning: At each retrieval layer, only the top- $K$ chunks with the highest similarity scores across all generated query embeddings are retained. This reduces the branching factor from $O(n^r)$ to a manageable linear growth.

3. Training and Data Construction

TreeHop is trained using contrastive learning (InfoNCE loss). The training data is constructed as Directed Acyclic Graphs (DAGs) representing multi-hop reasoning trajectories.

Data Pipeline: Single-hop questions are linked via bridge entities to form multi-hop queries. An LLM validates the reasoning topology and evidence integrity, ensuring the graph structure is correct.
Objective: The model learns to predict the next-step query embedding given a query-chunk pair, treating the correct next-hop evidence as positive targets and other corpus passages as negatives.
Efficiency: The model is compact (0.18 billion parameters) and can be trained on a single GPU.

Key Contributions

The paper identifies three primary contributions:

Novel Embedding-Updating Mechanism: Replaces LLM-driven iterative rewrites with lightweight neural operations (gated cross-attention on embeddings), enabling efficient computational complexity for MHQA tasks when combined with pruning.
Efficient Rule-Based Stopping Criterion: Introduces redundancy and layer-wise top- $K$ pruning to control branching factor growth while maintaining retrieval performance.
Empirical Validation: Demonstrates superior efficiency-accuracy trade-offs across four MHQA benchmarks (HotpotQA, 2WikiMultiHop, MuSiQue, MultiHop RAG).

Experimental Results

TreeHop was evaluated against baselines (Direct Retrieval) and advanced RAG solutions (Iter-RetGen, EfficientRAG, RAPTOR) across four datasets.

Efficiency: TreeHop reduces query latency by 92.8%–97.8% compared to LLM-based methods (e.g., Iter-RetGen) and significantly outperforms other efficient models. It achieves this with only 2.2%–29.4% of the parameter size of concurrent solutions.
Retrieval Performance:
- On 2WikiMultiHop and MultiHop RAG, TreeHop surpasses Iter-RetGen (backed by Llama3.1-8B) by 3.5%–5.2% in recall at the second iteration and 4.9%–9.1% at the third iteration, with slightly fewer chunks retrieved on average.
- On HotpotQA, it leads all iterative solutions in the second iteration and trails only RAPTOR at the third.
- On MuSiQue, it performs slightly lower than Iter-RetGen (0.4% difference), attributed to the dataset's requirement for synthesizing information from multiple chunks simultaneously, a limitation of the query-chunk-pair strategy.
End-to-End QA: When paired with downstream LLMs, TreeHop achieves competitive or superior accuracy, particularly on 2WikiMultiHop and MultiHop RAG, due to its ability to complete evidence chains effectively.

Significance and Claims

The paper claims that TreeHop offers a faster and more cost-effective solution for low-resource or latency-sensitive deployments. By collapsing the traditional "Retrieve-Rewrite-Vectorize-Retrieve" cycle into a streamlined "Retrieve-Embed-Retrieve" loop, TreeHop bridges the gap between computational efficiency and retrieval effectiveness.

The authors position TreeHop not as a replacement for all RAG methods, but as a scalable alternative for iterative retrieval that avoids the heavy computational overhead of LLM generation at every hop. The work suggests that embedding-centric rewriting is a practical balance for real-world multi-hop reasoning, particularly where latency and parameter efficiency are critical constraints. The paper acknowledges limitations, specifically that the current architecture relies on query-chunk pairs and may struggle with synthesizing information across multiple chunks simultaneously compared to methods that can process broader context windows.