Mellum：編輯器中快速、微型的程式碼完成模型

Mellum：編輯器中快速、微型的程式碼完成模型

整合開發環境（IDEs）長期以來都受益於自動完成功能的強大與精準。然而，現在出現了一位新的競爭者：人工智慧（AI）。AI 整合到 IDE 中已變得越來越普及，促使人們重新評估傳統工具。雖然這些較舊的方法仍然有效，但它們可能無法再提供開發人員現在所期望的尖端協助。

本文改編自 JetBrains 的一篇部落格文章，深入探討了 Mellum 的訓練與能力，Mellum 是一個支援雲端程式碼自動完成的模型。JetBrains 以其專為程式設計師量身打造的 IDE 和程式碼編輯器而聞名，他們開發了 Mellum 以提供快速且高效的程式碼自動完成功能。Mellum 的設計使其能夠在本地設備上有效運行，使其成為同類產品中最快且最精簡的模型之一。概念上相似的專案是微軟的 Phi。

解決顯而易見的挑戰

JetBrains 工程師在追求 AI 驅動的自動完成功能時，遇到了幾個關鍵挑戰：

• 速度與成本： 傳統的聊天模型因其高昂的計算成本和緩慢的反應時間而被證明是不切實際的。這些模型也缺乏對程式碼特定技術的了解，例如填空中間（FIM）或 token 修復。
• 輸出格式： 旗艦聊天模型經常產生格式不一致的資料，使得解析回應並將其無縫整合到編輯器中變得困難。
• 資料出處： 確定訓練資料的來源並減輕潛在的版權侵權問題，構成了一個重大的障礙。

Mellum：概述

JetBrains 的開發團隊意識到，創建自己的模型是最佳方法。他們的目標是設計一個在品質、推論成本和延遲之間取得平衡，同時在具有明確出處的資料上進行訓練的模型。初步研究表明，一個大約有 40 億個參數的模型可以為廣泛的場景和使用者提供稱職的自動完成功能。此外，透過僅在程式碼上訓練模型，他們可以建立一個特殊的 token 詞彙，擺脫不相關的資料。

Mellum 模型的訓練過程包括三個不同的階段，每個階段都貢獻了新的知識並提高了生成程式碼的品質。初始階段包括對大量個別檔案進行基本預訓練。第二階段包括使用較小的一組專業範例來完善模型。具有 AI 回饋的強化學習（RLAIF）在第三階段被採用，以使模型適應 IDE 特定的特性並消除不需要的輸出。

預訓練

為了避免與不明確的資料來源相關的潛在問題，該模型是從頭開始訓練的，需要全面介紹多種語言、程式設計語法、模式和核心概念。

資料集

預訓練的主要資料來源是 TheStack。此資料集確保資料在法律上合理且在實用上受益。

預訓練過程

在預訓練期間，資料集被多次取樣，以達到約 3 兆個 token。使用了 8192 個 token 的上下文視窗，資料集被分成此大小的片段。填空中間（FIM）轉換被應用於每個片段中一半的檔案，鼓勵模型考慮前面的和後面的程式碼。此技術密切模仿了真實世界的程式碼生成場景。

預訓練階段是在一個由 16 個節點組成的叢集上進行的，每個節點都配備了八個 H100 GPU。此階段大約需要 15 天才能完成，從而產生了具有 40 億個參數的 Mellum-base 模型。

預訓練創建了一個通用的程式碼自動完成模型，具有許多程式設計語言的廣泛知識。然而，在此階段，模型僅被訓練來預測隨機選擇的檔案片段中的下一個 token。如果沒有額外的上下文，模型就缺乏對程式碼結構的了解，也沒有機制來確定何時停止生成程式碼。

微調階段旨在解決這些限制。

上下文感知微調

增強的填空中間

與預訓練不同，在預訓練中，程式碼片段是隨機選擇以進行預測，而微調則側重於以更有意義的方式分割程式碼，教導模型提取「在野外」發生的程式碼片段。

