DatabricksIQ LLMの品質向上 - AIによるテーブル説明文の生成

最近、Unity CatalogのAI生成コメントをサポートする基礎となるアルゴリズムに大幅な改善を加えました。その結果を皆さんにお伝えできることを嬉しく思います。Databricks のデータインテリジェンスエンジンであるDatabricksIQを通じて、AIによって生成されたコメントは、顧客のUnity Catalogテーブルの新しいドキュメントの大部分をすでに生成しており、最近の機能強化は、この非常に人気のある機能をさらに強化します。

このブログでは、トレーニングデータの合成に更新されたオープンソースLLMを使用する方法、トレーニングデータのクリーニングにヒューリスティックフィルターを使用する方法、ファインチューニング用に更新されたベースモデルを使用する方法、および自動ベンチマークで利用される拡張評価セットを使用する方法について説明します。最小限の労力で、これらの変更により、オフライン ベンチマークで以前に導入されたモデルと比較して、優先率が2倍に増加しました。さらに広い意味では、この取り組みにより、 DatabricksIQはお客様の応用AIユースケースのあらゆる範囲に対応できるようになり、さらに強力になりました。

AI生成コメントが必要な理由

企業データにコメントやドキュメントを追加するのは、報われない作業ですが、組織のテーブルにドキュメントがまばらにしか存在しない場合、人間もAIエージェントも、データに関する質問に正確に答えるための適切なデータを見つけるのに苦労します。AIコメントは、生成AIの魔法によってテーブルや列に説明を追加する手動プロセスを自動化することで、この問題に対処します。

昨年の秋、私たちは2人のエンジニアが1か月かけて特注のLLMをトレーニングし、Unity Catalogのテーブル用のドキュメントを自動生成する問題に取り組んだ様子について書きました 。このモデルのタスクは、簡単に言えば、テーブルスキーマが提示されたときにテーブルの説明と列のコメントを生成することです。この以前の作業では、適度なサイズのMPT-7Bモデルを約3,600個の合成サンプルでファインチューニングし、次に62個のサンプルテーブルを使用した二重盲検評価でこのモデルをベンチマークしました。このコンテキストで、モデルをさらに改善するために何を変更したかを順に見ていきましょう。

トレーニングデータの改善

Mistral AIが専門家混合モデルを一般公開したのと同じ週に、私たちはMixtral 8x7B Instructを使用して合成データセットの生成を開始しました。 以前のアプローチと同様に、NAICSコードとDatabricksの内部分類法で定義されたユースケースから始めました。 次に、Mixtralを使用して、各ユースケースに対して複数のCREATE TABLEステートメントを生成し、再度Mixtralを使用して、生成された各スキーマのテーブルの説明と列のコメントを生成しました。 Mixtral 8x7Bは同規模の他のLLMよりも推論のレイテンシが低いため、わずか2時間足らずで9Kの合成トレーニングサンプルを生成できました。

この演習では、プロンプトを作成してLLMを使用して大量の合成データを生成し始めると、生成されたデータの一部が、スタイル的にも構文的にも適切な形式にならない可能性があることがわかりました。 たとえば、すべてのテーブルの説明に2 ～ 4文を含めるようにしたい場合、プロンプトでその旨を述べるだけでは、LLMが一貫して従うには不十分な可能性があります。 このため、生成中に、チャット形式でいくつかの例を提示し、Mixtralに希望する出力形式（スキーマを生成するときの有効なCREATE TABLEステートメントなど）を教えました。 この手法を使用して、Mixtralの命令チューニングを活用し、一貫性のある例を含むデータセットを作成しました。

フィルタリングの改善

少数のショットの例や命令のチューニングを行っても、出力は望ましくないことがよくあります。 たとえば、生成されるスキーマの一部には、現実的または教育的な十分な列がない場合があります。 説明では、列の名前を要約するのではなく、単に逆流させるだけであったり、微調整されたモデルでコピーしたくない双曲線的な言葉が含まれている可能性があります。 以下は、実際のフィルターの例です。

次の生成されたスキーマは、`CREATE TABLE …` 構文に従っていないため拒否されました。

CREATE TYPE agritech.weather.windDirection AS ENUM ('N', 'NNE', 'NE', 'ENE', 'E', 'ESE', 'SE', 'SSE', 'S', 'SSW', 'SW', 'WSW', 'W', 'WNW', 'NW', 'NNW')

設備のメンテナンスに関するテーブルに対する次の合成説明は、テーブルの列をリストアップしているだけなので却下されました。 これは、トレーニングデータのスタイル コンテンツに対するフィルターです。指定したスキーマをそのまま繰り返すモデルは望ましくありません。

「EquipmentRepair」テーブルは、すべての機器の修理記録を追跡し、修理を依頼した顧客、問題の機器モデル、修理日、修理を実施した技術者、修理された機器の状態、修理の待ち時間、修理費用、修理体験に関する顧客のレビューを詳細に記録します。 このデータを使用して、修理コスト、機器の信頼性、顧客満足度の傾向を特定できます。 また、機器のメンテナンススケジュールを管理したり、特定の機器モデルで最も一般的に報告されている問題を見つけたりするのにも役立ちます。

