想象一下，你是一名软件工程师，刚刚加入一个首次接触的庞大遗留代码库。你收到一张工单: “修复仪表盘上用户登录超时的错误。”你的第一个挑战不是修复代码，而是从成千上万个文件和数万个函数中找到该代码究竟在哪里 。

这种“大海捞针”的问题正是当今 AI 编码代理 (Coding Agents) 所面临的困境。虽然大型语言模型 (LLMs) 在编写代码方面表现出色，但它们在检索方面——即根据自然语言请求定位需要编辑的具体文件和函数——却举步维艰。

在这篇文章中，我们将深入探讨 CoRet , 这是一篇提出专门为代码编辑设计的稠密检索模型的新研究论文。通过利用仓库结构、调用图依赖关系以及专门的训练目标，CoRet 在性能上实现了对现有最先进模型的巨大飞跃。

标准检索器的问题

要理解为什么 CoRet 是必要的，我们首先需要看看现代代码检索通常是如何工作的。大多数系统使用“稠密检索” (Dense Retrieval) 方法:

  1. 编码器模型 (如 BERT 或 CodeBERT) 将用户查询 (例如，“修复登录超时”) 转换为向量。
  2. 同一个模型将代码片段也转换为向量。
  3. 系统计算查询向量和代码向量之间的相似度 (点积) ，以找到最佳匹配。

问题出在哪里？大多数预训练模型将代码仅仅视为扁平的文本。它们忽略了软件中固有的丰富结构信息:

  • 仓库层级结构: 文件存在于提供上下文的文件夹中 (auth/login.py 与 tests/login_test.py 截然不同) 。
  • 依赖关系: 函数会调用其他函数。要理解 process_login()，需要了解它调用的 validate_credentials()。

研究人员发现，现有的模型，即使是那些在代码上训练过的模型 (如 CodeSage 或 CodeBERT) ，在执行针对特定错误修复的仓库级检索任务时表现不佳。它们通常会检索出看起来语义相似，但并非实际需要修改逻辑的代码位置。

CoRet 解决方案: 代码库剖析

CoRet (Contextual Retriever，上下文检索器) 通过重新思考我们如何表示代码以及如何训练模型来寻找代码，从而解决了这些局限性。

1. 原子切分与仓库结构

一个标准的仓库文件可能有数千行长，包含多个类和数十个函数。检索整个文件太粗糙了；而检索任意的 512 个 token 的块通常会将逻辑切成两半。

CoRet 采用了一种“语义切分”策略。它不是基于任意行数，而是将代码解析为其“原子”: 函数、类和方法。

上图: 处理前的代码仓库。下图: 过滤和切分后的代码块。

如图 1 所示，仓库被分层分解。注意图片下半部分的两个关键细节:

  1. 粒度: 文件 jedi.py 被切分为独特的逻辑单元 (r2d2 函数，Jedi 类) 。
  2. 上下文保留: 至关重要的是, 文件路径被前置到了代码中。字符串 knights/jedi.py 成为了嵌入的一部分。

作者发现包含文件路径是检索的一个巨大信号。在他们的实验中，移除文件路径导致性能显著下降 (我们要结果部分会看到这一点) 。模型有效地学会了“阅读”目录结构以缩小搜索范围。

2. 调用图的力量

代码很少是孤立的。一个函数的意义通常由它调用什么 (下游依赖) 或什么调用它 (上游使用) 来定义。

CoRet 通过包含调用图上下文来丰富代码块的表示。具体来说，作者关注的是下游邻居——即当前块调用的函数。

给定一个函数 c_i 及其下游邻居 c_out，我们将字符串拼接，包括特殊的分隔符，并微调模型以获得 CoRet。

如图 2 所示，如果函数 lightsaber() (目标块 \(c_i\)) 调用了 lightsaber_on() (邻居 \(c_{out}\)) ，模型不仅能看到 lightsaber() 的代码。它看到的是一个拼接的字符串:

lightsaber code + [DOWN] + lightsaber_on code

这使得嵌入不仅能捕捉函数是什么，还能捕捉它在系统其余部分中做什么。

数学核心: 更好的损失函数

这篇论文最重要的贡献可能在于训练目标。标准的对比学习 (用于训练大多数检索器) 通常使用“批次内负样本” (in-batch negatives) 。这意味着如果模型正在查看来自仓库 A 的查询，它会将 (同一训练批次中) 来自仓库 B 的代码视为“错误的答案”。

这太简单了。区分登录函数和天气应用的数据库函数是很容易的。

代码检索的难点在于从同一个仓库中，区分实际的登录函数与登录辅助函数、登录测试函数或登出函数。

