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
用残差量化和k近邻伪标签建成可寻址的层次代码索引,再约束解码生成DocID,CoSQA上Hit@1达到0.744,比UniXcoder的0.518高出二十多个点。
代码检索长期把每段代码当成独立候选,用自然语言和代码的向量相似度一次排完。这对表面相近、结构不对的片段很吃亏:函数名和注释对上了,控制流和API用法却不是查询要的。结构感知编码器把数据流、标识符、AST塞进表示,检索空间本身仍是扁平的一堆片段,候选之间的远近没有被显式编进索引。
生成式检索(DSI这一路)改成直接生成文档标识符DocID,索引变成可学习的地址空间。直接套到代码上会出另一种问题:标识符前缀若不对应功能或实现上的语义区域,解码仍会跟着表面词走。
CodeHID把任务改写成从粗到细的语义地址生成:浅层前缀圈功能区,深层token再分开实现细节。
两块,先建地址,再学会走路。
第一块是伪邻居引导的DocID学习。用RQ-VAE把代码表示残差量化成四级离散码,码本大小256。同时在预训练编码器的连续空间上建k近邻图(k=8),按层级相似度阈值做成前缀共享伪标签:邻居在该层够像,就鼓励共享到这一层的前缀;非邻居则压深前缀重叠。正负约束用二元交叉熵对齐软码字分配算出的前缀共享概率。量化重建负责可编码,邻居约束负责让前缀有检索语义。
第二块是双阶段生成引导。训练时从合法DocID里挖对当前查询打分高的硬负例(K=4),用成对偏好让正例高于它们,再用教师的软排序(查询与代码的余弦)对学生的生成似然做秩蒸馏。推理时先按同一余弦从语料取出top-3000个DocID,建成查询相关的前缀trie,beam size 20的约束解码只走合法路径,并在每步加上前缀级词法加分(β=0.10)。编码器是冻结的GraphCodeBERT,解码器是Qwen2.5-Coder-7B加LoRA。
主结果在CoSQA和清洗后的ProCQA(Python/Java各约3万条,8:2:2划分)。
| 方法 | 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 |
CoSQA上Hit@1从最强匹配模型UniXcoder的0.518拉到0.744,MRR@20从0.654到0.766。Hit@5是另一面:UniXcoder 0.850,CodeHID 0.792,覆盖率没赢。ProCQA上四个指标都是第一。换成Qwen2.5-Coder-1.5B,CoSQA Hit@1仍有0.565,高于1.3B–2B的密集检索基线。
消融(CoSQA Hit@1):去掉DocID学习掉到0.399,去掉训练期查询信号掉到0.355,去掉推理期查询信号只掉到0.728。秩蒸馏和词法一致性各自拿掉后,Hit@1落到0.361和0.351。基线把目标召回进top-5却排不到第一的查询,CodeHID有76.3%能排到第一。
代码库里功能相近、实现不同的片段很多,扁平匹配最容易在第一名上栽。这篇说明生成式检索要先把标识符空间按代码邻域组织好,再在合法路径里做查询条件的排序,单靠换一个更大的生成器不够。1.5B变体仍强于同规模密集模型,增益不完全是7B解码器买来的。
工程上这是混合系统:推理仍要一次余弦短名单,再在trie里生成。DocID是全局静态的四级码。适合语料相对稳定、更在意第一名是否对的代码搜索,不适合每天都在涨的仓库索引。
推理侧B=3000的余弦约束说明它不是从全空间凭空生成地址,密集检索短名单仍在。CoSQA的Hit@5低于UniXcoder,第一名变准了,召回没变宽。ProCQA经过去重、规范化AST指纹和抽样,不是完整StackOverflow规模。静态DocID对持续演化的代码库不友好,作者自己把持续学习写成未来工作。k近邻伪标签继承底层编码器的邻域误差,码本和阈值是调过的(k=8,λ=0.03)。