如何訓練最強程式碼大模型?北大aiXcoder-7B貢獻前沿實踐

机器之心發表於2025-02-12

圖片

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年,機器之心AIxiv專欄接收報導了2000多篇內容,覆蓋全球各大高校與企業的頂級實驗室,有效促進了學術交流與傳播。如果您有優秀的工作想要分享,歡迎投稿或者聯絡報導。投稿郵箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

本文的通訊作者是北京大學計算機學院長聘教授李戈。

本文一作是 aiXcoder 蔣思源和北大李戈教授課題組博士生李佳,團隊重點關注融合深度學習與軟體工程的程式碼建模方法。

如何訓練一個程式碼大模型?這一過程看似簡單:獲取程式碼資料、清洗資料,最終啟動訓練。如今,開原始碼資料集層出不窮;資料清洗工具也已成熟,包括開源的許可證識別工具、MinHash 演算法、PII 識別模型等;而在分散式訓練方面,像 Megatron-LM、DeepSpeed 等框架也大大降低了技術門檻。看似我們只差計算資源,就能訓練出一個強大的程式碼大模型。

然而,訓練模型的初衷,應該始終從實際開發場景出發。作為開發者,我們不僅需要了解定義的各種 API 介面,還需要從入口函式模擬程式的執行過程,追蹤到每一行修改的程式碼。在複雜的專案中,任何小小的變動都可能影響整個系統的運轉。

但現有的程式碼大模型並未充分考慮到軟體開發的具體場景,它們往往將最終版本的程式碼簡單地視作自然語言文字,試圖透過複製自然語言處理的成功經驗來處理程式碼。這種方法忽略了程式碼的結構性和複雜的上下文關係,導致模型在實際開發中表現不佳。

北京大學 aiXcoder 團隊一直致力於探索如何將深度學習與軟體開發深度融合,推動軟體開發的自動化。2024 年 4 月,aiXcoder 開源了自研程式碼大模型 aiXcoder-7B,成為這一領域的一次重要嘗試,旨在將程式碼的抽象語法樹(AST)結構與大規模預訓練結合,以期提升模型對程式碼結構和上下文的理解能力。

近期,該篇論文被軟體工程領域國際頂級會議 ICSE 2025 收錄,將於 4 月 27 日 - 5 月 3 日赴加拿大渥太華參會分享研究成果。

此次論文錄用不僅是對 aiXcoder 7B 程式碼大模型技術前瞻性和應用創新性的高度認可,更標誌著該模型繼成功落地企業並獲各行業客戶廣泛認可後,再次於學術界獲得權威肯定,充分彰顯了 aiXcoder 在推動軟體工程發展中的前瞻性引領作用。

圖片

  • 論文地址:https://arxiv.org/pdf/2410.13187
  • 開源專案地址:https://github.com/aixcoder-plugin/aiXcoder-7B

程式碼資料,異於自然語言

相較於自然語言文字,程式是現實世界解決方案在計算機系統中的對映。因此,程式原始碼呈現出很多獨特的性質,例如:強結構性、可執行性等等。有效地表示和建模這些特性,對於程式碼生成等任務來說至關重要。

圖片

如上三行程式碼能夠嚴格解析為抽象語法樹格式

程式碼天然能被解析為抽象語法樹,其語法規則嚴格組織了程式碼語句之間的關係。在語法規則之上,也有很多方式描述程式碼之間的流轉關係,例如控制流圖、呼叫流圖等等。顧名思義,控制流圖會展示整個程式碼控制與條件關係,什麼樣的條件下哪個分支程式碼會執行。呼叫流圖則展示的是程式碼之間的呼叫關係,實現一個功能時在什麼樣的地方呼叫什麼樣的程式碼模組是能展示出來的。

圖片

控制流圖示例,程式碼執行條件與順序會解析成流程圖。

圖片

呼叫流圖示例,main 函式呼叫 calculate 函式計算兩個數之和,calculate 函式呼叫另外兩個函式 getFirst 和 getSecond 獲取參與計算的兩個加數。

