近段時間，AI 程式設計工具 Cursor 的風頭可說是一時無兩，其表現卓越、效能強大。近日，Cursor 一位重要研究者參與的一篇相關論文釋出了，其中提出了一種方法，可透過搜尋自然語言的規劃來提升 Claude 3.5 Sonnet 等 LLM 的程式碼生成能力。
具體來說，他們提出的方法名為 PlanSearch（規劃搜尋）。主導團隊是 Scale AI，本文一作為 Scale AI 研究者 Evan Wang。二作 Federico Cassano 現已加入如今炙手可熱的 AI 程式設計工具公司 Cursor。他曾參與創立了 GammaTau AI 專案，該專案的目標是實現 AI 程式設計的民主化。此外，他也是 BigCode 專案的活躍貢獻者，該專案負責開發用於 AI 程式設計的 StarCoder 系列大型語言模型。
論文標題：Planning In Natural Language Improves LLM Search For Code Generation
論文地址：https://arxiv.org/pdf/2409.03733
論文開篇，該團隊提到強化學習教父 Sutton 的經典文章《The Bitter Lesson（苦澀的教訓）》揭示的 Scaling Law 的兩大核心原則：學習和搜尋。隨著大型語言模型的迅猛發展，人們對於「學習」是否有效的疑慮已基本消除。然而，在傳統機器學習領域中表現出色的「搜尋」策略，將如何擴充大模型的能力，還是個未知數。
目前阻礙模型應用「搜尋」的主要難題是模型給出的答案過於雷同，缺乏多樣性。這可能是由於在預訓練的基礎上，模型會在特定的資料集上進行進一步的訓練，以適應特定的應用場景或任務所導致的。
經過大量實證研究證明，許多大語言模型往往會被最佳化，以產生一個正確的答案。比如下圖中所示，DeepSeek-Coder-V2-Lite-Base 的表現不如其基礎模型，但隨著回答的多樣性的減少，情況發生了逆轉。多個模型都存在這種現象：經過特別指令調整的模型在只生成一個答案的情況下（pass@1）通常比基礎模型表現得好很多，但當需要生成多個答案時，這種優勢就不明顯了 —— 在某些情況下，甚至完全相反。
模型在生成答案時缺乏多樣性，這對於搜尋的效果非常不利。特別是在極端情況，比如採用「貪心解碼」，模型給出的答案會非常相似，因為它們是從模型中重複抽取的。這種情況下，即使模型花費更多推理時間，也難以獲得更好的搜尋結果。
通行的大模型排行榜，例如例如 LMSYS Chatbot Arena、LiveCodeBench、OpenLLMLeaderboard，很難反應模型在回答多樣性方面的不足。這些排行榜主要關注模型在單一樣本上的透過率，沒有考慮到模型在更廣泛場景下的表現。由於模型需要很快地響應使用者的需求，單一樣本的回答質量是衡量一個聊天機器人的關鍵指標，但這一指標並不足以全面評估模型在允許更充裕推理時間時的綜合效能。
針對以上問題，研究人員對如何在大語言模型推理過程中提高回答的多樣性進行了探索。對此，他們提出了假設，想讓模型輸出的答案更加豐富，需要在自然語言的概念或想法的空間內進行搜尋。
為了驗證這個假設，研究人員進行了一系列實驗。首先，研究人員發現，如果給模型一些簡單的草圖（這些草圖是從已經能解決問題的程式碼中「回譯」而來），模型就能根據這些草圖寫出正確的最終程式。其次，研究人員還發現，如果讓模型在嘗試解決問題之前，先在 LiveCodeBench 上想出一些點子（這個過程叫做 IdeaSearch / 思路搜尋），然後看看模型能不能用這些點子解決問題。
結果發現，模型要麼完全解決不了問題（準確度為 0%），要麼就能完美解決問題（準確度為 100%）。這表明當模型嘗試解決一個問題時，成功與否主要取決於它最初的那個想法（草圖）對不對。
根據這兩個實驗的結果，研究人員認為一種提升 LLM 程式碼搜尋能力的自然方法是：搜尋正確的思路，然後實現它！
於是，規劃搜尋（PlanSearch）方法誕生了。
不同於之前的搜尋方法（通常是搜尋單個 token、程式碼行甚至整個程式）不一樣，規劃搜尋是搜尋解決當前問題的可能規劃。這裡，規劃（plan）的定義是：有助於解決某個特定問題的高層級觀察和草案的集合。
為了生成新規劃，規劃搜尋會生成大量有關該問題的觀察，然後再將這些觀察組合成用於解決問題的候選規劃。
這個操作需要對生成的觀察的每個可能子集都執行，以最大化地鼓勵在思路空間中進行探索，之後再將結果轉譯成最終的程式碼解決方案。
該團隊的實驗發現，在推理時有效使用計算方面，規劃搜尋方法優於標準的重複取樣方法以及直接搜尋思路的方法。
方法
在這項研究中，該團隊探索了多種不同方法，包括重複取樣（Repeated Sampling）、思路搜尋（IdeaSearch）以及新提出的規劃搜尋（PlanSearch）。其中前兩種方法顧名思義，比較直觀，這裡我們重點關注新提出的規劃搜尋。
該團隊觀察到，雖然重複取樣和思路搜尋能成功地提升基準評測的結果。但在很多案例中，多次提示（pass@k）（即使在溫度設定很高）只會導致輸出程式碼發生很小的變化，這些變化只會改變一些小方面，但無法改善思路中的缺陷。
下面來看具體的規劃搜尋過程：
1. 透過提示來獲取觀察
首先假設有一個問題陳述 P，透過向 LLM 傳送提示詞來獲取對該問題的「觀察」/ 提示。這裡將這些觀察記為 O^1_i，其中 i ∈ {1, . . . , n_1}；這是因為它們是一階觀察。通常而言，n_1 的數量級在 3 到 6 之間。具體數量取決於 LLM 輸出。為了利用這些觀察結果來啟發未來的思路，該團隊建立了 O^1_i 的集合 S^1 的且大小至多為 2 的所有子集。其中每個子集都是觀察結果的一個組合。這裡將每個子集記為 C^1_i，其中 i ∈ {1, . . . , l_1}，而
2. 推導新的觀察
這樣一來，所有觀察結果的集合都可以定義為深度為 1 的有向樹，其中根節點為 P，並且每個 C^1_i 都有一條從 P 指向 C^1_i 的邊。
然後，在每個葉節點 C^1_i 上重複上一步流程，從而生成一個二階觀察集 S^2。為了得到二階觀察，該團隊的做法是在給模型的提示詞中包含原始問題 P 和 C^1_i 中包含的所有觀察 —— 這些觀察被構造為解決 P 所必需的原始觀察。然後再提示 LLM，讓其使用 / 合併在 C^1_i 中找到的觀察來得出新的觀察。
這個過程可以繼續延伸，但由於計算限制，這裡在深度為 2 時對該樹進行了截斷操作。
3. 將觀察變成程式碼
在得到了觀察之後，必須先將它們實現成具體思路，然後再將它們轉譯成程式碼。
具體來說，對於每個葉節點，將所有觀察以及原始問題 P 放入提示詞來呼叫 LLM，以便生成問題 P 的自然語言解決方案。為了提升多樣性，對於每個生成的思路，該團隊透過假設該思路是錯誤的來生成一個額外的思路，並要求 LLM 給出批評 / 反饋，從而將提議的思路翻倍了。
然後，再將這些自然語言解決方案轉譯成虛擬碼；再把這些虛擬碼轉譯成真正的 Python 程式碼。
實驗
實驗採用了三個評估基準：MBPP+、HumanEval+ 和 LiveCodeBench。引數設定等細節請參閱原論文。
至於結果，該團隊報告了三種方法的結果，包括重複取樣、思路搜尋和規劃搜尋，見表 1、圖 1 和圖 5。
可以看到，規劃搜尋和思路搜尋的表現明顯優於基礎的取樣方法，其中規劃搜尋方法在所有實驗方法和模型上都取得了最佳分數。
圖 7、8、9 展示了在每個資料集上的詳細 pass@k 結果。
可以看到，在 Claude 3.5 Sonnet 上使用規劃搜尋方法時，在 LiveCodeBench 基準上得到了當前最佳的 pass@200 效能：77.0%。該表現優於不使用搜尋時獲得的最佳分數（pass@1 = 41.4%）以及標準的 best-of-n 取樣方法的分數（pass@200 = 60.6%）。
此外，使用小型模型（GPT-4o-mini）執行規劃搜尋時，僅僅 4 次嘗試後就能勝過未使用搜尋增強的大型模型。這佐證了近期一些使用小模型進行搜尋的有效性的研究成果。
在另外兩個程式設計基準 HumanEval+ 和 MBPP+ 上，規劃搜尋也能帶來類似的提升。
透過研究特定模型的差異，該團隊注意到 pass@k 曲線所呈現的趨勢在所有模型中並不統一；事實上，每條曲線看起都不一樣。該團隊猜想部分原因是思路多樣性的變化。
該團隊還得到了一個有趣的觀察結果：規劃搜尋並不利於某些模型的 pass@1 指標，其中最明顯的是 Sonnet 3.5 在 LiveCodeBench 上的表現 —— 這是實驗中表現最好的組合。
該團隊基於直覺給出瞭解釋：提升思路多樣性可能會降低生成任何特定思路的機率，同時增加在給定池中至少有一個正確思路的機率。因此，pass@1 可能會略低於平常，但也正是由於這個原因，pass@k 指標可能會優於缺乏多樣性的思路池。
另外，表 1 和圖 1 給出了在嘗試 / 完成上經過歸一化的主要結果。其中針對每個問題，每種搜尋方法都可以嘗試 k 次。
最後，該團隊還發現，在思路空間中觀察到的多樣性可用於預測搜尋效能，這可透過模型 / 方法的 pass@1 與其 pass@200 之間的相對改進計算得到，如圖 6 所示。
雖然熵是最常見的多樣性度量是，但由於種種原因，熵不足以精確衡量 LLM 的多樣性。
因此，該團隊測量多樣性的做法是在所有生成的程式上使用簡單的配對策略，將其置於思路空間中進行計算。具體演算法請訪問原論文。