Hierarchical Embedding Fusion for Retrieval-Augmented Code Generation

本文提出了分层嵌入融合（HEF）方法，通过离线构建可复用的分层向量缓存和在线将检索向量映射为固定数量的伪令牌，在保持仓库级代码生成精度的同时，显著降低了检索增强代码生成的推理延迟。

要点

这篇论文介绍了一种名为 HEF（分层嵌入融合） 的新方法，旨在解决一个非常具体的问题：如何让 AI 写代码时，既能参考整个项目（成千上万行代码），又不会反应太慢或“脑子”被塞满。

为了让你轻松理解，我们可以把 AI 写代码想象成一位才华横溢但记性有限的“建筑师”。

1. 痛点：建筑师面临的困境

想象一下，你是一位建筑师（AI 模型），正在设计一栋大楼的一层（写一段新代码）。

• 传统做法（单文件模式）： 你只看手头这一层的图纸。结果，你忘了隔壁楼层的承重墙在哪，或者不知道隔壁房间用的什么风格的窗户。你造出来的房子虽然结构没问题，但跟整个大楼不协调，甚至可能把承重墙给拆了（代码报错或逻辑冲突）。
• 现有的“检索增强”做法（RAG）： 为了帮你，有人把整栋大楼的所有图纸（整个项目的代码库）都复印了一份，直接塞到你面前。
  • 问题： 图纸太多了！你还没开始画图，光是翻阅这几千页的图纸就花光了时间（延迟高）。而且，图纸里混杂了很多无关紧要的垃圾信息（比如旧版本的草稿），把你搞晕了（噪声干扰）。
  • 更高级的“图检索”做法： 有人帮你画了一张复杂的“大楼关系图”，告诉你哪根线连着哪。但这需要现场重新计算和遍历，速度依然很慢。

2. 解决方案：HEF（分层嵌入融合）

HEF 的核心思想是：不要直接把几千页图纸塞给建筑师，而是先派一个“超级助理”把图纸浓缩成一张“精华记忆卡”。

这个过程分为两步：

第一步：离线准备（建图书馆）—— 把书读薄

在建筑师开始工作之前，我们有一个“离线阶段”。

• 切片与编码： 先把整个项目的代码切成小块（像把书撕成章节）。
• 分层融合（核心魔法）： 这里有一个叫“融合器（Fuser）”的小模型（就像一位超级图书管理员）。
  • 它把同一文件的代码块合并成“文件摘要”。
  • 把同一文件夹的代码合并成“模块摘要”。
  • 最后把整个项目合并成“项目总览”。
  • 比喻： 就像把一本 1000 页的厚书，压缩成了 32 张**“记忆卡片”。每张卡片不是文字，而是一种“数学味道”**（向量），代表了那部分代码的核心含义。
• 结果： 整个项目的信息被压缩成了一个分层的、紧凑的数据库，存好了随时待命。

第二步：在线工作（写代码）—— 只带卡片进场

当建筑师（AI 生成器）需要写代码时：

1. 提问： 它看着当前的代码，问：“我需要参考什么？”
2. 检索： 系统瞬间从刚才建好的“记忆卡片库”里，挑出最相关的 32 张卡片（而不是几千页原文）。
3. 注入： 这 32 张卡片被转换成一种特殊的“虚拟令牌”（Pseudo-tokens），直接贴在建筑师的指令前面。
4. 生成： 建筑师看着这 32 个“精华提示”，瞬间明白了整个项目的背景，开始写代码。

比喻： 以前是让你读整本百科全书再答题；现在是让你看一眼32 个精心提炼的“关键词提示”，然后答题。速度极快，而且信息量足够。

3. 为什么这个方法很厉害？

• 速度极快（低延迟）：
  • 以前的方法（如 DRACO、GraphCoder）像是要现场重新画一遍大楼的结构图，需要 10-17 秒。
  • HEF 只需要 0.68 秒（比眨眼还快）。因为它不需要重新分析代码结构，直接调用存好的“记忆卡片”。
