Search papers, labs, and topics across Lattice.
This paper introduces CodeHID, a novel generative code retrieval framework that transitions from flat candidate matching to a hierarchical, semantic address generation approach. By employing Pseudo-Neighbor Guided DocID Learning and Dual-Phase DocID Generation Guidance, CodeHID constructs a hierarchical index that enhances the retrieval of semantically related code snippets while maintaining their distinctiveness. Experimental results on CoSQA and ProCQA benchmarks show that CodeHID significantly outperforms traditional retrieval methods, achieving notable improvements in rank-one retrieval metrics.
CodeHID achieves a paradigm shift in code retrieval by organizing semantically similar snippets into a meaningful hierarchical index, leading to superior retrieval performance.
Code retrieval models have predominantly relied on a flat matching paradigm that treats code snippets as independent candidates, making them less capable of distinguishing similar code candidates. Generative retrieval offers a solution by constructing a learnable index over the code corpus, guiding the retriever to better understand how code candidates are semantically organized and addressed. However, naively applying generative retrieval in the code retrieval task may result in operating over an identifier space whose prefixes do not correspond to meaningful code-semantic regions. In this paper, we propose CodeHID, a generative code retrieval framework that reformulates the code retrieval task from flat candidate matching to coarse-to-fine semantic address generation. CodeHID relies on two core components. First, Pseudo-Neighbor Guided DocID Learning constructs a globally static hierarchical index by applying multi-level residual quantization and $k$-nearest-neighbor pseudo-supervision, ensuring that semantically related code snippets share prefixes while preserving target-level separability. Second, Dual-Phase DocID Generation Guidance reliably navigates this fixed index by combining training-side ranking enhancements, using hard negatives and rank distillation, with inference-side candidate constraints and prefix-aware decoding. Extensive experiments on CoSQA and ProCQA benchmarks demonstrate that CodeHID outperforms existing sparse retrieval, pre-trained code models, dense code retrieval, and generative retrieval baselines by a large margin in most cases, achieving particularly strong improvements in rank-one retrieval metrics.
