Databricks は最近、データ作業専用に構築された自律型 AI パートナーである Genie Code を発表しました。Genie Code は Databricks Assistant を置き換え、複数のエージェントを統合し、新しい統合ポイントと機能を提供します。Unity Catalog との深い統合により、テーブル、列、リネージ、メトリクスビュー、ビジネス定義（セマンティクス）を理解します。このコンテキスト認識により、Genie Code は汎用的なチャットボットよりもデータ実務者にとってはるかに有用です。
Genie Code が「チャーン予測モデルを構築する」などの従来の ML タスク用のノートブックを生成するとき、適切な Python ライブラリのインストール、データの探索と前処理、トレーニング、チューニング、モデルの登録とデプロイ、パフォーマンス評価を含む、プロダクション対応のワークフローが生成されることを期待します。また、各ステップが真にデータに基づいていること、例えば Genie Code が二値分類問題における不均衡クラスが劇的に異なるワークフローと成功指標をもたらすことを理解していることを期待します。
Genie Code が Databricks ネイティブのベストプラクティスに一貫して従い、例えば交差検証のスキップ、データ漏洩の見落とし、不適切なデータ補完などを回避するために、生成されたコードが実際に良いものであるかどうかを判断する厳密な方法が必要でした。生成されるノートブックは顧客が解決しようとしている問題に大きく依存し、顧客によって問題は大きく異なるため、これは非常に自明ではない問題です。
この投稿では、Genie Code の従来の ML 機能のための評価パイプラインを構築した方法と、MLflow の新しいオープンソースアライメントフレームワークである MemAlign を使用して、LLM 評価者と人間の専門家の間に見られた大きなギャップを埋めた方法について説明します。改善された評価者は、見逃していたであろう Genie Code の ML ガイダンスのギャップを特定し修正するのに役立ちました。
評価フレームワークの構築
堅牢な評価フレームワークは以下のために必要です：
- ヒルクライミング：プロンプト、ツール、スキル、アーキテクチャの変更が出力に与える影響を定量化する。
- 回帰の防止：「モデルトレーニング」の改善が誤って「データ探索」を劣化させないようにする。
- ベンチマーク：異なる基盤モデル（LLM バックエンド）がノートブックの品質にどのように影響するかを測定する。
- CI：基盤となるエージェントループの変更が最終的な ML タスクにどのように波及するかを監視する。
従来の ML ノートブックの評価は、コード品質、ML のベストプラクティス、データ駆動型の適応/調整を網羅するため、最も複雑な評価タスクの1つです。ML ノートブックの評価のような広範で複雑なタスクを処理するために、LLM を評価者として使用します。人間によって良いノートブックとは何かを教えられた LLM「専門家」です。9つの評価者を作成し、ほとんどの ML ワークフローに現れる9つの次元に沿って ML ノートブックを評価するようにプロンプトを与えました：ライブラリインストール、探索的データ分析、データ補完、特徴量エンジニアリング、モデルトレーニング、モデルサービング、メトリクス評価、MLflow ログ、セル編成。
各次元について、1から3のスコアを割り当て、0は「該当なし」とするスコアリングルーブリック（人間の評価者と LLM 評価者の間で再利用）を作成しました：
- 3（良い）：ノートブックがその次元で高い基準を満たしており、ベストプラクティスを示し、期待される範囲をカバーし、エッジケースを適切に処理している。
- 2（平均）：許容できるがギャップがある。基本は存在するが、経験豊富な実務者が期待する洗練さが欠けている。
- 1（悪い）：根本的な問題。重要なステップが欠落している、誤っている、または誤った結論につながる方法で適用されている。
- N/A（該当なし）：その次元はこのプロンプトに適用できない（例えば、データセットに欠損値がない場合、データ補完次元は適用できない）。
評価者とともに、さまざまな ML タスク（分類、回帰、予測）をカバーし、異なるデータセットサイズ、ドメイン、複雑さにわたる評価テストケースのセットを維持しています。各テストケースには、Genie Code に指定されたデータセットで解決すべき ML タスクを伝えるユーザープロンプトが含まれています。評価ループは、各テストケースに対して Genie Code を使用して1つまたは複数のノートブックを生成し、次にすべての該当する次元に沿って各ノートブックをスコアリングすることで構成されます。
評価システムの評価
LLM 評価者を人間の代わりに使用して Genie Code の成果物を評価することで、本質的に1つの困難な問題を別の問題に置き換えました。既製の評価者は目の前のタスクに熟練しておらず、人間の評価と一致していません。私たちの問題は、LLM 評価者のスコアを人間の評価者のスコアと一致させることです。
LLM 評価者評価のための評価セットには、50の Genie Code 生成ノートブック（「テストケース」）が含まれており、人間の専門家がすべての該当する次元を評価し、スコアと短い正当化を提供してグラウンドトゥルースとしました。2つのスコアの間のグレーゾーンでは、評価者は自分の判断を表現することが許されましたが、スキーマはこれが稀であるように書かれていました。
人間と機械のアライメントの尺度は、各次元のスコア間の平均絶対誤差（MAE）です。結果はまちまちで、一部の次元では強いアライメント（4つの次元で MAE が0.3未満）が見られましたが、他の次元ではそうではありませんでした。最もアライメントが悪かった次元は、「データ補完」（MAE 0.68）、「モデルサービング」（0.62）、「メトリクス評価」（0.69）でした。これらの次元は ML ワークフローの重要なボトルネックであるため、高い誤差率は許容できませんでした。
MemAlign の登場
MemAlign（メモリアライメント）は、MLflow 2.16 以降で利用可能なオープンソースフレームワークであり、LLM 評価者を人間の好みに合わせるために使用します。これは推論時に2種類のメモリを使用して LLM のコンテキストを強化します：
- セマンティックメモリ：利用可能なレシピの長期記憶。これには、明示的なルール、事実、知識（例：「すべての ML 実験には交差検証と「評価」メトリクスを含めるべき」）が含まれます。
- エピソディックメモリ：過去の経験の短期記憶。これには、誤った判断から学んだ例（例：「前回 LLM はデータ漏洩を見逃したため誤ったスコアを付けた——今ではデータ漏洩をチェックすることを知っている」）が含まれます。
MemAlign を使用するには、ラベル付きサンプル（人間のスコアと説明）のセットを MLflow のベンチマーク評価セットに追加し、推論時に簡単な API 呼び出しでこれらのメモリを注入するだけです。MemAlign はこれらのメモリを使用して LLM 評価者の出力を自動的に調整し、人間の好みに合わせます。微調整やプロンプトエンジニアリングは一切不要です。
最もアライメントが悪かった3つの次元（データ補完、モデルサービング、メトリクス評価）に MemAlign を適用しました。各次元について、ラベル付きサンプルからランダムに約10個を選択し（50個から抽出）、アライメント効果をテストしました。結果は印象的でした。データ補完では MAE が0.68から0.18に低下（74%減）、モデルサービングでは0.62から0.07に低下（89%減）、メトリクス評価では0.69から0.14に低下（80%減）しました。この改善は非常に顕著で、LLM 評価者の評価を、回帰の検出やベンチマークのために人間の評価の代理として安全に使用できるようになりました。
その後の研究では、各メモリタイプがアライメントに与える影響をテストしました。このタスクでは両方のメモリが重要であることがわかりました。セマンティックメモリのみを使用した場合、データ補完の MAE は0.47、エピソディックメモリのみでは0.42でした。両方を組み合わせると MAE は0.18に低下しました。他の次元でも同様のパターンが見られました。
教訓
- 一貫したスコアリングは難しい：LLM 評価者と人間のずれは、多くの場合、プロンプトが曖昧すぎる、ルーブリックが不明確、またはコンテキストが不十分であることに起因します。専門家が正確なスコアと明確な説明を提供して、あいまいさを減らすようにしてください。
- 専門家間の一致：人間のアノテーターが互いに一致しない場合（データ探索回帰で見られたように）、アライメントは学習可能な一貫した信号を持ちません。
MemAlign は MLflow に同梱されており、わずか約50のラベル付きサンプルで機能しました。LLM 評価者が専門家と一致しない場合は、午後を使って試す価値があります。