• 效果一样好（高精度）：
  • 虽然只看了 32 张卡片，但 AI 写代码的准确率（Exact-Match）竟然和那些看了几千页原文的方法差不多，甚至更好。
  • 这说明：“浓缩的才是精华”。AI 不需要看到所有文字，只需要理解核心逻辑。
• 抗干扰能力强：
  • 如果检索到的卡片里有垃圾信息（比如错误的代码片段），HEF 的机制能很好地过滤掉，不会像直接塞原文那样把 AI 搞糊涂。

4. 总结：这就好比……

想象你要去一个陌生的城市旅行：

• 旧方法： 给你一本该城市所有街道的详细地图集（几千页），让你出发前背下来。你背得累死，出发时还晕头转向。
• HEF 方法： 给你一张智能导航卡。出发前，系统已经帮你把城市的所有重要地标、路线逻辑都压缩成了几个关键坐标。你只需要看这几个坐标，就能瞬间知道该往哪走，既快又准。

一句话总结：
这篇论文发明了一种**“代码压缩术”，把庞大的项目代码库变成了一组“魔法记忆卡片”。这让 AI 写代码时，既能拥有“全知全能”的项目视野，又能保持“秒回”的超快反应速度，完美平衡了质量与速度**。

技术摘要

这是一篇关于**分层嵌入融合（Hierarchical Embedding Fusion, HEF）**的论文技术总结，旨在解决检索增强代码生成（RACC）中的延迟与噪声问题。

1. 研究背景与问题 (Problem)

在仓库级代码补全（Repository-level Code Completion）任务中，模型不仅需要理解当前文件的上下文，还需要访问跨文件的信息（如导入的类、函数签名、类型定义等）。现有的检索增强生成（RAG）方法主要存在以下痛点：

• 延迟与成本耦合：传统的“片段注入”（Snippet-injection）方法将检索到的原始代码片段直接拼接到提示词（Prompt）中。这导致在线推理成本与检索到的 Token 数量成正比，当检索大量上下文时，推理延迟显著增加。
• 上下文噪声：将大量原始文本注入上下文窗口会引入无关片段，干扰模型生成，降低准确性。
• 现有替代方案的局限：
  • 基于图或迭代检索的方法（如 DRACO, GraphCoder）虽然提高了相关性，但需要在线进行图遍历或多次模型调用，延迟极高。
  • 长上下文模型（如 RepoFusion）虽然能处理全库，但推理成本随仓库大小线性增长，且无法动态选择最相关的信息。

2. 方法论 (Methodology)

HEF 提出了一种两阶段的仓库表示方法，将仓库信息压缩为可重用的分层稠密向量，并通过**伪 Token（Pseudo-tokens）**接口与生成模型交互。

A. 离线阶段：分层缓存构建 (Offline Stage)

1. 分块与嵌入：将仓库文件切分为 $\le 512$ Token 的语义块（Chunk），使用冻结的强编码器（Embedder，如 Qwen3-Embedding-8B）将其映射为稠密向量。
2. 分层融合 (Fusion)：引入一个小型的“融合器”（Fuser，如 Qwen-2.5-Coder-0.5B）。
   • 该模型递归地将子节点向量（如代码块）合并为父节点向量（如文件、模块、仓库级表示）。
   • 构建出一个树状分层结构：Chunk -> File -> Module -> Repository。
   • 每个节点存储元数据，可回溯到原始代码范围。
3. 索引：所有层级的节点向量被存入 HNSW 索引中，供在线检索使用。

B. 在线阶段：查询与生成 (Online Stage)

1. 查询构建：将当前代码前缀（Prefix）的后 512 Token 编码为查询向量。
2. 检索：在分层索引中检索 Top-K（例如 K=32）个最相关的节点（可能来自不同层级）。
3. 伪 Token 投影：
   • 使用一个投影器（Projector，MLP）将检索到的 $K$ 个稠密向量映射为 $K$ 个伪 Token（Pseudo-tokens）。
   • 这些伪 Token 作为软提示（Soft Prompt）拼接到生成模型的输入序列前端。