專業範例

實際上，程式碼自動完成需要了解周圍的檔案和更廣泛的上下文，可能包括整個專案。

對於資料預處理，該公司啟動了一個內部專案，代號為 Code Engine：一個跨平台 SDK 和一組控制台實用程式，旨在直接從普通檔案建立上下文，而無需完整的專案索引。此 SDK 被部署在一個內部 MapReduce 叢集上，並用於處理數千個公共儲存庫，從而在合理的時程內產生了許多有用的訓練範例。

找到正確的演算法需要一些嘗試和錯誤。

針對特定語言的微調

小型模型可以從針對特定語言的專業化中獲益匪淺。雖然基礎模型是在 80 多種語言上訓練的，但大多數使用者通常只使用一到兩種。為了解決這個問題，JetBrains 創建了多個專業模型：

• mellum-all： 支援 JetBrains IDE 中提供的大多數語言和方言，但自動完成品質低於專業模型。
• mellum-python： 專門用於 Python 和 Jupyter。
• mellum-kotlin： 專門用於 Java 和 Kotlin。
• mellum-web： 專門用於 Web 技術。

最後一步：RLAIF

最後，必須解決學習目標與使用者期望不符的情況。一個額外的訓練階段，RLAIF — 具有 AI 回饋的強化學習，被用於解決這些問題。

該模型從使用者互動中學習，並了解如何更好地反映使用者偏好。

這種方法不僅提高了整體品質分數，還減少了令人煩惱的生成偽影的數量。

Mellum 有多好？

該模型的效能非常好，因為它的尺寸很小。以下是如何評估它的：

• 首先，該模型在一個內部基準測試上進行了評估，該基準測試的代號為「JetBrains BigCode」。
• 然後，它在眾所周知的公共基準測試（例如 SAFIM）上進行了測試。
• 最後，收集了功能的使用統計資料，並計算了使用者指標。

離線評估

收集資料是一項複雜的任務，但是創建一個好的指標來比較原始建議與神經網路提出的新建議更具挑戰性。我們進行了一項小型研究，最終確定了兩個主要指標的組合：

EM：

• 完全匹配是一個非常流行的想法。
• 如果完成的第一行與原始行的第一行匹配（經過最小的預處理），則認為預測是好的。

KK：

• 該指標以其作者的名字命名。
• 原始建議的行數除以提議完成的行數。

JetBrains BigCode

該模型是根據使用內部 JetBrains BigCode 工具獲得的基準資料集進行評估的。

透過保持對我們資料集的完全控制，而不是依賴公共基準測試，可以可靠地評估各種程式設計風格和實踐的模型品質。

我們的 JetBrains BigCode 評估結果顯示，品質與流行的模型相當，但是 Mellum 更小，效率更高。

單行建議的品質（EM 指標）

公共基準測試

該模型不僅在內部資料集上進行了評估，而且在各種公共基準測試上進行了評估，例如多語言基準測試 SAFIM（語法感知填空中間）。

線上評估

主要指標稱為已完成程式碼的比率（RoCC）。它被定義為使用程式碼自動完成編寫的程式碼字元數與編輯器中程式碼總量的比率。

另一個重要的指標是接受率（AR），它是根據接受的建議數量除以顯示的所有建議數量來計算的。

這是一段複雜的旅程，但是 JetBrains 的專家們以莊嚴的態度完成了它。最後，獲得了一個通用模型和幾個專業模型，這些模型可透過 JetBrains AI 平台獲得。它們現在已在 JetBrains AI Assistant 中成功運行。

下一步是什麼？

• JetBrains 工程師目前正在開發用於 Web 開發語言的模型。它可能在不久的將來公開發布。
• 有計劃同時增加參數數量和資料多樣性。程式設計中有許多不同的任務 - Mellum 也將能夠執行它們。服務效能仍然是一個關鍵指標，因此模型的擴展將在合理的範圍內。