程式碼知識原來這麼重要。
如今說起大語言模型（LLM），寫程式碼能力恐怕是「君子六藝」必不可少的一項。
在預訓練資料集中包含程式碼，即使對於並非專門為程式碼設計的大模型來說，也已是必不可少的事。雖然從業者們普遍認為程式碼資料在通用 LLM 的效能中起著至關重要的作用，但分析程式碼對非程式碼任務的精確影響的工作卻非常有限。
在最近由 Cohere 等機構提交的一項工作中，研究者系統地研究了程式碼資料對通用大模型效能的影響。
論文連結：https://arxiv.org/abs/2408.10914
設問「預訓練中使用的程式碼資料對程式碼生成以外的各種下游任務有何影響」。作者對範圍廣泛的自然語言推理任務、世界知識任務、程式碼基準和 LLM-as-a-judge 勝率進行了廣泛的消融和評估，模型的引數大小從 4.7 億到 2.8 億個引數不等。
在各種配置中，我們可以看到存在一致的結果：程式碼是泛化的關鍵模組，遠遠超出了編碼任務的範圍，並且程式碼質量的改進對所有任務都有巨大影響。預訓練期間投資程式碼質量和保留程式碼資料，可以產生積極影響。
這裡有幾個因素很重要，包括確保程式碼比例正確、透過包含合成程式碼和程式碼相鄰資料（例如 commits）來提高程式碼質量，以及在冷卻期間等多個訓練階段利用程式碼。該研究結果表明，程式碼是泛化的關鍵構建塊，遠遠超出了編碼任務的範圍，程式碼質量的提高對效能有巨大的影響。
再進一步，作者對廣泛的基準進行了廣泛的評估，涵蓋世界知識任務、自然語言推理和程式碼生成，以及 LLM 作為評判者的勝率。在對 4.7 億到 28 億引數模型進行實驗後，以下是詳細結果：
1. 程式碼為非程式碼任務的效能提供了重大改進。使用程式碼預訓練模型進行初始化可提高自然語言任務的效能。特別是，與純文字預訓練相比，新增程式碼可使自然語言推理能力相對增加 8.2%，世界知識增加 4.2%，生成勝率提高 6.6%，程式碼效能提高 12 倍。
2. 程式碼質量和屬性很重要。使用標記樣式的程式語言、程式碼相鄰資料集（例如 GitHub commits）和合成生成的程式碼可提高預訓練的效能。特別是，與預訓練中的基於 Web 的程式碼資料相比，在更高質量的合成生成的程式碼資料集上進行訓練可使自然語言推理和程式碼效能分別提高 9% 和 44%。此外，與不包含程式碼資料的程式碼模型初始化相比，包含合成資料的程式碼模型持續預訓練分別使自然語言推理和程式碼效能相對提高 1.9% 和 41%。
3. 冷卻中的程式碼可進一步改善所有任務。在預訓練冷卻中包含程式碼資料，其中高質量資料集被加權，與冷卻前的模型相比，自然語言推理效能增加 3.6%，世界知識增加 10.1%，程式碼效能增加 20%。更重要的是，包含程式碼的冷卻比基線（無冷卻的模型）的勝率高出 52.3%，其中勝率比無程式碼的冷卻高出 4.1%。
方法概覽
在方法部分，研究者從預訓練資料、評估、訓練與模型細節三個部分著手進行介紹。下圖 1 為高階實驗框架。
預訓練資料
研究者描述了預訓練和冷卻（cooldown）資料集的細節。目標是在當前 SOTA 實踐的標準下，評估程式碼在預訓練中的作用。因此，他們考慮了由以下兩個階段組成的預訓練執行，即持續預訓練和冷卻。
其中持續預訓練是指訓練一個從預訓練模型初始化而來並在固定 token 預算下訓練的模型。冷卻是指在訓練的最後階段，提高高質量資料集的權重並對相對較少數量的 token 進行學習率的退火。
關於文字資料集，研究者使用 SlimPajama 預訓練語料庫作為他們的自然語言文字資料來源。
關於程式碼資料集，為了探索不同屬性的程式碼資料的影響，研究者使用了不同型別的程式碼源，包括如下：
基於 web 的程式碼資料，這是主要的程式碼資料來源，包括用於訓練 StarCoder 的 Stack 資料集。該資料集包含了爬取自 GitHub 的自由授權的程式碼資料。研究者使用了質量過濾器，並選定了基於文件數（document count）的前 25 種程式語言。在走完所有過濾步驟後，僅程式碼和 markup 子集的規模為 139B tokens。
Markdown 資料。研究者單獨處理了 mark-up 風格的語言，比如 Markdown、CSS 和 HTML。走完所有過濾步驟後，markup 子集的規模為 180B tokens。
合成程式碼資料。為了對程式碼資料集進行消融測試，研究者使用了專門的合成生成程式碼資料集， 包含已經正式驗證過的 Python 程式設計問題。他們將該資料集作為高質量程式碼資料來源，最終的合成資料集規模為 3.2B tokens。
相鄰程式碼資料。為了探索不同屬性的程式碼資料，研究者還使用了包含 GitHub 提交、jupyter notebooks、StackExchange threads 等輔助資料的程式碼資料。這類資料的規模為 21.4B tokens。
預訓練冷卻資料集。冷卻包含在預訓練最後階段提高更高質量資料集的權重。對此，研究者選擇了包含高質量文字、數學、程式碼和指令型文字資料集的預訓練冷卻混合。
評估
本文的目標是系統地理解程式碼對通用任務效能的影響，因此使用了一個廣泛的評估元件，涵蓋了包含程式碼生成在內的多樣下游任務。
為此，研究者在包含 1）世界知識、2）自然語言推理和 3）程式碼效能的基準上對模型進行了評估。此外，他們還報告了透過 LLM-as-a-judge 評估的勝率（win-rates）。
下表 1 展示了完整的評估元件以及相應的任務、資料集、指標。
研究者對不同規模的模型（從 470M 到 2.8B 引數）展開了效能評估。由於最小規模的模型能力有限，因此為了保證公平比較，他們只比較了所有模型均能達到隨機以上效能的基準。
除了特定於任務的判別式效能，研究者評估了使用 LLM-as-a-judge 勝率的生成式效能。
訓練與模型細節
如上文所說，對於預訓練模型，研究者使用了 470M 到 2.8B 引數的 decoder-only 自迴歸 Transformer 模型，它們按照標準語言建模目標來訓練。
具體來講，研究者使用了並行注意力層、SwiGLU 啟用、沒有偏差的密集層和包含 256000 個詞彙的位元組對編碼（BPE）tokenizer。所有模型使用 AdamW 最佳化器進行預訓練，批大小為 512，餘弦學習率排程器的預熱步為 1325，最大序列長度為 8192。
在基礎設施方面，研究者使用 TPU v5e 晶片進行訓練和評估。所有模型在訓練中使用了 FAX 框架。為了嚴格進行消融評估，研究者總共預訓練了 64 個模型。每次預訓練執行使用 200B tokens，470M 引數模型用了 4736 TPU 晶片時，2.8B 引數模型用了 13824 TPU 晶片時。每次冷卻執行使用了 40B tokens，470M 引數模型用了 1024 TPU 晶片時。
實驗結果
該研究展開了系統的實驗，探究了以下幾方面的影響：
使用程式碼預訓練模型初始化 LLM
模型規模
預訓練資料中程式碼的不同比例
程式碼資料的質量和屬性
預訓練冷卻中的程式碼資料
為了探究使用具有大量程式碼資料的 LM 作為初始化是否可以提高模型效能，該研究針對不同的預訓練模型初始化進行了實驗。如圖 2 所示，使用 100% 程式碼預訓練模型（code→text）進行初始化能讓模型在自然語言 (NL) 推理基準上獲得最佳效能，緊隨其後的是 balanced→text 模型。
為了瞭解上述結果是否可以遷移到更大的模型，該研究以 470M 模型相同的 token 預算，訓練了 2.8B 引數模型。下圖顯示了 2.8B 模型與 470M 模型的比較結果。
該研究探究了預訓練中程式碼資料比例對不同任務模型效能的影響，觀察到隨著預訓練程式碼資料比例的增加，程式碼任務的效能呈線性提高，而對於 NL 推理任務和世界知識任務則存在效益最明顯的最佳程式碼資料比例範圍。
如圖 5 (a) 所示，在評估程式碼質量和程式碼構成的影響方面，該研究觀察到，包含不同的程式碼源和合成程式碼，都會導致自然語言效能的提高，但是，只有合成生成的程式碼才能提高程式碼效能。
如圖 5 (b) 所示，在 NL 推理任務和程式碼任務中，balanced+synth→text 比 balanced→text 分別實現了 2% 和 35% 的相對改進。這進一步證實，即使是一小部分的高質量程式碼資料，也可以提高程式碼和非程式碼任務的效能。
如圖 6 所示，該研究發現：在預訓練冷卻中包含程式碼資料，模型的NL推理效能增加 3.6%，世界知識效能增加 10.1%，程式碼效能增加 20%。
如圖 7 所示，正如預期的那樣，冷卻對以勝率衡量的生成效能有重大影響。
採用不同預訓練方案的模型的效能比較結果如表 2 所示，Balanced→Text 實現了最佳的 NL 效能，而 Balanced-only 在程式碼生成方面明顯表現更好。