跳至内容

用在线强化学习改进 Cursor Tab

作者: Jacob Jackson, Phillip Kravtsov & Shomil Jain归入研究

在 Cursor,我们的目标是让开发者的效率提升一个数量级。实现这一目标的重要组成部分是 Cursor Tab——我们的系统会在你的代码库中预测你的下一步操作。每当你输入一个字符或在编辑器中移动插入符号时,我们的 Tab 模型都会尝试预测你的后续操作;如果其置信度足够高,我们会将该预测显示为建议,你可以按 Tab 键接受。

我们的 Tab 模型在每一次用户操作时都会运行,每天处理超过 4 亿次请求。因此,我们积累了大量关于用户接受或拒绝哪些建议的数据。本文将介绍我们如何利用这些数据,通过在线强化学习改进 Tab。

我们的做法不同寻常,因为我们会在一天内频繁向用户推出新模型,并使用这些数据进行训练。大多数其他 LLM 提供商要么在静态数据集上训练,要么使用付费标注员,并且仅在每隔数个月进行一次具名模型发布时,才向用户推出新模型。

噪声建议的问题

我们致力于保持 Tab 建议的接受率较高。如果接受率较低,说明我们展示了过多不正确的建议,这会分散注意力并打断编码的流畅性。

要实现高接受率,不仅要让模型更智能,还要知道何时该提出建议、何时不该提出建议。有时信息不足,无法判断用户将采取什么操作:即使模型拥有完美的知识和推理能力,也无法预知用户会怎么做。在这种情况下,我们不应给出任何建议。

为了提高模型建议的接受率,一个简单的方法是训练一个单独的模型来预测该建议是否会被接受。2022 年,Parth Thakkar发现 GitHub Copilot 采用了这种方法,使用逻辑回归模型基于 11 个特征推导出“上下文过滤分数”,这些特征包括编程语言、上一条建议是否被接受或拒绝、用户插入点前的尾随字符,以及其他特征。目前尚不清楚该模型训练预测的具体信号,但我们最好的猜测是:它预测在显示建议时用户接受该建议的可能性。当该分数低于 15% 时,建议会被跳过,不会显示任何内容。

这个解决方案可行,但我们希望采用一种更通用的机制,复用 Tab 模型对代码所学得的强大表示。我们并不想事后过滤掉糟糕的建议,而是希望从源头上调整 Tab 模型,避免生成不佳的建议。因此,我们转而采用了策略梯度方法。

策略梯度

策略梯度方法是一类通用手段,用于优化“策略”(此处指 Tab 模型)以提升“奖励”。奖励是我们为策略所采取的每个动作分配的一个数值。通过使用策略梯度算法,我们可以更新策略,使其在未来获得更高的平均奖励。

这些算法的工作方式是:让策略以一定随机性进行行为,观察哪些动作带来高或低回报,然后对带来高回报的动作进行正向强化,对带来低回报的动作进行负向强化。

为了用策略梯度方法改进 Tab,我们定义了一个奖励函数:鼓励被接受的建议,同时惩罚向用户展示但未被接受的建议。假设我们希望当建议被接受的概率至少为 25% 时,模型才展示该建议。那么,我们可以为被接受的建议赋予 0.75 的奖励,为被拒绝的建议赋予 -0.25 的奖励,而在不展示建议时奖励为 0。若接受概率为p,则在展示建议时的期望奖励为0.75p0.25(1p)0.75 \cdot p - 0.25 \cdot (1 - p),该值为正当且仅当p > 0.25。因此,一个以最大化奖励为目标的策略会在其估计接受概率至少为 25% 时给出建议,否则不展示任何内容。

在实际应用中,我们使用更复杂的回报函数,它不仅考虑建议的长度/规模,还考虑跳转到代码中其他位置并展示更多建议的可能性。但这已经说明了核心思路:我们不显式建模接受率,而是学习一个以特定接受率为目标的策略。可以推测,模型会在其内部表示中学到关于接受概率的模型(或者至少学到它是否超过 25%),但我们将其交由优化器处理。

On-policy 数据的重要性

为了获得策略更新,我们依赖一个被称为“策略梯度定理”的重要事实:如果一个策略π(as,θ)\pi(a \mid s, \theta)在状态(例如用户代码库的状态)α\alpha上对动作给出了一个分布sP(s)s \sim P(s),并由参数θ\theta参数化,且奖励为J(θ)=EsP(s),aπ(as,θ)[R(s,a)]J(\theta) = \mathbf{E}_{s \sim P(s),\, a \sim \pi(a \mid s,\theta)}\bigl[R(s,a)\bigr],那么奖励的梯度为:

θJ(θ)=EsP(s),aπ(as,θ)[θlogπ(as,θ)R(s,a)]\nabla_\theta \, J(\theta)= \mathbf{E}_{s\sim P(s),a\sim \pi(a \mid s,\theta)}\left[\nabla_\theta \, \log \, \pi(a \mid s,\theta) \cdot R(s,a)\right]

这很有用,因为右侧易于估计:我们可以通过对用户请求展示的建议进行采样来获取状态与动作的样本sP(s),  aπ(as,θ)s \sim P(s), \; a \sim \pi(a \mid s, \theta),我们可以使用类似 PyTorch 的框架来计算θlogπ(as,θ)\nabla_{\theta} \,\log \,\pi(a \mid s, \theta),并且我们可以通过查看用户是否接受了该建议来计算R(s,a)R(s,a)。因此我们可以使用该等式获得θJ(θ)\nabla_{\theta} J(\theta)的无偏估计,从而通过随机梯度下降改进策略。

然而,这只有在从正在优化的策略中采样动作时才有效。一旦我们更新了策略,就不再拥有来自当前优化策略的样本——我们手上只有来自先前策略的样本。为了获得新的“on-policy(同策略)”样本,我们需要将新模型部署给用户并观察其行为。这意味着我们需要可靠的基础设施,能够快速部署新的检查点,并尽量缩短从向用户展示建议到该数据进入我们训练流程下一步之间的时间。目前,我们需要大约 1.5 到 2 小时来推送一个检查点并收集用于下一步的数据。虽然这在 AI 行业相对算快,但仍有提升空间,使其更快。

全新的 Tab 模型

使用这里介绍的方法,我们训练了一个新的 Tab 模型,现已成为 Cursor 的默认模型。相比上一代模型,该模型的建议数量减少了 21%,但其给出的建议接受率提升了 28%。我们希望这能提升你的编码体验,并计划在未来进一步完善这些方法。

Graph showing the percentage improvement of the new Tab model

归类于: 研究

作者s: Jacob Jackson, Phillip Kravtsov & Shomil Jain