CoRet 使用了一种专门为这种“实例内” (in-instance) 检索设计的损失函数。

公式 1: 优化检索正确块可能性的损失函数。

在这个公式中:

  • \(\mathbf{q}_i\) 是查询嵌入。
  • \(\mathbf{c}^*\) 是正确的代码块 (正样本) 。
  • \(\Gamma\) 是归一化因子 (与所有其他块的相似度之和) 。

这看起来像是一个标准的交叉熵损失，但魔力在于他们如何计算分母 \(\Gamma\)。由于对整个仓库 (10,000+ 个块) 进行求和在计算上非常昂贵，他们使用从同一仓库采样的困难负样本集来近似它。

公式 2: 使用实例内负样本近似归一化因子。

这里，\(\mathcal{B}\) 代表“实例内负样本”的随机样本。通过强制模型区分正确块与同一项目中的其他块，模型学会了细粒度的区分，而不仅仅是通用的话题匹配。

实验结果

作者在 SWE-bench 上评估了 CoRet，这是一个基于真实 GitHub 问题 (错误和功能请求) 的基准测试。他们将 CoRet 与强大的基线 (如 BM25 关键字搜索) 和预训练模型 (如 CodeBERT 和 CodeSage) 进行了比较。

使用的主要指标是 Recall@k : 如果模型检索 \(k\) 个块，正确的那个是否在列表中？

召回率指标定义

巨大的性能提升

结果令人震惊。CoRet 不仅仅是略微胜出，而是占据了主导地位。

显示 CoRet 在 Recall@k 上显著优于基线的图表。

在图 3 中，请看棕色虚线 (CoRet) 。它急剧上升，在仅 \(k=5\) 时就达到了超过 50% 的召回率 , 而之前最好的模型 (CodeSage S，红线) 仅在 35% 左右。

这对 AI 代理来说是一个关键的改进。代理的上下文窗口有限；它不能阅读 100 个文件。它需要在前 5 或 10 个结果中获得正确的代码。

作者还引入了一个严格的指标，称为 Perfect Recall (完美召回率) , 用于衡量模型检索到修复错误所需所有块的频率 (因为错误通常跨越多个函数) 。

表 1: 块级别的完美召回率。

表 1 强调了 CoRet (最后一行) 实现了 0.54 的 Perfect Recall@5 , 而基础 CodeSage 模型仅为 0.34。这意味着在找到解决问题所需的所有代码片段的能力上，提高了 20 个百分点。

为什么它有效？消融研究

作者系统地移除了某些特征，以观察什么最重要。

1. 文件路径的重要性: 当从输入中移除文件路径时，性能显著下降。

表 2: 比较有无文件路径的 CoRet。

表 2 显示，如果没有文件路径，Recall@5 从 0.53 下降到 0.42。这证实了目录结构提供了必要的语义线索。

2. “实例内”负样本的重要性: 最有趣的发现之一是训练期间负样本采样策略的影响。

图 5: 负样本数量及其来源的影响。

图 5 显示，使用“实例内”负样本 (紫色/绿色线) 始终优于“跨实例”负样本 (橙色线) 。此外，简单地增加这些负样本的数量 (从 8 增加到 1024) 会推高性能。训练任务越难 (在同一仓库的更多块之间进行区分) ，检索器就变得越聪明。

可视化注意力

为了证明模型实际上是在使用文件路径，而不仅仅是运气好，作者可视化了 Transformer 模型的注意力图 (Attention Maps) 。

注意力图显示模型聚焦于文件路径 token。

  • (图 6: 包含文件路径的查询的注意力图) *

注意力图: 包含文件路径的代码块。

  • (图 7: 代码块的注意力图) *

这些热力图 (图 6 和图 7) 显示了代表文件路径的 token 周围有强烈的激活 (较浅的颜色) 。模型有效地学会了在问题描述中提到或暗示的文件路径与仓库中的实际文件之间执行“模糊匹配”。

结论与启示

CoRet 代表了“仓库级”理解向前迈出的重要一步。它使我们不再将代码视为词袋，而是将其视为结构化、互连的图。

主要收获:

  1. 结构即信号: 包含文件路径和调用图邻居彻底改变了检索性能。
  2. 训练上下文很重要: 对于此任务，训练模型区分同一仓库中的文件比标准的对比学习方法有效得多。
  3. 代理需要帮助: 为了让 AI 软件工程师发挥作用，他们需要可靠的检索。CoRet 提供了让这些代理高效定位错误的“眼睛”。

随着 AI 编码助手从简单的自动完成演变为能够导航和修复复杂仓库的自主代理，像 CoRet 这样的专用检索系统将成为其成功的基石。毕竟，在大海中捞针的能力，是穿针引线的第一步。