DatabricksのインターンがAIにコードのクイック修正を強化する方法を教える

人間が反復練習と間違いから学ぶことでスキルを磨くように、人工知能モデルも現在、コード修復のような複雑な領域でその技術を完成させるために訓練されています。この基本的なフィードバックループは、Databricksの最近の夏季インターンシッププロジェクトの中心にあり、その目的はCode Assistantの機能を向上させることでした。この取り組みは、最適​​なコード修正を識別し優先順位付けするために、特殊な「報酬モデル」を教えることに焦点を当てていました。これは、洗練された開発者ツールを構築する上で極めて重要なステップです。

中核となるイノベーションは、Databricksの「Quick Fix」機能であり、同社のNotebooksおよびSQLエディタにシームレスに統合されています。迅速かつ高信頼性の解決策のために設計されたQuick Fixは、構文エラー、スペルミスのある列名、単純なランタイム問題など、一般的な問題を対象としています。開発者がエラーに遭遇すると、Quick Fixが作動し、問題のあるコードとその付随するエラーメッセージを分析します。その後、大規模言語モデル（LLM）を活用して、通常1〜3秒以内にターゲットとなるソリューションを迅速に生成します。

Quick Fixはすでに貴重な支援を提供していましたが、課題はその精度を向上させることにありました。修正が生成され、基本的な構文チェックを通過した後でも、システムはそれがユーザーにとって最も関連性が高く正確なソリューションであることをどのように保証できるでしょうか？その答えは、「best-of-kサンプリング」として知られる技術を通じて現れました。このアプローチでは、複数の潜在的な修正提案を生成し、その後、洗練された報酬モデルを使用して単一の最良のオプションを評価および選択します。

このプロジェクトは、バックエンドエンジニアリングと実験的研究の両方を網羅していました。最初のフェーズでは、Quick Fixシステムを拡張して、多様な提案を並行して生成することに焦点を当てました。これには、「思考の連鎖」推論、予測出力推論、システムプロンプトのバリエーションなどの技術、および選択的なデータベースコンテキストを含む、さまざまなプロンプトとコンテキスト情報の実験が含まれていました。これらの方法論は提案の品質と多様性を高めましたが、処理遅延もいくらか導入し、出力品質と運用速度の間の重要なトレードオフを浮き彫りにしました。

複数の提案が生成された後、次のステップはユーザーに提示する最も有望なものを特定することでした。初期の試みは単純な多数決システムでしたが、これは単独のテストではうまく機能したものの、実際のユーザーA/Bテストでは significant improvements をもたらさなかったため、展開されませんでした。より効果的な解決策は、専用の報酬モデルを開発および訓練することでした。これらのモデルは、ユーザーが最終的に受け入れ、正常に実行する提案された修正を予測するように設計されました。興味深いことに、開発プロセスでは、ロジスティック回帰や勾配ブースティング決定木（特にLightGBMパッケージを使用）などの伝統的な機械学習手法と、より現代的なファインチューニングされた大規模言語モデルの両方が利用されました。

評価から注目すべき発見がありました。ユーザーの受け入れと実行の成功を予測するという特定のタスクにおいて、オフライン評価では、古典的な機械学習モデルが、よりリソースを消費するファインチューニングされたLLMと同等の性能を示しました。特に決定木モデルは非常に効果的であることが証明されました。その成功は、Quick Fixが対処するエラーの種類について、正しく見えるコード編集がしばしば実際に正しいという事実に起因しています。その決定を通知した主要な特徴には、元のコード行と生成された修正との類似性、および遭遇した特定のエラータイプが含まれていました。

その比較可能な性能、そして決定的に、その著しく速い推論時間により、LightGBM決定木モデルは最終的に本番環境への展開のために選択されました。Quick Fixにとって速度は最優先事項です。開発者が手動で介入する前に、提案はほぼ瞬時に表示される必要があります。なぜなら、どのような遅延も迅速に解決できるエラーの数を減らすからです。LightGBMモデルのコンパクトなサイズも、リソース効率が非常に高く、既存のインフラストラクチャへの統合が容易になりました。モデルとインフラストラクチャの最適化を組み合わせることで、平均推論時間は劇的に約100倍削減されました。このbest-of-kアプローチの実装は、報酬モデルと相まって、Databricksの内部受け入れ率の測定可能な増加につながり、システムレイテンシーを許容範囲内に保ちながら、より高い品質とよりシームレスなユーザーエクスペリエンスに直接貢献しました。