Zeta2 の構築：本番環境で編集予測モデルを訓練する

Zed の編集予測 AI「Zeta2」が本番環境で高品質な学習データを生成する仕組み

Zed エディタの編集予測責任者、ベン・クンケル氏が解説した「Zeta2」というモデルは、単なるコード補完を超え、開発者の次のアクションを即座に予測するための小規模専門モデルです。彼らが本番環境でこのモデルを訓練する際に見せたのは、膨大なノイズを含むユーザー行動データを、いかにして高品質な教師データへと変換するかという、実用的かつ画期的なアプローチでした。

編集予測の定義と「即座に動く」要件

Zed の編集予測機能は、カーソル周辺のコード、直近の編集履歴、型定義、そしてエラー情報などを入力として受け取り、「次にユーザーが何をするか」を予測する仕組みです。このタスクの最大の特徴は、すべてのキー入力ごとに実行される必要がある点にあります。

すべてのキー入力ごとに即座に動作するためには、小規模で特化したモデルが不可欠です。

そのため、汎用的な大規模言語モデル（LLM）ではなく、この特定のタスクのみを処理するためにファインチューニングされた軽量モデル「Zeta2」が開発されました。このモデルは、複雑な推論よりも、高速なレスポンスと高い精度のバランスが求められます。

最先端モデルによる「蒸留」と修復プロセス

Zeta2 の訓練には、「蒸留（Distillation）」と呼ばれる手法が採用されています。これは、すでに高度に学習された最先端の教師モデル（Teacher Model）を用いて予測を生成し、その出力を Zeta2（Student Model）の正解データとして転送するプロセスです。

しかし、最先端モデルに単純に入力を与えても、必ずしも良い結果が返ってくるわけではありません。そこで Zed チームは、以下の「修復ステップ」を設けてデータの質を担保しています。

1. 静的評価: 直前の入力を取り消していないか、編集可能な領域の境界を無視していないかなど、ヒューリスティックを用いて予測を検証します。
2. 自動修復: 検証で違反が見つかったケースは、再度教師モデルに「失敗しました。修正できますか？」と問い合わせて修正させます。

このようにして得られた高品質なデータのみが、Zeta2 の訓練用ラベルとして利用されます。このプロセスは実験間で共通するためキャッシュ可能であり、効率的に学習データを生成しています。

「セットル状態」を活用したノイズ除去の知恵

本番環境から収集するユーザー行動データには、当然ながらノイズが含まれます。例えば、ユーザーがコードを書いている途中で考えが変わったり、別のツールで書き直したりする場合です。このような「中途半端な状態」を学習データにしてしまうと、モデルは不安定になります。

そこで Zed チームが目をつけたのが「セットル状態（Settled State）」です。これは、ユーザーがそのコードブロックに対して編集を完了し、一時的に操作を停止した直後の安定した状態を指します。

ユーザーが編集を止めるまで待ち、その領域のスナップショットを取得することで、ノイズの少ない高品質な学習サンプルを抽出できます。

ただし、この「セットル」を検出する際にも注意が必要です。単に 10 秒間操作が止まっただけでなく、ユーザーが本当に満足してその状態を維持しているかを判断する必要があります。また、完全に書き直されるようなケースもノイズとして除去する必要があります。

ノイズフィルタリングの具体的な手法

セットル状態の判定には、以下の指標が用いられています。

• Levenshtein 距離（編集距離）: セットルした後のコードと、教師モデルや Zeta2 が生成した予測との類似度を計算します。もしセットル後のコードが予測結果に非常に近ければ、それは「ノイズ」ではなく「意図的な編集」として扱います。
• コスト削減の工夫: 本来は教師モデルで 10 万回もリクエストを繰り返す必要がありますが、Zeta2 がすでに教師に近い精度を持っていることを利用し、学生モデル（Zeta2）のチェックポイントを使って同様のフィルタリングを行うことで、計算コストを実質ゼロに抑えています。

この手法により、「関数 add A を入力すると結果は B になる」といった明確なパターンや、学習データのカットオフを超えた新しい関数の記述など、モデルが学ぶべき「中間の興味深いセクション」のみを抽出できます。

評価と本番環境への展開

モデルの評価には、保持されたテストセット上でオフライン評価を行い、過学習を防ぎます。また、正解が一つとは限らないタスクであるため、教師モデルが生成した複数の異なる回答（3 つ）と比較し、そのいずれかに近い予測であれば「良い予測」として扱います。

本番環境への展開では、まず 15% のトラフィックでサンプリングし、受容率やレイテンシなどの指標を確認。問題なければ徐々に割合を上げていき、最終的に全ユーザーに提供されるモデルとして稼働させます。このように、Zed は実験データと実運用データをシームレスにつなげながら、モデルの精度を継続的に向上させています。

まとめ

Zeta2 の訓練プロセスは、単なる大量データの収集から「実環境での行動データをどうフィルタリングし、高価値な学習セットに変換するか」というパラダイムシフトを示しています。ノイズの多いユーザーデータの中から、編集が落ち着く直前の状態（セットル）や修復ステップを活用して質を高めるこのアプローチは、開発者支援ツールにおける AI 化の標準的なベストプラクティスとして、業界全体に大きな影響を与えるでしょう。