Cursor では、開発者の生産性を桁違いに高めることを目指しています。その重要な要素が Cursor Tab です。これはコードベース全体で次の操作を予測するシステムです。エディタ内で文字を入力したり、キャレットを移動したりするたびに、Tab モデルが次に行う操作を予測し、十分な確信度がある場合は、その予測を提案として表示します。提案は Tab キーで受け入れられます。

私たちのTabモデルは、ユーザーのあらゆる操作で動作し、1日に4億件以上のリクエストを処理しています。その結果、ユーザーがどの提案を採用し、どれを却下したかに関する豊富なデータが蓄積されています。本記事では、このデータをオンライン強化学習で活用し、Tabをどのように改善しているかを説明します。

私たちの手法は、1日のうちに高頻度で新しいモデルをユーザーに順次展開し、そのデータを学習に活用するという点で独特です。多くの他の LLM プロバイダーは静的なデータセットでの学習や有償ラベラーの利用に依存しており、数か月ごとに名前付きモデルのリリースとして新モデルをユーザーに提供するだけです。

ノイズの多いサジェストの問題

私たちは、Tab の提案の受け入れ率を高く保つよう努めています。受け入れ率が低い場合は、不正確な提案を過剰に表示しており、注意をそらし、コーディングのフローを妨げていることを意味します。

高い受け入れ率を達成するには、モデルを賢くするだけでなく、いつ提案すべきか・すべきでないかを見極めることも重要です。ユーザーが次にどのような行動を取るかを判断するのに十分な情報がない場合があります。たとえモデルが完全な知識と推論能力を持っていたとしても、ユーザーが何をするかは分かりません。こうした状況では、何も提案すべきではありません。

モデルの提案の受け入れ率を高めるための簡便な方法のひとつは、その提案が受け入れられるかどうかを予測する別のモデルを学習させることです。2022年、Parth Thakkar はGitHub Copilot がこの手法を用いていることを発見しました。これはロジスティック回帰モデルで11個の特徴量（プログラミング言語、直前の提案が受け入れられたか却下されたか、ユーザーのキャレット直前の末尾文字列、その他の特徴）を入力として「コンテキストフィルター・スコア」を算出するというものです。このモデルが何というシグナルを予測するように学習されたかは不明ですが、最もあり得る推測は、提案が表示された場合にユーザーがそれを受け入れる確率を予測しているというものです。スコアが15%未満の場合、その提案はスキップされ、何も表示されません。

この解決策は実用的ではありますが、私たちはTabモデルが学習した強力なコード表現を再利用できる、より汎用的な仕組みを求めていました。望ましくないサジェストを後からフィルタリングするのではなく、そもそも悪いサジェストを生成しないようにTabモデル自体を調整したかったのです。そこで、私たちは代わりにポリシー勾配法の採用を選びました。

方策勾配

ポリシー勾配法は、「報酬」を高めるために「ポリシー」（この場合は Tab モデル）を最適化するための一般的な手法です。報酬とは、ポリシーが取る各アクションに割り当てる数値のことです。ポリシー勾配アルゴリズムを用いることで、将来の平均報酬がより高くなるようにポリシーを更新できます。

これらのアルゴリズムは、方策にランダムに振る舞わせ、どの行動が高い報酬や低い報酬につながるかを観察し、高報酬につながった行動を正に強化し、低報酬につながった行動を負に強化することで機能します。

Tab を改善するために方策勾配（policy gradient）手法を用いるにあたり、受け入れられた提案を促し、受け入れられない提案をユーザーに表示することを抑制する報酬を定義しました。たとえば、受け入れられる確率が少なくとも 25% の場合にモデルが提案を表示してほしいとします。このとき、受け入れられた提案には 0.75 の報酬、拒否された提案には -0.25 の報酬、何も表示しなかった場合には 0 の報酬を与えることができます。受け入れ確率がpのとき、提案を表示した場合の期待報酬は$0.75 \cdot p - 0.25 \cdot (1 - p)$となり、これはp > 0.25 のときにちょうど正になります。したがって、報酬を最大化する方策は、受け入れ確率が少なくとも 25% と見積もられるときに提案を表示し、それ以外は何も表示しません。

実際には、提案の長さ（規模）に加えて、コード内の他の箇所へジャンプして追加の提案を表示する可能性も考慮する、より複雑な報酬を用いています。しかし、ここで示したい中核は同じです。受理率を明示的にモデリングするのではなく、特定の受理率を目標とする方策（ポリシー）を学習します。おそらくモデルは内部表現において受理確率のモデル（少なくともそれが25％を超えるかどうかのモデル）を学習しますが、その点はオプティマイザに委ねます。

オンポリシー・データの重要性

ポリシー更新を得るために、私たちはPolicy Gradient Theorem（方策勾配定理）と呼ばれる注目すべき事実に依拠します。この定理は、ポリシー$\pi(a \mid s, \theta)$が状態（すなわちユーザーのコードベースの状態）における行動$\alpha$の分布を$s \sim P(s)$により指定し、パラメータが$\theta$であるとき、報酬が$J(\theta) = \mathbf{E}_{s \sim P(s),\, a \sim \pi(a \mid s,\theta)}\bigl[R(s,a)\bigr]$で与えられるなら、報酬の勾配は次のように表されると述べています。

これは有用です。というのも、右辺は推定しやすいからです。ユーザーのリクエストに表示された提案からサンプルを取得することで、状態と行動のサンプル$s \sim P(s), \; a \sim \pi(a \mid s, \theta)$を得ることができ、PyTorch のようなフレームワークを用いれば$\nabla_{\theta} \,\log \,\pi(a \mid s, \theta)$を計算でき、ユーザーが提案を受け入れたかどうかを見ることで$R(s,a)$を算出できます。したがって、この方程式を用いて$\nabla_{\theta} J(\theta)$の不偏推定量を得られ、確率的勾配降下法によってポリシーを改善できます。

ただし、これは最適化対象のポリシーからサンプリングされたアクションに対してのみ機能します。ポリシーを更新すると、もはや最適化中のポリシーからのサンプルはなく、前のポリシーからのサンプルしかありません。新しい「オンポリシー」サンプルを得るには、新しいモデルをユーザーにデプロイし、その挙動を観察する必要があります。つまり、新しいチェックポイントを素早くデプロイし、ユーザーにサジェストが表示されてからそのデータが学習プロセスの次段階に届くまでの時間を最小化するための良いインフラが必要でした。現在、チェックポイントのロールアウトと次のステップ用データの収集には1.5〜2時間かかります。これはAI業界で一般的な水準と比べれば速いものの、さらに高速化できる余地があります。

新しいTabモデル

ここで説明した手法を用いて、新しい Tab モデルを学習し、Cursor のデフォルトにしました。このモデルは、以前のモデルと比べて提案数が21%少なくなる一方で、行った提案の受け入れ率が28%高くなります。これによりコーディング体験が向上することを願っており、今後もこれらの手法をさらに発展させていく予定です。