跳至主內容

以線上強化學習改進 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 模型,讓它一開始就避免產生不佳的建議。因此,我們改採用策略梯度(policy gradient)方法。

方策梯度

策略梯度方法是一種通用的方式,用來優化「策略」(此處指 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