TreeHop: Efficient Embedding-Level Query Rewriter
TreeHop is a lightweight, embedding-level query rewriter for multi-hop retrieval-augmented generation that updates and prunes query embeddings directly, instead of repeatedly calling an LLM to rewrite the query between retrieval hops.
Multi-hop questions often need several rounds of retrieval, and the conventional approach rewrites the query with an LLM between each round, which is slow and costly; TreeHop instead updates the query embedding directly and prunes unpromising branches, producing a streamlined retrieve-embed-retrieve loop. The project reports 92.8 to 97.8 percent lower retrieval latency than iterative LLM-based rewriters, using a model of only 0.18 billion parameters compared to the billions of parameters typical rewriters need, while maintaining comparable recall.
The repository includes evaluation embedding databases for reproducing published results, code for training a TreeHop model from scratch, and instructions for skipping the large data download when only the codebase itself is needed.
It requires a GPU or Apple Metal with at least 32GB of RAM, is written in Python, and includes citation information for its accompanying paper.