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

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

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

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

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

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

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

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

方策勾配

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

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

Tab を改善するためにポリシー勾配法を用いるにあたり、受け入れられた提案を促しつつ、受け入れられない提案をユーザーに表示することを抑制する報酬を定義しました。たとえば、受け入れ確率が少なくとも 25% の場合にのみ提案を表示したいとします。このとき、受け入れられた提案には 0.75 の報酬、拒否された提案には -0.25 の報酬、提案を表示しない場合には 0 の報酬を割り当てます。受け入れ確率を p とすると、提案を表示したときの期待報酬は 0.75*p-0.25*(1-p) となり、これは p > 0.25 のときに限って正になります。したがって、報酬最大化を行うポリシーは、受け入れ確率が少なくとも 25% と見積もられるときに提案を表示し、それ以外の場合は何も表示しません。

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

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

ポリシー更新を得るために、私たちはPolicy Gradient Theoremと呼ばれる注目すべき事実に依拠します。これは、ポリシー (a | s,) が状態（すなわち、ユーザーのコードベースの状態） sP(s) における行動 a の分布を theta によりパラメータ化して定め、報酬が J()=EsP,a[R(s,a)] で与えられるとき、報酬の勾配は次で与えられる、というものです。

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

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

新しいTabモデル

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