LogSentinel: DatabricksがLLMを活用したPII検出とガバナンスにDatabricksを利用する方法

Databricksほどの規模になると、社内のログやデータセットは絶えず変化します。スキーマは進化し、新しい列が追加され、データのセマンティクスもdriftしていきます。このブログでは、Databricks社内でDatabricksをどのように活用し、プラットフォームが変化する中でPIIやその他の機密データに正しくラベル付けを行っているかを解説します。

このために、私たちは Databricks 上で動作する LLM を活用したデータ分類システムである LogSentinel を構築しました。このシステムは、スキーマ進化を追跡し、ラベル付けのdriftを検出し、高品質のラベルをガバナンスおよびセキュリティ制御にフィードします。当社では MLflow を使用して実験を追跡し、パフォーマンスを長期的に監視しています。また、顧客が同じアプローチからメリットを得られるように、LogSentinel の優れたアイデアを Databricks データ分類製品に再統合しています。

このシステムが重要な理由

このシステムは、プラットフォーム、データ、セキュリティの各チーム向けに、3つの具体的なビジネスレバーを動かすように設計されています。

• コンプライアンスサイクルの短縮：以前はアナリストが数週間を要していた定期的なレビュータスクが、人間が確認する前に列に事前にラベル付けとトリアージが行われるため、数時間で完了するようになりました。
• 運用リスクの低減：システムがラベル付けのドリフトやスキーマの変更を継続的に検出するため、不正確なタグや欠落したタグが付いた機密フィールドが気付かれないまま見過ごされる可能性が低くなります。
• ポリシー適用の強化: 信頼性の高いラベルによって、マスキング、アクセス制御、保持、所在地ルールが直接実行され、かつての「ベストエフォート ガバナンス」が実行可能なポリシーに変わります。

実際には、チームはドメインごとに特注の分類器を構築することなく、新しいテーブルを標準パイプラインに接続し、driftメトリクスと例外を監視し、システムにPIIと所在地の制約を適用させることができます。

システムアーキテクチャの概要

当社はDatabricks上にLLMを活用した列分類システムを構築しました。このシステムは、社内のデータ分類法を使用してテーブルに継続的にアノテーションを付け、ラベリングのドリフトを検出し、問題が発見された場合には修正チケットを発行します。システムに関わるさまざまなコンポーネントを以下に示します (MLFlowを使用して追跡および評価)。

• データ取り込み: さまざまなデータソース (Unity Catalogの列データ、ラベル分類データ、正解データを含む) の取り込み
• データ拡張：Databricks Vector SearchとAIコメント生成を使用したデータ拡張
• LLMオーケストレーション
• 階層型ラベリングシステム
• モデルのバージョン管理: 複数のモデルを並列実行
• ラベル予測: Mixture of Experts (MoE) アプローチを使用して最終的なラベルを予測
• チケット作成: 違反を検出し、JIRA チケットを生成します

エンドツーエンドのワークフローを下図に示します。

データの取り込み

注釈を付けるログタイプやデータセットごとに、すべての列から値をランダムにサンプリングし、テーブル名、列名、型、既存のコメント、少数のサンプル値といったメタデータをシステムに送信します。LLM のコストを削減してスループットを向上させるため、同じテーブルの複数の列を 1 回のリクエストでまとめてバッチ処理します。

当社の分類法はProtocol Buffersを使用して定義されており、現在100を超える階層的なデータラベルが含まれています。チームが追加のカテゴリを必要とする場合には、カスタム拡張の余地があります。これにより、ガバナンスおよびプラットフォームのステークホルダーは、少数の正規表現にとどまらず、「PII」と「機密」が何を意味するかについて共通の定義を持つことができます。

データ拡張

2つの拡張戦略により、分類の品質が大幅に向上します。

• AI による列コメントの生成: コメントがない場合は、Databricks の AI 生成コメントを使用して、LLM と将来のテーブル利用者の両方に役立つ、簡潔で人間が判読可能な説明を合成します。
• Few-shot事例生成: 正解データセットを維持し、静的な事例と、Vector Searchで取得した動的な事例の両方を使用します。各列について名前、型、コメント、コンテキストから埋め込みを構築し、類似度が高い上位K個のラベル付き列を取得してプロンプトに含めます。

静的プロンプティングは、初期段階やラベル付きデータが限られている場合に最適であり、一貫性と再現性を提供します。成熟したシステムでは動的プロンプティングがより効果的です。ベクトル検索を使用して類似の例を取得し、大規模で多様なデータセットの新しいスキーマやデータドメインに適応します。

LLMオーケストレーション

システムの中核には、本番環境のスケールで LLM コールを管理する軽量なオーケストレーションレイヤーがあります。

主な機能は次のとおりです:

• 社内でホストされている LLM（例: Llama、Claude、GPT ベースのモデル）間でのマルチモデルルーティングと、モデルが利用できない場合の自動フォールバック。
• 一時的な障害やレート制限に対する、指数バックオフを用いた再試行ロジック。
• 空、無効、またはハルシネーションによって生成されたラベルを検出し、バックアップモデルでそれらのケースを再ランする検証フック。
• コンテキストを失うことなくトークンの使用量を最適化するために、複数の列に一度にアノテーションを付けるバッチ処理。

階層型ラベリングシステム

列ごとに次の 3 種類のラベルを予測します：

