CodeHID: Learning an Addressable Hierarchical Code Index for Generative Code Retrieval
Zhen Li, Yuhong Chen, Wenhao Xu, Xiaodong Li, Hui Li
cs.IR
2026-08-25
CodeHID builds a hierarchical code index with residual quantization and k-NN prefixes, then generates DocIDs; CoSQA Hit@1 reaches 0.744 versus UniXcoder's 0.518.
Code retrieval still scores each snippet independently against a natural-language query. Surface overlap then beats structure: similar names and comments outrank the snippet whose control flow and APIs actually match the intent. Structure-aware encoders inject data flow, identifiers, or ASTs into the vector, but the candidate set stays flat. Relations among snippets never enter the index.
Generative retrieval (DSI and follow-ons) emits document identifiers instead of scores, so the identifier space itself is learnable. Applied naively to code, prefixes often fail to line up with functional or implementation neighborhoods, and decoding still tracks surface form.
CodeHID recasts the task as coarse-to-fine semantic addressing: shallow prefixes locate a functional region, deeper tokens separate implementations.
Two pieces: build the address space, then learn to walk it.
Pseudo-Neighbor Guided DocID Learning quantizes code vectors with a four-level RQ-VAE, codebook size 256. A k-NN graph (k=8) in a pretrained encoder's continuous space supplies prefix-sharing pseudo-labels at increasing similarity thresholds: neighbors similar enough at a level should share that prefix; non-neighbors should not share deep prefixes. Binary cross-entropy aligns those labels with prefix-sharing probabilities from soft codeword assignments. Reconstruction keeps the codes writable; the neighbor term gives prefixes retrieval meaning.
Dual-Phase DocID Generation Guidance handles navigation. At training time the model mines K=4 high-scoring legal negatives for each query, pushes the positive above them with pairwise preference, and distills the teacher's cosine ranking into the generator's likelihoods. At inference it first takes the top-3000 DocIDs by the same cosine score, builds a query-specific prefix trie, and runs constrained beam search (beam 20) with a prefix lexical bonus (β=0.10). The encoder is frozen GraphCodeBERT; the decoder is Qwen2.5-Coder-7B with LoRA.
Benchmarks are CoSQA and a cleaned ProCQA split (about 30k Python and 30k Java, 8:2:2).
| Method | CoSQA Hit@1 | ProCQA-Py Hit@1 | ProCQA-Java Hit@1 |
| CodeHID (7B) | 0.744 | 0.597 | 0.520 |
| UniXcoder | 0.518 | 0.362 | 0.363 |
| CodeXEmbed | 0.470 | 0.208 | 0.216 |
| NCI | 0.498 | 0.406 | 0.325 |
| BM25 | 0.252 | 0.316 | 0.360 |
On CoSQA, Hit@1 moves from UniXcoder's 0.518 to 0.744, MRR@20 from 0.654 to 0.766. Hit@5 goes the other way: UniXcoder 0.850, CodeHID 0.792. On ProCQA, CodeHID leads every metric. A 1.5B decoder variant still reaches 0.565 Hit@1 on CoSQA, above 1.3B–2B dense baselines.
Ablations on CoSQA Hit@1: drop DocID learning and land at 0.399; drop train-time query signals, 0.355; drop inference-time query signals, only 0.728. Removing rank distillation or lexical consistency collapses Hit@1 to 0.361 and 0.351. On queries where a baseline recalls the target in the top-5 but not at rank one, CodeHID places it first 76.3% of the time.
Code corpora are full of near-functional duplicates. Flat matching fails most often at rank one. The result here is that generative retrieval needs a neighborhood-organized identifier space plus query-conditioned ranking among legal paths; a larger generator is not enough. The 1.5B variant beating same-scale dense models suggests the gain is not just the 7B decoder.
In production this is a hybrid: a cosine shortlist, then trie-constrained generation, over a globally static four-level code. It fits a relatively stable corpus where first-hit accuracy matters more than recall at 5. It does not fit a repo index that changes daily.
The B=3000 cosine constraint means the model is not inventing addresses over the full corpus; dense shortlisting is still in the loop. CoSQA Hit@5 trails UniXcoder: rank-one got sharper, coverage did not. ProCQA is deduplicated, AST-normalized, and subsampled, not full StackOverflow scale. Static DocIDs sit poorly with evolving codebases; the authors flag continual learning as future work. k-NN pseudo-labels inherit encoder neighborhood errors, and k=8, λ=0.03 were tuned.