TreeHop: Efficient Embedding-Level Query Rewriter

Retrieval-Augmented Generation (RAG) systems face significant challenges in multi-hop question answering (MHQA), where complex queries require synthesizing information across multiple document chunks. Conventional approaches to this problem typically rely on iterative LLM-based query rewriting and routing, resulting in high computational costs due to repeated LLM invocations and multi-stage processes.

To address these limitations, the authors propose TreeHop, an embedding-level framework without the need for LLMs in query refinement. TreeHop dynamically updates query embeddings by fusing semantic information from prior queries and retrieved documents, enabling iterative retrieval through embedding-space operations alone. This method collapses the traditional 'Retrieve-Rewrite-Vectorize-Retrieve' cycle... into a streamlined 'Retrieve-Embed-Retrieve' loop, significantly reducing computational overhead.

The core of the TreeHop architecture is a pair of cross-attention gates, the ForgetGate and the UpdateGate, which together modulate how information from prior queries q and retrieved chunks c is discarded or is retained. The update rule is defined as:

q r+1 = q r - ForgetGate(q r, c ir) + UpdateGate(q r, c ir)

Specifically, the ForgetGate term suppresses the semantics that the retrieved chunk has already resolved, preventing redundant retrieval of information already captured. Conversely, the UpdateGate selectively incorporates new information from c ir to form the next query. Because the frozen encoder emits a single vector per query and per chunk, the model replace[s] the dot products with element-wise products and take[s] the softmax along the d embedding dimensions instead, so that each gate produces one weight per dimension which scales its value vector independently.

To maintain efficiency and prevent unchecked iteration [that] grows the retrieved set exponentially, TreeHop introduces two rule-based stopping criteria:

1. Redundancy Pruning: This terminates paths where the document chunk has been retrieved in the previous iterations.

2. Layer-wise Top-K Pruning: This retains only the top-K chunks with highest similarity scores across all generated query embeddings, which reduces the branching factor from O(n r) to O(nr).

The model is trained using contrastive learning on a dataset constructed through a three-stage pipeline consisting of Multi-Hop Composition, Topology and Evidence Integrity Validation, and the construction of a Contrastive Learning Graph. The training utilizes the BGE-m3 multilingual embedding model, which remains frozen during training.

Experimental results indicate that "TreeHop rivals advanced RAG methods across four open-domain MHQA datasets, achieving comparable performance with only 2.2%-29.4% of the parameter size of concurrent solutions and reducing the query latency by 92.8%-97.8%. The authors note that TreeHop reduces retrieval latency by 92.8%-97.8% compared to LLM-based methods while sacrificing only 0.9% of the recall rate at maximum."

A limitation of the framework is that TreeHop’s query-rewriting mechanism depends on query-chunk pair embeddings, limiting its capability to synthesize information across multiple chunks simultaneously.

1. Multi-Chunk Cross-Attention Synthesis

• Improvement: Replace the pairwise query-chunk update rule with a multi-chunk aggregation mechanism that uses a transformer-style cross-attention layer to fuse a set of K retrieved chunks into a single contextual representation before applying the ForgetGate and UpdateGate.

• Capability: The AI system can perform simultaneous synthesis, allowing it to resolve complex queries that require comparing, contrasting, or aggregating information spread across multiple disparate documents (e.g., Compare the fiscal policies of the US and EU during the 2008 crisis) rather than just following a linear, one-chunk-at-a-time reasoning path.

2. Tiered Hybrid Reasoning Orchestration

• Improvement: Implement a two-tier architecture where the TreeHop embedding-level framework acts as a high-speed reasoning scout to navigate the document space and prune irrelevant branches, followed by a single, high-capacity LLM invocation only for the final answer synthesis.

• Capability: The AI system can provide LLM-grade reasoning quality for multi-hop questions while maintaining the extreme low latency and low computational cost of an embedding-only model, making complex RAG viable for real-time edge computing and high-throughput enterprise applications.

3. Entropy-Based Adaptive Pruning

• Improvement: Enhance the Layer-wise Top-K Pruning by replacing static similarity scores with an information-gain metric that calculates the semantic entropy or novelty of a chunk relative to the current query embedding q r.

• Capability: The AI system can dynamically adjust its search breadth based on query complexity, preventing it from getting stuck in semantic loops (retrieving highly similar but non-progressive information) and allowing it to navigate much deeper reasoning paths without the exponential growth of the retrieved set.

4. Logical-Distractor Contrastive Training

• Improvement: Augment the Contrastive Learning Graph training pipeline with logical hard negatives—document chunks that are semantically highly similar to the query but contain logically incorrect or non-bridgeable information for the next hop.

• Capability: The AI system can achieve much higher precision in highly technical domains (legal, medical, or scientific research) by learning to distinguish between semantically related text and logically necessary evidence, significantly reducing retrieval-based hallucinations.

Retrieval-augmented generation (RAG) systems face significant challenges in multi-hop question answering (MHQA), where complex queries require synthesizing information across multiple document chunks. Existing approaches typically rely on iterative LLM-based query rewriting and routing, resulting in high computational costs due to repeated LLM invocations and multi-stage processes. To address these limitations, we propose TreeHop, an embedding-level framework without the need for LLMs in query refinement. TreeHop dynamically updates query embeddings by fusing semantic information from prior queries and retrieved documents, enabling iterative retrieval through embedding-space operations alone. This method replaces the traditional "Retrieve-Rewrite-Vectorize-Retrieve" cycle with a streamlined "Retrieve-Embed-Retrieve" loop, significantly reducing computational overhead. Moreover, a rule-based stopping criterion is introduced to further prune redundant retrievals, balancing efficiency and recall rate. Experimental results show that TreeHop rivals advanced RAG methods across four open-domain MHQA datasets, achieving comparable performance with only 2.2%-29.4% of the parameter size of concurrent solutions and reducing the query latency by 92.8%-97.8%. This makes TreeHop a faster and more cost-effective solution for low-resource or latency-sensitive deployment. For reproducibility purposes, codes and data are available here: https://github.com/allen-li1231/TreeHop-RAG