• マスキング、リダクション、厳格なアクセス制御を強化する、100以上のきめ細かいオプションから選択される詳細なラベル。
• 階層ラベル: 関連する詳細なラベルを、モニタリングとレポートに適した、より広範なカテゴリに集約します。
• データがリージョン内に留まる必要があるか、リージョンを越えて移動できるかを示し、データ移動ポリシーに直接反映されるレジデンシーラベル。

​予測の一貫性を保ち、ハルシネーションを減らすために、2段階のフローを採用しています。広範な分類ステップで上位カテゴリを割り当て、次に絞り込みステップでそのカテゴリ内の正確なラベルを選択します。これは、人間のレビュー担当者がまず「これはワークスペースデータである」と判断し、次に特定のワークスペース識別子ラベルを選択する手順を反映しています。

モデルのバージョン管理とラベル予測

単一の「最良」な構成に依存するのではなく、各モデルのセットアップは、カラムのラベリングを競い合うエキスパートとして扱われます。​

複数のモデルバージョンは、以下の点が異なり、並行して実行されます:

• プライマリおよびフォールバックのLLMの選択。
• 生成されたコメント 対 生のメタデータ
• プロンプト戦略 (静的few-shot 対 動的few-shot)。
• ラベルの粒度と分類体系のサブセット。

各エキスパートは、ラベルと0から100の間の信頼度スコアを生成します。次に、システムは信頼度が最も高いエキスパートのラベルを選択します。これは専門家混合 (Mixture-of-Experts) スタイルのアプローチであり、精度を向上させ、単一の構成からの時折の不正確な予測の影響を低減します。

この設計により、安全なエクスペリメントが可能になります。新しいモデルやプロンプト戦略を導入して既存のものと並行して実行し、defaultとして採用する前にメトリクスと下流のチケット量の両方で評価できます。

チケットの作成

パイプラインは現在のスキーマアノテーションと LLM の予測を継続的に比較し、意味のある逸脱を表面化させます。

一般的なケースは次のとおりです。

• 注釈なしで追加された新しい列。
• 列のコンテンツと一致しなくなった既存のアノテーション。
• リージョン間の移動が可能としてラベル付けされている、機密値を含む列。

システムが違反を検出すると、ポリシーエントリを作成し、テーブル、列、提案されたラベル、および信頼度に関するコンテキストを付けて、所有チームに対して JIRA チケットを発行します。これにより、データ分類の問題は、チームが他の本番運用インシデントを追跡するのと同じ方法で追跡および解決できる、継続的なワークフローに変わります。

影響と評価

このシステムは2,258個のラベル付きサンプルで評価されました。そのうち1,010個はPIIを含み、1,248個はPIIを含まないものでした。このデータセットにおいて、PII検出で最大92%の適合率と95%の再現率を達成しました。

ステークホルダーにとってさらに重要なことは、このデプロイメントが必要な運用上の成果を生み出したことです。

• レビュアーが生のスキーマからではなく、高品質な提案ラベルから作業を開始できるため、大規模な監査サイクルごとに手動レビューにかかる労力が数週間から数時間に短縮されました。
• 年次レビューで発見されるのではなく、スキーマの進化に合わせてラベリングのドリフトが継続的に検出されるようになりました。
• 安全であると誤ってラベル付けされた機密データに関するアラートが、より的を絞ったものになるため、セキュリティチームはノイズの多いルールベースのスキャナーのトリアージを行う代わりに、迅速に対応できます。
• マスキングとデータ所在地のポリシーは、アナリティクスとレポーティングで利用されるのと同じラベル分類体系を使用して、大規模に適用されます。

​適合率と再現率はガードレールとして機能しますが、システムはレビュー時間、ドリフト検出の遅延、週に生成される実用的なチケットの量などの成果を中心に調整されています。

まとめ

分類体系駆動のラベリングとMoEスタイルの評価フレームワークを組み合わせ、MLflowを使用してエクスペリメントとデプロイを管理することで、Databricksにおける既存のエンジニアリングおよびガバナンスのワークフローを有効化しました。これにより、スキーマが変更されてもラベルが最新の状態に保たれ、コンプライアンスレビューがより迅速かつ的確になり、プラットフォーム全体でマスキングルールとデータ所在地ルールを一貫して適用するために必要な実行フックが提供されます。

この作業の最もエキサイティングな部分は、社内での学びをデータ分類製品に直接統合することです。LogSentinel内でこれらの技術を運用化および検証するにつれて、我々の技術をDatabricksのデータ分類に直接組み込んでいます。

メタデータとサンプルの取り込み、コンテキストの拡張、複数のLLMのオーケストレーション、予測のポリシーおよびチケッティングシステムへのフィードという同じパターンは、信頼性が高く、進化し続けるデータ理解が必要とされるあらゆる場所で再利用できます。これらの知見を当社の中核的な製品に組み込むことで、すべての組織がDatabricksと同じ精度と規模で、コンプライアンスとガバナンスのためにデータインテリジェンスを活用できるようになります。

謝辞

このプロジェクトは、複数のエンジニアリングチームの協力によって実現しました。Anirudh Kondaveeti、Sittichai Jiampojamarn、Zefan Xu、Li Yang、Xiaohui Sun、Dibyendu Karmakar、Chenen Liang、Viswesh Periyasamy、Chengzu Ou、Evion Kim、Matthew Hayes、Benjamin Ebanks、Sudeep Srivastava の皆様のサポートと貢献に感謝します。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事