Arctic-SnowCoder揭秘:小數(shù)據(jù)如何煉成高性能代碼模型?
一、結論寫在前面
論文標題:Arctic-SnowCoder: Demystifying High-Quality Data in Code Pretraining
論文鏈接:??https://arxiv.org/pdf/2409.02326??
高質量數(shù)據(jù)對于語言模型的有效預訓練至關重要。然而,“高質量”的精確定義仍未得到充分探索。
聚焦于代碼領域,論文引入了Arctic-SnowCoder-1.3B,這是一個數(shù)據(jù)高效的基礎代碼模型,通過三個階段的逐步精煉數(shù)據(jù)進行預訓練,共處理了555B token:(1) 使用500B個標準質量代碼token進行通用預訓練,經(jīng)過基本過濾、去重和去污染預處理;(2) 使用50B個高質量token進行持續(xù)預訓練,這些token從第一階段中通過BERT風格的質量注釋器選出,該注釋器經(jīng)過訓練以區(qū)分優(yōu)質代碼與隨機數(shù)據(jù),使用從高質量代碼文件中提取的正例,以及來自Magicoder和StarCoder2-Instruct的指令數(shù)據(jù);(3) 使用5B個合成數(shù)據(jù)進行增強預訓練,這些數(shù)據(jù)由Llama-3.1-70B使用第二階段數(shù)據(jù)作為種子生成,采用Magicoder的預訓練方法。
僅僅訓練有限數(shù)據(jù)集,Arctic-SnowCoder-1.3B在BigCodeBench上取得了SOTA的結果,這是一個專注于實際和具有挑戰(zhàn)性的編程任務的編碼基準,在訓練了小于等于 1T tokens的類似大小的模型中。特別是,它比Phi-1.5-1.3B [20]高出36%。盡管訓練數(shù)據(jù)量為555B tokens,與其他在萬億tokens上訓練的最先進的小型代碼模型相比,Arctic-SnowCoder在多個基準測試中達到了或超過了這些模型的性能。
例如,Arctic-SnowCoder-1.3B在所有評估的基準測試中擊敗了StarCoderBase-3B,后者訓練數(shù)據(jù)超過1T tokens。Arctic-SnowCoder-1.3B在HumanEval+ (28.0 vs. 27.4)上優(yōu)于StarCoder2-3B,后者訓練數(shù)據(jù)超過3T tokens,這是一個評估函數(shù)級代碼生成的基準,同時在BigCodeBench(19.4 vs. 21.4)上保持競爭力。論文進行了全面的消融研究,以驗證訓練Arctic-SnowCoder背后的設計決策:
?首先,論文的研究結果表明,在一般預訓練中,將文件級數(shù)據(jù)按編程語言劃分成倉庫后,顯著優(yōu)于僅按倉庫名稱分組數(shù)據(jù)的方法。
?此外,論文確定了最佳的學習率計劃,包括重新預熱階段后進行線性衰減,以及在持續(xù)預訓練期間高質量數(shù)據(jù)的理想重復次數(shù),論文發(fā)現(xiàn)是四次。
?更重要的是,論文對基于模型的質量注釋器的比較,這些注釋器在各種數(shù)據(jù)組合上進行訓練,強調了預訓練數(shù)據(jù)與下游任務的對齊對于實現(xiàn)卓越性能至關重要。
圖1:Arctic-SnowCoder-1.3B的三階段預訓練,逐步使用更高質量的數(shù)據(jù)。
二、論文的簡單介紹
2.1 論文的背景
預訓練LLMs通常依賴于大量數(shù)據(jù)。在代碼等專業(yè)領域,這種對數(shù)據(jù)量的重視尤為明顯,研究人員通過抓取GitHub等平臺獲取大量代碼預訓練數(shù)據(jù)集。然而,最近的研究越來越表明,高質量的數(shù)據(jù)對于有效的預訓練至關重要, 包括代碼領域。
在通用領域,研究人員探索了多種技術來策劃高質量的語言模型預訓練數(shù)據(jù)。FineWeb-Edu使用基于Snowflake-arctic-embed-m嵌入構建的線性回歸器來評估網(wǎng)頁的教育價值并選擇高質量內容,而DCLM方法則采用基于fastText的過濾器,該過濾器在高質量在線資源和指令數(shù)據(jù)的正例以及隨機負例網(wǎng)頁上進行訓練,以識別高質量文本。與使用未過濾的大規(guī)模數(shù)據(jù)集相比,這些基于模型的質量過濾器顯著提高了語言模型在下游任務中的性能。
同樣,研究人員也認識到高質量代碼數(shù)據(jù)對于預訓練的重要性,Phi-1使用隨機森林分類器在Code-Gen嵌入上選擇教育性代碼樣本,DeepSeek-Coder-V2采用多階段fastText管道來召回與網(wǎng)絡相關的代碼數(shù)據(jù)和GitHub上的高質量代碼,實現(xiàn)了最先進的編碼性能。
在本文中,論文介紹了Arctic-SnowCoder-1.3B,這是一個高性能的小型代碼模型,通過一種新穎的三步訓練方法創(chuàng)建,專注于數(shù)據(jù)質量的逐步提升。得益于這種方法,Arctic-SnowCoder-1.3B在所有評估的基準測試中均優(yōu)于StarCoderBase-3B [19],并在復雜且實用的BigCodeBench基準測試 [46] 上超越Phi-1.5-1.3B [20] 36%,該基準對實際編程至關重要。
2.2 論文的方法--Arctic-SnowCoder
這里論文將詳細解釋Arctic-SnowCoder-1.3B的訓練方法,如圖1所示。論文首先討論原始訓練數(shù)據(jù)的組成(見圖1),然后概述通用預訓練階段。接下來,論文描述使用高質量數(shù)據(jù)的協(xié)同預訓練過程,最后,論文詳細闡述使用合成數(shù)據(jù)的增強預訓練。模型架構基于Llama-2,具體細節(jié)見表1。
表1:Arctic-SnowCoder的模型架構細節(jié)。
2.2.1原始數(shù)據(jù)
用于訓練Arctic-SnowCoder-1.3B的原始預訓練數(shù)據(jù)僅包含代碼,主要來源于用于訓練Snowflake Arctic的編碼數(shù)據(jù)。該數(shù)據(jù)結合了The Stack v1和GitHub爬取數(shù)據(jù)的清洗后版本。從這些數(shù)據(jù)中,論文選擇了18種流行的編程語言進行訓練,類似于StarCoder2-3B。這些語言包括Python、Java、C++、C 、JavaScript、PHP、C#、Go、TypeScript、SQL、Ruby、Rust、Jupyter Notebook、Scala、Kotlin、Shell、vart、Swift,總計400B唯一token。
2.2.2 通用預訓練
在通用預訓練階段,模型使用Adam進行訓練,訓練500B token,序列長度為8192,批量大小為512。學習率在600次迭代后進行線性預熱,隨后遵循余弦衰減。論文設置最大學習率為5.3x 10^{-4},最小學習率為5.3 x 10^{-5},這一設置參考了DeepSeek-Coder。在此階段,論文使用了全部400B原始數(shù)據(jù),未進行額外的質量過濾。論文首先按編程語言對代碼文件進行分區(qū),按倉庫進行分組,然后以隨機順序連接它們,類似于StarCoder2的方法。論文展示了首先按編程語言對代碼文件進行分區(qū)的優(yōu)勢。論文將此階段生成的模型命名為Arctic-SnowCoder-alpha。
2.2.3 高質量數(shù)據(jù)繼續(xù)預訓練(Continued pretraining)
在通用預訓練之后,論文使用從相同原始預訓練語料庫中提取的50B高質量token繼續(xù)預訓練Arctic-SnowCoder-alpha。這50B高質量token是通過將12.5B個由論文的代碼質量注釋器評分的前百分位代碼文件token重復4次形成的。
受FineWeb-Edu和DCLM的啟發(fā),論文在基于BERT的先進嵌入模型Snowflake-arctic-embed-m之上訓練了一個線性分類頭。訓練數(shù)據(jù)包括30萬個正樣本,采樣自22萬個高質量開源代碼文件、8萬個來自Magicoder和StarCoder2-Instruct的高質量指令數(shù)據(jù),以及從預訓練語料庫中隨機選擇的300個代碼文檔。
關于代碼質量的先前研究,如Phi-1,往往過分強調代碼的“教育價值”,使模型偏向于像HumanEva這樣的簡單基準。論文展示了論文的注釋方法能夠帶來更平衡的模型能力提升。
此外,鑒于這些代碼文檔通常超過1000個token,超過了BERT的512個token的上下文窗口大小,論文改進了FineWeb-Edu的流程,通過平均質量注釋器產生的頂部、中部和底部部分的分數(shù)來計算每個文件的分數(shù)。在此階段,論文從0到最大預訓練學習率5.3x 10^{-4}進行1000次迭代的學習率預熱,隨后進行線性衰減至0。此階段生成的模型稱為Arctic-SnowCoder-beta。
2.2.4 使用合成數(shù)據(jù)增強預訓練
在增強預訓練階段,論文利用Llama-3.1-70B-Instruct 生成比持續(xù)預訓練階段更高質量的數(shù)據(jù),并將Python混合比例提高到約50%,同時保持其他語言的比例不變。Phi-1 [13]表明,類似教科書的合成預訓練數(shù)據(jù)可以顯著提升模型性能。
然而,過度依賴此類數(shù)據(jù)可能會導致模型分布偏斜,從而可能損害其在實際編碼任務中的有效性。例如,論文后面展示的,Phi-1.5在HumanEvalt和MBPP+上表現(xiàn)出色,這些任務類似于教科書練習,但在BigCodeBench [46]中更復雜和實用的編碼任務上表現(xiàn)較差。為了解決這個問題,論文改編了Magicoder的OSS-Instruct方法用于預訓練。最初,OSS-Instruct旨在通過提示模型創(chuàng)建受開源代碼片段啟發(fā)的問答對來生成現(xiàn)實的指令調優(yōu)數(shù)據(jù)。
相比之下,論文通過使用Llama-3.1-70B-Instruct生成高質量和面向問題解決的代碼文件來生成高質量的合成預訓練數(shù)據(jù),這些代碼文件以持續(xù)預訓練階段中評分最高的代碼文檔為種子。后面論文展示了每個預訓練階段都顯著優(yōu)于前一個階段,突顯了逐步提高數(shù)據(jù)質量的有效性。
2.3 論文的效果
這里論文將Arctic-SnowCoder與最先進的小型語言模型進行比較,并展示了每個預訓練階段的性能提升,評估了兩種形成通用預訓練中倉庫級別數(shù)據(jù)的策略,并對持續(xù)預訓練中的設計選擇進行了詳細的消融分析。
2.3.1 實驗設置
論文考慮以下四個多樣化的編程基準,以全面評估不同代碼模型的代碼生成能力:
?HumanEval+和MBPP+ 。HumanEval和MBPP是用于函數(shù)級代碼生成的兩個最廣泛使用的基準。論文采用了EvalPlus增強的版本,提供了80倍/35倍的更多測試用例以進行嚴格評估。HumanEvalt和MBPP+分別包含164和378個編碼問題。
?EvoEval 是一個程序合成基準測試套件,通過將現(xiàn)有基準測試演化為不同的目標領域而創(chuàng)建。論文采用了其五個默認的轉換類別,即困難、創(chuàng)造性、微妙、組合和工具使用,總計500個任務。
?BigCodeBench 通過實際且具有挑戰(zhàn)性的編程任務評估語言模型。它包含1140個編程任務,每個任務都是通過人類與語言模型的協(xié)作創(chuàng)建的,任務質量由人類專家保證。
2.3.2 基線比較與三階段預訓練的有效性
表 2 :將Arctic-SnowCoder與最先進的小型語言模型 ( \mathrm{< 3 B} ) 進行比較,按訓練計算量 > 1T 標記進行劃分。Arctic-SnowCoder-alpha和Arctic-SnowCoder-beta分別是通用預訓練和繼續(xù)使用高質量數(shù)據(jù)預訓練后的檢查點。Arctic-SnowCoder是使用合成數(shù)據(jù)增強預訓練后的最終檢查點。
表2展示了多種小型語言模型(參數(shù)少于3B)在多個編碼基準上的綜合比較,按其訓練計算是否超過1T tokens進行分類。值得注意的是,Arctic-SnowCoder表現(xiàn)出色,特別是在其有限的訓練數(shù)據(jù)下。Arctic-SnowCoder-1.3B在BigCodeBench上達到了與訓練數(shù)據(jù)不超過1T token的同類模型相比的最新性能,顯著優(yōu)于StarCoderBase-3B、SmolLM-1.7B和Phi-1.5-1.3B。特別是,盡管Phi-1.5-1.3B在“教科書式”基準測試如HumanEval+、MBPP+和EvoEval上具有優(yōu)勢,但Arctic-SnowCoder-1.3B在更復雜和實用的Big-CodeBench上以36%的優(yōu)勢超越了Phi-1.5-1.3B。
此外,Arctic-SnowCoder-1.3B在所有評估基準上均擊敗了StarCoderBase-3B,后者是StarCoder2-3B的前身,訓練數(shù)據(jù)為1T tokens。盡管僅訓練了555B tokens,Arctic-SnowCoder-1.3B在HumanEvalt上與經(jīng)過更廣泛訓練的模型如StarCoder2-3B、StableCode-3B、CodeGemma-2B-v1.0和Qwen1.5-1.8B相媲美甚至超越。在EvoEval和BigCodeBench上,Arctic-SnowCoder仍然具有競爭力。
此外,該表還突顯了Arctic-SnowCoder在其訓練階段的持續(xù)改進:Arctic-SnowCoder-alpha、Arctic-SnowCoder-beta和最終的Arctic-SnowCoder。每個階段都建立在前一階段的基礎上,Arctic-SnowCoder在所有基準測試中均取得了最高分數(shù)。這種穩(wěn)步提升強調了高質量和合成數(shù)據(jù)在最終階段的關鍵作用。盡管從相同的數(shù)據(jù)開始,Arctic-SnowCoder的每次迭代都縮小了與最先進模型的差距,展示了整體訓練方法的有效性。
2.3.3 通用預訓練中的倉庫級數(shù)據(jù)
在通用預訓練階段,論文采用了StarCoder2的方法,通過隨機拼接文件內容將文件級數(shù)據(jù)隨機分組到倉庫中。在表3中,論文研究了兩種方法:(1)僅按倉庫名稱對文件進行分組,這意味著每個訓練文檔可以是多語言代碼文件的混合,如果倉庫是用不同語言編寫的;(2)在將文件分組到倉庫之前,先按編程語言對文件進行分區(qū),這意味著每個訓練文檔僅關注一種單一語言。
表3:兩種預訓練方法對倉庫級別數(shù)據(jù)分組的比較。(1)“按倉庫分組”將每個倉庫視為一個可能混合多種語言的單一訓練單元,以及(2)“按語言和倉庫分組”在按倉庫分組之前先按編程語言對數(shù)據(jù)進行分區(qū)。
論文可以觀察到,第二種方法,即論文在一般預訓練中最終采用的方法,明顯優(yōu)于第一種方法。
2.3.4繼續(xù)預訓練中的設計選擇
在繼續(xù)預訓練中,論文從預訓練語料庫中提取高質量的token,并訓練一個改進的基礎模型。為了獲得高質量的token,論文采用了基于模型的質量標注器。在本節(jié)中,論文實驗了各種設計選擇,包括標注器的訓練數(shù)據(jù)、繼續(xù)預訓練中使用的學習率以及高質量token的最佳重復次數(shù)。
基于模型的質量標注器 類似于FineWeb-Edu ,論文在Snowf lake-arctic-embed-m嵌入模型之上訓練一個線性頭來為每個代碼文件評分。在表4中,論文實驗了4種變體:
?ANN-EDU:論文提示Mixtral-8x7B-Instruct標注每個代碼文件的教育價值(1到5)。使用400k標注數(shù)據(jù)訓練一個線性回歸頭。對于以下變體,類似于DCLM ,論文隨機采樣負文檔并僅更改正部分。使用線性分類頭。
?ANN-INs:正樣本是來自ANN-EDU的100k教育數(shù)據(jù)(3.5+)和來自Magicoder [41]和StarCoder2-Instruct [40]的100k高質量指令數(shù)據(jù)的混合。
?ANN-HQ: 正樣本為220k開源、合成、高質量代碼文件。
?ANN-HQINs: 正樣本為220k ANN-HQ訓練數(shù)據(jù)與80k來自Magicoder[41]和StarCoder2-Instruct[40]的指令數(shù)據(jù)的混合。
表4: 通過應用基于模型的質量標注器(使用不同配方訓練)進行10B繼續(xù)預訓練的下游性能比較
訓練標注器后,論文首先將每個標注器應用于整個預訓練語料庫,為每個文件獲取一個分數(shù)。與僅掃描前2k字符的FineWeb-Edu不同,論文掃描代碼文件的頂部、中部和底部部分,并平均這些分數(shù)。然后,論文根據(jù)這些分數(shù)按語言對代碼文件進行排名,并選擇前百分位的文檔,直到達到大約10 B token。論文保持與預訓練中使用的相同混合比例。表中顯示,結合高質量文件和指令數(shù)據(jù)的ANN-HQINS實現(xiàn)了最佳的下游性能。
論文在圖2中進行了額外的分析。對于每個標注者,論文創(chuàng)建了一個驗證數(shù)據(jù)集,其中正樣本來自代碼解決方案基準,負樣本來自訓練期間未見過的隨機預訓練數(shù)據(jù)。論文使用ROC-AUC [6](受試者工作特征曲線下面積)分數(shù)來評估標注者在基準數(shù)據(jù)排名中的表現(xiàn)。該圖展示了每個基準的ROC-AUC分數(shù)與基準通過率之間的相關性。幾乎一致的趨勢是:較高的ROC-AUC分數(shù)導致更好的基準性能。良好的ROC-AUC分數(shù)表明標注者有效地塑造了下游任務的分布。因此,高質量的關鍵在于與下游應用分布的對齊。
圖2:注釋者ROC-AUC評分與基準測試pass@1之間的相關性。
學習率調度 論文還在表5中研究了不同的學習率調度策略,包括(1)從最小預訓練學習率線性退火至零,(2)使用最小預訓練學習率的恒定調度,以及(3)重新預熱至最大預訓練學習率后線性衰減至零。根據(jù)經(jīng)驗,論文發(fā)現(xiàn)重新預熱方法表現(xiàn)最佳,并在所有其他關于繼續(xù)預訓練的實驗中一致使用。
高質量數(shù)據(jù)重復次數(shù) 最后,論文將預訓練的token范圍從10 B擴展到50 B。剩下的一個問題是如何確定高質量token的最佳重復次數(shù)。論文通過選擇由ANN-HQINS排名的前百分位token進行實驗,重復次數(shù)從1到5,如表6所示。在這種情況下,前百分位token是最高質量的token。例如,1 x 50B表示前50B token的一次重復,而4 X12.5B表示前12.5B token的四次重復,確保所選token的質量最佳。
根據(jù)表中的結果,重復高質量token四次(4 x 12.5B)在下游多個評估指標中任務中表現(xiàn)最佳,四次重復(4 x 12.5B)在HumanEval和EvoEval中得分最高。兩次重復( 2 X25.0B )和三次重復( 3 X16.7B )也表現(xiàn)出強勁的性能,特別是在mbpp中。五次重復( 5 x10.0B )在MBPP中得分最高,但在總體指標上未超過四次重復。一次重復( 1 x50.0B )與多次重復相比,改進最小。
表6:使用ANN-HQINS在50B繼續(xù)預訓練中不同高質量數(shù)據(jù)重復次數(shù)的下游性能。
本文轉載自 ??AI帝國??,作者: 無影寺
