Mellum：エディターに統合可能な高速・軽量コード補完モデル
統合開発環境（IDE）の世界では、長年にわたりオートコンプリート機能がその強力さと正確性を支えてきました。しかし、新たな挑戦者として人工知能（AI）が登場しました。AIのIDEへの統合はますます普及しており、従来のツールを再評価する必要が生じています。これらの古い手法はまだ機能しますが、開発者が現在期待する最先端の支援を提供できなくなっている可能性があります。
この記事はJetBrainsのブログ記事を元に、クラウドベースのオートコンプリートを支えるモデルであるMellumのトレーニングと機能について掘り下げています。プログラマー向けに調整されたIDEとコードエディターで有名なJetBrainsは、迅速かつ効率的なコードオートコンプリートを提供するためにMellumを開発しました。Mellumはローカルデバイス上で効果的に動作するように設計されており、同クラスで最も高速かつコンパクトなモデルの一つとして際立っています。概念的に類似したプロジェクトとしてはMicrosoft Phiが挙げられます。
明らかな課題への取り組み
JetBrainsのエンジニアは、AIを活用したオートコンプリートの追求において、いくつかの重要な課題に直面しました。
- 速度とコスト: 従来のチャットモデルは、高い計算コストと遅い応答時間のため、実用的ではありませんでした。これらのモデルは、fill-in-the-middle (FIM)やtoken healingのようなコード固有の技術に関する認識も欠いていました。
- 出力形式: 主要なチャットモデルは、一貫性のない形式でデータを生成することが多く、応答を解析してエディターにシームレスに統合することが困難でした。
- データの出所: トレーニングデータの出所を特定し、潜在的な著作権侵害の問題を軽減することは、大きなハードルとなりました。
Mellum：概要
JetBrainsの開発チームは、独自のモデルを作成することが最適なアプローチであることに気づきました。彼らの目標は、品質、推論コスト、レイテンシのバランスを取りながら、明確な出所を持つデータでトレーニングされたモデルを設計することでした。予備調査の結果、約40億のパラメーターを持つモデルが、幅広いシナリオやユーザーに対して有能なオートコンプリート機能を提供できることが示唆されました。さらに、モデルをコードのみでトレーニングすることにより、無関係なデータを含まないトークンの特殊な語彙を確立することができました。
Mellumモデルのトレーニングプロセスは、3つの異なる段階で構成されており、各段階が新しい知識を提供し、生成されたコードの品質を向上させます。最初の段階では、個々のファイルの大規模なコーパスに対して基本的な事前トレーニングが行われます。2番目の段階は、より小さな専門的な例のセットでモデルを改良することです。AIフィードバックを用いた強化学習（RLAIF）は、3番目の段階で採用され、モデルをIDE固有の特性に適応させ、不要な出力を排除します。
事前トレーニング
曖昧なデータの出所に関連する潜在的な問題を回避するために、モデルはゼロからトレーニングされ、多数の言語、プログラミング構文、パターン、およびコアコンセプトを包括的に導入する必要がありました。
データセット
事前トレーニングの主なデータソースはTheStackでした。このデータセットは、データが法的に健全であり、実質的に有益であることを保証します。
事前トレーニングプロセス
事前トレーニング中、データセットは複数回サンプリングされ、約3兆のトークンに達しました。8192トークンのコンテキストウィンドウが使用され、データセットはこのサイズのフラグメントに分割されました。fill-in-the-middle (FIM)変換は、各フラグメントのファイルの半分に適用され、モデルが先行するコードと後続のコードの両方を考慮するように促しました。この手法は、現実世界のコード生成シナリオを非常に模倣しています。
事前トレーニングフェーズは、それぞれ8つのH100 GPUを搭載した16ノードのクラスターで実行されました。この段階の完了には約15日かかり、40億パラメーターのMellum-baseモデルが完成しました。
事前トレーニングにより、多くのプログラミング言語に関する広範な知識を備えた汎用的なコードオートコンプリートモデルが作成されます。ただし、この段階では、モデルはランダムに選択されたファイルセグメント内の次のトークンを予測するようにのみトレーニングされています。追加のコンテキストがないと、モデルはコード構造を認識せず、コードの生成をいつ停止するかを判断するメカニズムもありません。
ファインチューニング段階は、これらの制限に対処するように設計されています。
コンテキストを意識したファインチューニング
強化されたFill-in-the-Middle
予測のためにコードフラグメントがランダムに選択される事前トレーニングとは異なり、ファインチューニングでは、より意味のある方法でコードをセグメント化することに重点を置き、モデルに「実世界」で発生するコードフラグメントを抽出するように教えます。
特殊な例
実際には、コードオートコンプリートは、周囲のファイルやより広範なコンテキスト、場合によってはプロジェクト全体を理解する必要があります。
データの前処理のために、同社はCode Engineというコードネームの内部プロジェクトを開始しました。これは、プロジェクト全体のインデックス作成を必要とせずに、通常のファイルから直接コンテキストを構築するために開発されたクロスプラットフォームSDKと一連のコンソールユーティリティです。このSDKは内部のMapReduceクラスターにデプロイされ、数千もの公開リポジトリを処理するために使用され、妥当な時間枠でトレーニングに役立つ多くの例が生成されました。
正しいアルゴリズムを見つけるには、試行錯誤が必要でした。
特定の言語向けのチューニング
小さなモデルは、特定の言語に特化することで大きなメリットを得ることができます。ベースモデルは80以上の言語でトレーニングされていますが、ほとんどのユーザーは通常、1つまたは2つの言語のみを使用します。これに対処するために、JetBrainsは複数の特殊モデルを作成しました。
- mellum-all: JetBrains IDEで利用可能なほとんどの言語と方言をサポートしていますが、オートコンプリートの品質は特殊モデルよりも低くなっています。
- mellum-python: PythonとJupyterに特化しています。
- mellum-kotlin: JavaとKotlinに特化しています。
- mellum-web: Webテクノロジーに特化しています。
最終ステップ：RLAIF
最後に、学習目標がユーザーの期待と一致しないケースを解決する必要があります。追加のトレーニングフェーズであるRLAIF（AIフィードバックを用いた強化学習）を使用して、そのような問題を解決します。
モデルはユーザーのインタラクションから学習し、ユーザーの好みをより適切に反映する方法を理解します。
このアプローチは、全体的な品質スコアを向上させるだけでなく、迷惑な生成アーティファクトの数を減らします。
Mellumの性能
このモデルは、そのサイズに対して非常に優れた性能を発揮します。評価方法は次のとおりです。
- 最初に、モデルは「JetBrains BigCode」というコードネームの内部ベンチマークで評価されました。
- 次に、SAFIMのような有名な公開ベンチマークでテストされました。
- 最後に、機能の使用状況統計が収集され、ユーザーメトリックが計算されました。
オフライン評価
データを収集することは複雑なタスクですが、元の提案とニューラルネットワークによって提案された新しい提案を比較する優れたメトリックを作成することはさらに困難です。小規模な調査を実施し、最終的に2つの主要なメトリックの組み合わせに落ち着きました。
EM:
- Exact Matchは非常に人気のあるアイデアです。
- 最小限の前処理で、完了の最初の行が元の最初の行と一致する場合、予測は良好と見なされます。
KK:
- メトリックはその作者にちなんで名付けられました。
- 元の提案された行数を除算することによって完了した行数。
JetBrains BigCode
モデルは、内部JetBrains BigCodeツールを使用して取得されたベンチマークデータセットに対して評価されました。
公開ベンチマークに依存するのではなく、データセットを完全に制御することで、さまざまなコーディングスタイルとプラクティスに対するモデルの品質を確実に評価できます。
JetBrains BigCode評価の結果は、一般的なモデルと同等の品質を示していますが、Mellumはより小さく、より効率的です。
単一行提案の品質（EMメトリック）
公開ベンチマーク
モデルは、内部データセットだけでなく、多言語ベンチマークSAFIM（構文を意識したfill in the middle）などのさまざまな公開ベンチマークでも評価されました。
オンライン評価
主なメトリックは、完了したコードの比率（RoCC）と呼ばれます。これは、コードオートコンプリートを使用して記述されたコード文字の比率を、エディター内のコードの合計量で割ったものとして定義されます。
もう1つの重要なメトリックは、承諾率（AR）であり、表示されたすべての提案の数で割った承認された提案の数として計算されます。
これは複雑な道のりでしたが、JetBrainsのスペシャリストは立派にそれを完了しました。最終的に、1つの一般的なモデルといくつかの特殊モデルが取得され、JetBrains AIプラットフォームを通じて利用可能になりました。現在、JetBrains AI Assistantで正常に動作しています。
次に何が来るのでしょうか？
- JetBrainsのエンジニアは現在、Web開発言語用のモデルに取り組んでいます。近い将来、一般に公開される可能性があります。
- パラメーターの数とデータの多様性の両方を同時に増やす計画があります。コーディングにはさまざまなタスクがあり、Mellumもそれらを実行できるようになります。サービスのパフォーマンスは依然として重要なメトリックであるため、モデルの拡張は合理的な範囲内にとどまります。