AIによるコード修正機能の改善：DatabricksのQuick Fix開発体験記

私たち人間は、バレエやボクシング（私がこの夏に挑戦した両方の活動です！）のような新しいことを試行錯誤を通じて学びます。私たちは物事を試し、失敗から学び、指導に耳を傾けることで成長していきます。私はこのフィードバックループをよく知っています。私のインターンプロジェクトの一部は、Databricksがトップレベルのコードアシスタントを構築するための取り組みの一環として、ユーザーに表示するより良いコード修正を識別する報酬モデルを教えることでした。

しかし、私のモデルだけが試行錯誤を通じて学んでいたわけではありません。私がモデルに良いコード修正と悪いコード修正を区別する方法を教えている間、私は堅牢なコードを書く方法、影響力のある製品のためのレイテンシと品質の問題をバランスする方法、大きなチームに明確にコミュニケーションを取る方法、そして何よりも、途中で楽しむ方法を学びました。

Databricksアシスタントクイックフィックス

コードを書いて実行しようとした際に厄介なエラーに遭遇したことがあるなら、Quick Fixの価値を実感していただけるでしょう。DatabricksのノートブックとSQLエディタに組み込まれたQuick Fixは、1-3秒で生成できる高信頼性の修正に設計されています。構文エラー、列名のスペルミス、簡単なランタイムエラーなどに最適です。クイックフィックスがトリガーされると、コードとエラーメッセージを取り、LLMを使用してエラーを解決するためのターゲットフィックスを生成します。

私のインターンプロジェクトはどのような問題に取り組んだのでしょうか？

Quick Fixはすでに存在し、Databricksのユーザーがコードを修正するのを助けていましたが、それをさらに良くする方法はたくさんありました！例えば、コード修正を生成し、それが構文規則を満たしているかの基本的なチェックを行った後、ユーザーに表示する修正が最も関連性があり、正確であることをどのように保証できるでしょうか？そこでbest-of-kサンプリングの導入です。複数の修正提案を生成し、報酬モデルを使用して最良のものを選択する手法です。

私のプロジェクト構造

私のプロジェクトはバックエンドの実装と研究実験の混合で、楽しく学びがいのあるものでした。

複数の提案を生成する

まず、Quick Fixのバックエンドフローを拡張して、異なるプロンプトとコンテキストを使用して並行して多様な提案を生成しました。私は、提案の質と多様性を最大化するために、思考の連鎖的な推論、予測された出力の推論、システムプロンプトのバリエーション、選択的なデータベースコンテキストの追加などの技術を試しました。追加の推論を含む提案生成により品質指標は向上しましたが、レイテンシの増加というトレードオフも発生することがわかりました。

ユーザーに表示する最適な修正提案を選択する

複数の提案が生成された後、返す最善のものを選択する必要があります。私は、最も頻繁に提案される修正をユーザーに提示するという、単純な多数決の基準を実装することから始めました。これは、より一般的に生成される解決策が最も効果的である可能性が高いという原則に基づいています。このベースライン手法はオフライン評価では良好な性能を示しましたが、オンラインでのA/Bテストでは既存実装に対して有意な改善が見られなかったため、本番環境への展開は見送りました。

さらに、私は最も有望な提案をランク付けし選択するための報酬モデルを開発しました。ユーザーが受け入れ、成功裏に実行する修正を予測するためにモデルを訓練しました。私たちは古典的な機械学習のアプローチ（LightGBMパッケージを使用したロジスティック回帰と勾配ブーステッド決定木）と微調整されたLLMを使用しました。

結果と影響

驚いたことに、候補修正のユーザー受け入れと実行成功を予測するタスクにおいて、古典的なモデルはオフライン評価で微調整されたLLMと同等のパフォーマンスを発揮しました。特に決定木モデルは、Quick Fixが扱うエラーの種類に対して「正しい」と見えるコード編集が実際には正しい傾向があるため、うまく機能したかもしれません。特に有効だった特徴量は、元のコード行と生成された修正案との類似度、およびエラーの種類でした。

このパフォーマンスを考慮に入れて、決定木（LightGBM）モデルを本番環境にデプロイすることにしました。LightGBMモデルを支持するもう一つの要素は、微調整されたLLMと比較して大幅に高速な推論時間でした。Quick Fixにとって速度は重要な要素です。ユーザーが手動でコードを編集する前に提案を表示する必要があり、レイテンシの増加は結果的に修正されるエラー数の減少につながります。LightGBMモデルの小さなサイズは、リソース効率が高く、製品化が容易になりました。モデルとインフラストラクチャの最適化により、平均推論時間を約100分の1に短縮することができました。

最良のkアプローチと報酬モデルを実装することで、私たちは内部受け入れ率を上げ、ユーザーの品質を向上させることができました。また、レイテンシを元の実装と同等の許容範囲内に維持することができました。

Databricks Assistantについて詳しく知りたい場合は、ランディングページまたはAssistant Quick Fix Announcementをご覧ください。

私のインターンシップ体験

Databricksの文化を体験

このインターンシップは、高影響力の製品に直接貢献する素晴らしい経験でした。Databricksの文化が、システムと製品の品質に対する高い基準を維持しながら、強い行動への偏見を奨励する方法について、私は直接的な洞察を得ました。

最初から、皆がいかに賢くて謙虚であるかに気づきました。その印象は時間と共に強くなり、チームがどれほど真剣にサポートしているかを見るたびに強まりました。非常に経験豊富なエンジニアでさえも、技術的な課題を話し合ったり、思慮深いフィードバックを提供したり、過去のアプローチや学びを共有したりすることで、私が成功するのを手助けしてくれました。

特に、私のメンターであるWill Tipton、マネージャーのPhil EichmannとShanshan Zheng、非公式のメンターであるRishabh SinghとMatt Hayes、エディター/アシスタントチーム、Applied AIチーム、MosaicMLの皆さんに感謝したいと思います。彼らから学んだ貴重なスキルと人生の教訓は、私がこれからのキャリアで持ち続けるものです。

他の素晴らしいインターンたち！

最後に、他のインターンたちと知り合うのがとても楽しかったです！採用チームは私たちがつながるのを助けるために、多くの楽しいイベントを開催しました。私のお気に入りの一つはインターンオリンピック（下に写真）でした。昼食を共有したり、地元のワークアウトクラスを試したり、カラオケで誕生日を祝ったりするなど、インターングループが仕事の中でも外でも支え合っていることを本当に感謝しています。

インターンオリンピック！チーム2を応援しよう！

ボクシングを一緒に試してくれた他のインターンに感謝します！

この夏、私が学んだことは、本物の問題と本物の制約を解決しながら、特に賢くて、意欲的で、支援的な人々に囲まれているときに最高の学習が行われるということです。インターンシップで最も報われた部分は、モデルの訓練を完了させたり、チームに興味深い結果を発表したりすることだけではなく、より良い質問をする能力が向上し、設計のトレードオフを理解し、Databricksのような広く使用されているプラットフォームで具体的な機能を最初から最後まで提供することを実現したことでした。

先端的なプロジェクトに取り組み、素晴らしいチームメイトと一緒に働きたいのであれば、Databricksでの仕事を申し込むことをお勧めします！Databricks Careers pageを訪れて、会社全体での求人情報を詳しく知ることができます。