Hierarchy-Aware Training for Generative Retrieval

Author

Gaurav Sinha

Gaurav Sinha

Gaurav Sinha is a Principal Researcher at Microsoft Research, working in the areas of Information Retrieval, Reinforcement Learning, Causal Inference and Theory. He received his Ph.D. in Mathematics from the California Institute of Technology in 2016, where he was advised by Prof. Eric Rains and his Integrated M.Sc. in Mathematics and Scientific Computing from Indian Institute of Technology (IIT) Kanpur in 2011.

Project

Generative retrieval is a search method where a language model directly generates a document’s unique identifier instead of ranking candidates. To produce valid IDs, these models navigate a search hierarchy, such as a prefix tree. Standard training objectives treat these IDs as flat text sequences, ignoring the hierarchical relationship between tokens thereby creating a structural mismatch. We wish to propose a training objective that explicitly incorporates this tree structure into the loss function. Instead of a uniform penalty, the error signal is weighed by the token’s specific position and depth within the hierarchy. Our hope is that by aligning the objective with these discrete constraints, the model can traverse a more accurate trajectory through the search hierarchy.