程式語言與自然語言之間存在顯著差異。儘管大模型透過大規模自迴歸訓練任務在通用知識學習上取得了巨大成功,但這並不意味著可以簡單地將程式碼資料視為「自然語言」,並將其拉長為一維 Token 序列進行自迴歸訓練,就能複製自然語言處理的成功。

事實上,當使用自迴歸模型或「Fill in the middle」任務訓練基礎模型時,會發現實際在程式碼補全任務中,模型生成的結果往往與人類程式設計師的程式設計方式不符,我們還需要更符合程式碼的預訓練方法。

aiXcoder-7B:創新在 LLM 上引入程式碼特性

正因為當前程式碼大模型很少將程式碼特性引入到 LLM 的訓練過程中,程式碼大模型在企業真實專案中表現得不盡人意,所以我們創新將一些傳統軟體工程方法引入到大規模預訓練中,希望能生成更符合真實場景的程式碼內容。

為此,aiXcoder-7B 主要從以下幾個方面最佳化預訓練:

  • 資料預處理:軟工工具保證程式碼資料語法正確且不存在嚴重 Bug
  • 結構化 FIM:按照語法結構組織預訓練任務
  • 多檔案排序:保證單專案內,檔案排序既考慮內容相似,又考慮呼叫關係

資料預處理

aiXcoder 核心資料集主要用於強化程式碼大模型在以上程式語言上的效果,其經過大量的過濾與篩選過程。相比於其它程式碼大模型,aiXcoder-7B 預訓練資料既採用常規的資料處理,例如資料去重、自動生成程式碼去除、透過 Star 量、正則等規則去除低質量程式碼、敏感資訊等,同時藉助軟體工程方法進行更精細的資料處理。

圖片

具體而言,aiXcoder-7B 預訓練資料採用語法分析和靜態分析兩大類工具預處理資料。對於語法分析,重點解析五十種主流語言的語法結構,並排除存在語法錯誤、簡單 Bug 、大面積被註釋掉的程式碼等。

圖片

語法分析能天然解析並處理明顯不合理的程式碼

對於靜態分析,則側重解析十餘種最主流程式語言的嚴重錯誤,即當出現這一些型別錯誤時,程式碼大機率在執行過程中會出現比較大的問題。具體而言,掃描並定位影響程式碼可靠性和可維護性的 161 種 Bug,影響程式碼安全性的 197 種安全漏洞。

圖片

靜態分析能檢測出很多更深層缺陷與漏洞的程式碼。

結合軟體工程分析方法以及過濾規則,能夠將存在明顯問題的程式碼刪除掉,明顯提升整體程式碼質量。

結構化 FIM

在實際開發過程中,程式碼具有類、方法、條件程式碼塊、迴圈程式碼塊等眾多結構。研究團隊期待讓程式碼大模型天然能學會這樣的結構,而不是放任程式碼大模型向下一直生成,或者從字元層面上擷取一個片段,期待補全該字元片段。

為此,團隊結合語法分析方法,將程式碼解析為抽象語法樹,並基於語法樹的結構構建訓練任務。具體而言,程式碼檔案中的每個位置都對應著抽象語法樹中的某個節點。在訓練過程中,團隊挖掉該節點的子節點,或者挖掉該節點所在父節點剩餘的部分,然後針對被挖掉的程式碼塊做一個先驗約束:挖掉的程式碼塊橫跨一個或少數幾個完整的程式碼結構。將這部分完整程式碼結構用來計算損失訓練模型,就能一定程度上讓程式碼模型理解部分語法結構。