合成データセットの全体的な品質を向上させるために、トレーニングに望ましくないスキーマ、説明、列コメントを削除するヒューリスティックフィルターをいくつか定義しました。 このプロセスの後、元の9Kサンプルのうち約7Kが残りました。

ファインチューニング

テーブルの説明と列のコメントの両方について7Kサンプルのフィルター処理された合成データセットを使用して、以前の本番運用モデルMistral（MPT-7B）とほぼ同じサイズとレイテンシを持つMistral-7B v1 Instructモデルをファインチューニングしました。 最終的に、MPT-7B ではなく Mistral-7B を選択しました。これは、当時チャットボット分野で最高ランクの 7B パラメータ モデルであったためです。トレーニングデータセットで2つのエポックにわたって ドメイン効率の高いファインチューニングを実行しました。完了するまでに約 45 分かかり、モデルを評価する準備が整いました。

モデル評価の改善

以前のアプローチでは、合成（以前は表示されていなかった）テーブルと、独自のDatabricksワークスペースからキュレートした実際のテーブルが混在する62個のスキーマの評価セットをブートストラップしました。 合成サンプルを評価に使用することは、データを見つけるのが難しい場合に最適ですが、多くの場合、実際の特性が欠けています。 この問題に対処するために、Databricksワークスペースで最も使用されているテーブル500個をサンプリングし、実際の使用状況をより正確に表す評価セットを作成しました。

この新しいモデルを発売前に検証するために、500個の評価テーブルすべてについて、以前のモデルと新しいモデルの両方を含むテーブルの説明を生成しました。 二重盲検評価フレームワークを使用して、生成された記述が評価者に好まれるかどうかを測定しました。 開発サイクル中に人間が500個のサンプルに注釈を付けるのはコストがかかるため、評価タスクを明確に定義するプロンプトを使用して、別のLLMを評価者として活用するようにフレームワークを更新しました。 自動評価には、注目すべき機能がいくつかあります。

• モデルが 2 つの出力から最適なものを選択するように要求するのではなく、同順位を許可しました。 私たち自身が評価タスクを実行したとき、2 つの出力があまりにも近すぎて判断できないことが時々あったため、モデルにもこの柔軟性を与えることが重要でした。
• モデルは、その答えの背後にある説明を提供する必要がありました。 この手法は、より正確な結果につながることが研究や業界で発見されています。 これらすべての説明を記録し、その動作を確認し、モデルが正しい推論に基づいて選択していることを確認する機会を与えました。
• 各サンプルに対して複数の評価を生成し、多数決を使用して最終的な評価を選択しました。

モデルの説明により、評価タスクを定義するプロンプトを改善できる箇所についての知見が得られ、人間の判断とより一致する評価につながりました。 アプローチの微調整が終わったら、2人の人間の評価者に100個の無作為評価サンプルを審査してもらいました。 人間の評価者が LLM 評価者と同等の好みを持つ可能性は低いですが、両方の評価者が新しいモデルを明確に好みます。 実際、人間の評価者は、以前の本番運用モデルの出力よりも約 2 倍頻繁に新しいモデルの出力を好みました。

パフォーマンス上の理由からLLM評価のみを使用して、列コメントを生成するタスクに対しても同じ評価プロセスを繰り返しました。 当社のLLM評価者は、以前のモデルよりも新しいモデルを3倍近く好んで評価しました。

新しいLLMの導入

オフライン評価を分析した後、新しいモデルをデプロイする段階になりました。以前と同じ手順に従って、モデルをUnity Catalog（UC）に登録し、 Delta Sharingを活用してすべてのリージョンで公開し、Databricksの最適化されたLLMサービスを通じて提供しました。

主要なポイント

このブログで説明した手法により、DatabricksIQのAIによるテーブルと列のコメントの自動生成機能が大幅に向上しました。 これを実現するために、オープンソースのLLMを活用して大規模なトレーニング データセットを作成しました。 次に、人間の直感と経験を活用し、ヒューリスティックフィルターを使用してデータセットをクリーンアップしました。 最後に、現在のモデルに代わる別の小さなLLMをファインチューニングしました。 現実的なサンプルを多数含めて評価セットを改善することで、オフライン環境での成果が顧客にとってより高品質な結果につながるという自信が深まりました。

これらの学習が、オープンソースLLMを活用して具体的なタスクに取り組みたいと考えているすべての人にとって役立つことを願っています。プロセスの各ステップには注意と分析が必要ですが、Databricks基盤モデルAPIや最適化された LLM サービスなどのツールを使用すると、LLMを活用することがこれまでになく簡単になりました。

AI 生成コメントを使用する準備はできていますか?

AIによって生成されたコメントにアクセスするには、カタログエクスプローラーを開き、Unity Catalogによって管理されているテーブルを選択します。 そこから、AIコメントを生成して承諾するオプションが表示されます。

AI によって生成されたコメントアイコンが表示されない場合は、管理者がこちらに記載されている手順に従って、DatabricksアカウントとワークスペースでDatabricksIQ機能を有効にすることができます。