4. 生成：生成器（Generator）基于前缀和伪 Token 进行代码补全。
   • 核心优势：无论检索了多少原始代码，输入给生成器的 Token 数量是固定的（例如 32 个），从而解耦了仓库大小与推理延迟。

C. 训练策略

• 对比预训练：使用 InfoNCE 损失函数训练 Fuser，使其能压缩出具有区分度的仓库表示。
• 联合端到端优化：解冻 Fuser、Projector 和 Generator，联合优化语言模型损失（Causal LM Loss），使分层向量与生成概率更好地对齐。
• 数据过滤 (UWL)：引入“效用加权似然”（Utility-Weighted Likelihood）信号，仅保留那些能提升生成似然度的上下文进行训练，过滤有害或无关的检索内容。

3. 关键贡献 (Key Contributions)

1. 分层稠密缓存 + 伪 Token 接口：提出了一种将仓库大小与在线提示词长度解耦的架构，用固定数量的向量 Token 替代了成千上万的原始代码 Token。
2. 端到端集成方案：将强编码器、轻量级融合器、伪 Token 接口以及无监督的数据构建流程（UWL）整合为一个完整的仓库级代码补全流水线。
3. 训练机制分析：对比了“对比预训练 + 微调”与“端到端联合训练”两种模式，并证明了端到端训练结合 UWL 数据过滤能显著提升效果。
4. 鲁棒性验证：证明了该方法在面对有害或无关检索内容时，比传统 RAG 具有更强的抗噪能力。

4. 实验结果 (Results)

在 RepoBench 和 RepoEval 两个基准测试上，HEF 表现优异：

• 准确性：
  • 在 RepoBench 上，HEF（端到端版）达到了 61.3% 的精确匹配（Exact-Match）准确率。
  • 相比基线模型 Qwen-2.5-Coder-1.3B 提升了 12.2 个百分点，比低延迟基线 RepoFusion 提升了 21.5 个百分点。
  • 与高延迟的图检索模型（如 GraphCoder, 16B 参数）相比，HEF（1.8B 参数）仅低了 2.8 个百分点，但参数量仅为对方的 1/8。
• 延迟与效率：
  • 推理延迟：中位延迟为 0.68 秒（单 A100）。
  • 速度对比：比 GraphCoder 快 26 倍，比 DRACO 快 13 倍。
  • 离线成本：构建仓库缓存仅需 35 秒/项目，远低于迭代检索方法的开销。
• 消融实验：
  • 发现约 30-40 个 融合 Token 即可捕获大部分仓库级信息，超过 60 个会导致性能下降。
  • 使用 UWL 过滤数据能显著提升 CodeBLEU 分数（从 24.53 提升至 36.20）。
  • 即使使用较小的 Fuser（0.5B），也能获得大部分收益，无需更大的融合模型。

5. 意义与结论 (Significance)

• 打破延迟 - 准确性权衡：HEF 证明了通过分层稠密缓存和伪 Token 接口，可以在保持亚秒级低延迟的同时，实现接近甚至超越高延迟复杂检索系统的代码生成质量。
• 实用性强：该方法不依赖复杂的在线图构建或多次模型调用，适合对响应速度要求高的交互式代码补全场景。
• 未来方向：为低延迟、仓库感知的代码生成提供了一种新的范式，即“压缩表示”优于“原始文本注入”。未来的工作可探索自适应层级构建及符号结构与连续缓存的混合方法。

总结：HEF 通过离线构建分层向量索引，在线仅检索并投影少量“伪 Token"，成功解决了检索增强代码生成中“上下文过长导致延迟高、噪声大”的核心矛盾，在保持高精度的同时实现了极低的推理延迟。