更形象地解釋,常規的 Fill in the middle 會構造很多不合法的程式碼片段,例如下圖「or i in range (2」,常規的做法只是從字元上隨機取一個片段。但論文研究團隊提出的 Structured Fill-In-the-Middle (SFIM) 會隨機先選定一個語法節點「IF」,並在 IF 節點向下取了「Compare」程式碼片段「i % 5 == 0:」

圖片

最終團隊在預訓練中根據 SFIM 構建整體訓練損失計算,以此更好地學習程式碼的語法結構資訊。

圖片

多檔案排序

當前主流的程式碼開源資料集,例如 TheStackV1 、 TheStackV2 或者 The Pile 中程式碼部分,都是根據單個語言,甚至單一字尾名組織資料,致使整個訓練樣本的構造侷限在單語言檔案中。而此次研究團隊構建的訓練資料以專案為單位,保留與處理多種程式語言的程式碼檔案,確保訓練資料中程式語言的分佈與真實開發一致。

此時有一個重要的問題:專案內不同檔案該如何排序?

為了提升模型對專案內多程式碼檔案關係的充分建模能力,並在推理過程中更高效地抽取有用的上下文資訊,研究團隊透過相似性關係和依賴性關係對程式碼檔案排序:相似性關係即模型在預訓練中能學會仿寫相似的程式碼;依賴性關係即模型在預訓練中能學會 API 呼叫或者函式呼叫的關係。

圖片

預訓練中,專案內檔案排序演算法

如 Algorithm 1 所示,本論文給出了一個專案內檔案排序虛擬碼。簡單理解,以 0.3 的機率採用檔案內容相似排序,即透過 KMeans 聚類演算法將檔案聚成不同的簇,並且同一個簇排列在一起;此外,以 0.3 的機率進行路徑相似排序,把同一目錄下的檔案,或者被測程式碼與測試程式碼等路徑相關的檔案能排列在一起;最後還以 0.3 的機率構建函式呼叫流圖,並根據圖的葉節點一路向根節點建立程式依賴路徑,將路徑上的程式碼檔案排列在一起。

aiXcoder 7B 獨特的效果優勢

藉助軟體工程方法,研究團隊透過更符合程式碼大模型的預訓練方法,提升了其在程式碼資料上的理解與生成能力。例如論文表 5 中的 Fill-in-the-middle 評測集顯示,經過高質量程式碼資料的 SFIM 任務訓練,不同語言的程式碼補全能力有明顯的提升。

圖片

為了進一步測評 aixcoder-7B 在多種情況下的程式碼補全能力,團隊從方法簽名、方法體、方法區域性、條件塊、迴圈塊、異常捕捉塊等維度評估了模型在程式碼補全上的效果,如論文圖 4 所示。對比 DeepSeekcoder-6.7B,aixcoder-7B 大部分的補全位置都擁有更好的效果。

圖片

此外,因為預訓練任務充分考慮了程式碼的語法結構,模型在推理過程中對程式碼的上下文結構展現出更出色的感知能力,能夠準確判斷需要補全完整的語法結構,並傾向生成更短的程式碼片段。如論文表 6 所示,模型生成的 Token 數與 GroundTruth Token 數的比值,aiXcoder -7B 更小,表明 SFIM 預訓練任務有效指導了模型更好學會如何終止預測。

圖片

對於程式碼補全任務,另外一個比較重要的是跨檔案上下文的理解能力。aiXcoder -7b 在預訓練中以專案為單位對專案內的程式碼檔案進行排序,獲得了更好的檔案間建模能力。如表 4 所示,aiXcoder -7b 在 CrossCodeEval 評測集上擁有更好的效果,表明其利用多檔案的上下文資訊,補全當前程式碼檔案能力更有優勢。

圖片

後續改進方向

在真實軟體開發場景中,還有很多能力是大模型未曾學習到的,重中之重即程式碼上下文。

實際程式碼補全往往需要基於不同型別的上下文(如:當前檔案的上文、跨檔案上下文、相似程式碼段),去預測後續的程式碼。這種複雜的上下文形式與基礎模型預訓練時的上下文形式不一致,從而限制了基礎模型在實際應用時的程式碼補全準確率。

為解決這個問題,研究團隊在 aiXcoder 7B 上做了更多的對齊訓練實驗。該對齊訓練有效地將模型對齊到真實軟體開發場景中的上下文形式,顯著地提升了模型在多種語言上的程式碼補全準確率。例如,在四種語言(Python、Java、C++和Go)的多行補全上,相較於aiXcoder-7B,經過最佳化的新模型在Exact Match(完全匹配)指標上平均取得了 13 個點的絕對提升。

當前,充分利用數十年積累的軟體工程經驗,將程式碼大模型真正應用於軟體開發的實際場景中,仍然是一項艱鉅而複雜的任務。然而,隨著不斷深入的研究,程式碼大模型已經讓「軟體開發自動化」這一宏偉目標變得愈加觸手可及。

相關文章