Amazon SageMaker AI と MLflow を用いた判別型機械学習モデルの監視

機械学習（ML）モデルの有効性と精度は、トレーニングジョブが完了した直後からほぼ低下し始めます。消費者行動の変化、新製品のリリース、センサー技術のアップグレード、そして変化する経済・政治情勢など、モデルがトレーニング中に学習したパターンや確率を変化させる制御不能な要因の例が挙げられます。生産環境で展開されたモデルを精度と基本統計値の変化について積極的に監視することで、精度低下が問題となる前に介入することができます。モデル監視は、レイテンシ、アプリケーションの可用性、およびシステム全体の課題を特定するために使用されるその他の指標を追跡する AI 観測性ツール（AI observability tools）と組み合わせることができます。

本記事では、分類および回帰ユースケースに使用される判別型機械学習モデル（discriminative machine learning models）に焦点を当てます。生成 AI モデルについては、Amazon SageMaker AI エンドポイント推論における大規模言語モデル（LLM）の運用準備完了リアルタイム監視ソリューション を参照してください。判別型 ML モデルの品質低下を引き起こす要因は、おおまかに 2 つのカテゴリーに分類できます：

• データドリフトとは、入力データの統計的特性の変化を指します。これは、上流のデータソースにおける予期せぬ変化により列のデータ型が整数から浮動小数点に変化するといった単純なものから、全く新しい製品ラインがリリースされるような複雑なケースまで様々です。データドリフトは、トレーニングデータセットに対してベースライン統計量を計算し、これを時間経過とともに収集された生産環境のデータで計算した同じ統計量と比較することで測定できます。

• モデルドリフトとは、モデルが学習した確率的パターンが入力されるデータに適合しなくなったために、モデルが生み出す予測の精度が変化する現象を指します。例えば、経済の改善に伴う消費者行動の変化などが原因となり得ます。モデルドリフトは、正解ラベル（ground truth labels）を集めてモデル品質指標を計算し、これらの指標をモデルトレーニングプロセス中に計算された同じ指標と比較することで測定できます。

図 1: データドリフトとモデルドリフトが ML ワークフローにおいてどのように位置づけられるか

Amazon SageMaker AI は、組織が判別型および生成型の機械学習（ML）モデルを構築、トレーニング、デプロイ、管理するためのフルマネージド型サービスです。SageMaker AI はフルマネージドですが、よりカスタマイズ性の高いアプローチが必要な場合もあります。例えば、コスト効果の高い方法でモデルのライフサイクル全体を管理したい、マネージドサービスがサポートしていない独自のユースケースを監視したい、またはモデル監視を他の UI や観測パイプラインに統合したいといった要望があるかもしれません。そこで本記事では、データドリフトおよびモデルドリフトの計算のために、オープンソースの Evidently Python ライブラリ と Amazon SageMaker AI with MLflow を基盤としたモデル監視アーキテクチャをご紹介します。このモデル監視ソリューションの結果は、お好みのダッシュボードに統合したり、関連するステークホルダーへアラートを送信したり、自動的なモデル再トレーニングパイプラインをトリガーしたりするために使用できます。

ソリューションの概要

本ソリューションでは、モデルのトレーニングからデプロイに至るまでの機械学習ワークフローにおけるモデル監視の実装方法を示します。図 2 は、バッチ推論ユースケースにおけるワークフローを示しており、Slack へのアラート送信および MLflow でのモデル結果の可視化までを含みます。

図 2: バッチ推論用のモデル監視ワークフロー

ワークフローは以下の手順から構成されています：

• Amazon Simple Storage Service (Amazon S3) バケットから入力データを用いてモデルを学習させるトレーニングジョブ。同じトレーニングジョブはモデルメトリクスも計算し、これらは MLflow に保存できます。モデルの学習に使用されたベースラインデータセットは、本番環境でのモニタリングに利用できるよう、S3 内に別途保存されます。

• 本番ワークロードにおける推論のためのバッチトランスフォーム。バッチトランスフォームジョブの結果は S3 バケットに格納されます。

• バッチトランスフォームジョブの結果、および元のデータセットとモデルメトリクスを用いて、オープンソースの Evidently プリセット (Evidently presets) を使用し、データドリフトとモデル品質メトリクスを計算する処理ジョブが実行されます。なお、Evidently はモデルメトリクスを計算しますが、トレーニングランからの前回のメトリクスと比較して直接モデルドリフトを計算するわけではありません。ただし、カスタムコードを追加することで、この処理ジョブを拡張しモデルドリフトの計算を行うことが可能です。

バッチ変換および処理ジョブは、Amazon EventBridge Scheduler を使用してスケジュール管理可能なパイプラインにラップできます。

• すべての監視メトリクスと Evidently レポートは MLflow に保存され、ここでは実行履歴を時間経過とともに追跡し、実行間での結果を比較したり、レポートを可視化したりすることができます。

• 任意で、ドリフトが検出された場合、パイプラインは Amazon Simple Notification Service (Amazon SNS) を介してメールによりユーザーに通知を送信するトリガーを発動できます。

モデルデプロイに リアルタイムエンドポイント が使用されるユースケースでも、同様のアプローチを採用可能です。図 3 はリアルタイムエンドポイントのワークフローを示しています。重要な違いは、エンドポイントの入力と出力を S3 バケットにログ記録するために データキャプチャ が有効化されている必要がある点です。処理ジョブの代わりに、このアーキテクチャでは AWS Lambda 関数を使用して監視コードをデプロイします。どちらのアーキテクチャも、ご自身の好みに応じて処理ジョブまたは Lambda のいずれかを使用できます。データドリフトおよびモデルドリフトを計算する AWS Lambda 関数は、定期スケジュールで実行することも可能です。また、エンドポイントからのデータが S3 バケットに到着した際にトリガーして実行することもできます。

図 3: リアルタイムエンドポイントにおけるモデル監視ワークフロー

このリアルタイムエンドポイントのパターンは、Amazon SageMaker Hyperpod でも利用可能です。ここでは、モデル学習およびモデル推論に使用するためのクラスターをプロビジョニングできます。データキャプチャは、エンドポイントレベル、ロードバランサーレベル、またはモデルポッドレベルで有効化できます。推論リクエストとレスポンスは自動的に S3 バケットに書き込まれるため、そこからアクセスしてモデル監視を行うことができます。

上記の 2 つの図はいずれもモデル監視プロセスのみを対象としていますが、通常、モデル監視はより広範な MLOps ワークフローの一部として実装されます。図 4 は、このモデル監視ソリューションがエンドツーエンドの MLOps シナリオにどのように組み込まれるかを示しています。ここでは、ドリフト計算に使用される処理ジョブが、CI/CD ワークフローを通じてエンドポイントと共にデプロイされます。SageMaker AI を用いた MLOps の実装に関する詳細については、Amazon SageMaker AI を活用した企業向け MLOps 基盤ロードマップをご覧ください。

図 4: モデル監視が MLOps アーキテクチャとどのように統合されるか

ウォークスルー

このセクションでは、バッチトランスフォーム用のモデル監視ソリューションを設定するための手順を段階的に説明します。リアルタイムエンドポイント のサンプルを含む他の例については、完全なリポジトリ をご覧ください。

前提条件

このウォークスルーの手順に従う前に、以下の準備が必要です：

• このサービスが利用可能な AWS リージョン内の Amazon SageMaker AI ドメイン。
• SageMaker Studio 上の MLflow アプリ。
• SageMaker Studio 内の JupyterLab スペース（本ソリューションには ml.t3.medium インスタンスで十分です）。
• JupyterLab 上のソリューションリポジトリのコピー。スペースが起動した後、リポジトリをクローンしてください。

ウォークスルーを続けるには、ノートブック predictive_ml_experimentation_data_model_monitoring_evidently.ipynb を開いてください。なお、本ソリューションでは SageMaker Python SDK v3 を使用しています。DefaultMLFlowApp を使用する場合は、ノートブックコードが自動的に正しいアプリを識別します。それ以外の場合は、MLflow アプリの名前を指定してコードを編集してください。

モデルのトレーニングと推論

例示のノートブックでは、UCI Machine Learning Repository からの Bank Marketing dataset を使用しています。このデータセットには、ポルトガルの銀行を代表して行われたマーケティング電話に関する情報が含まれており、顧客が定期預金に加入するかどうかを予測することを目的としています。これは二値分類のユースケースであり、ターゲット変数は顧客が加入した場合に 1、加入しなかった場合に 0 の値を持ちます。

ノートブックの前半では、データのクリーニング、処理、およびトレーニング、バリデーション、テストセットへの分割が行われます。XGBoost モデルがトレーニングおよびバリデーションデータセット上で学習され、ログとメトリクスは MLflow に送信され、最終的なモデルオブジェクトは MLflow のモデルレジストリに登録されます。

このデータ処理およびモデル学習セクションにおける 2 つの重要なアクションが、後続のプロダクション環境でのモニタリングのためにモデルを準備します:

• トレーニングデータセットは、データドリフト（データ分布の変化）の計算に使用されるベースラインデータセットとして S3 バケットに保存されます。
• モデルメトリクスは MLflow に保存され、モデルドリフト（モデル性能の変化）の計算に利用できます。

図 5: MLflow で表示されるモデルメトリクス（精度、再現率、AUC）

次に、ノートブックはトレーニングジョブの出力からモデルオブジェクトを作成し、テスト用特徴量をモデルに通して推論を行うバッチトランスフォームジョブを設定します。

データドリフトとモデル品質の計算

Evidently には、さまざまな種類の データドリフト を計算するためのプリセットが用意されており、これらは特定の機械学習ユースケースに合わせてカスタマイズすることが可能です。サンプルノートブックでは DataDriftPreset と DataSummaryPreset を使用し、Evidently レポート（HTML 形式および JSON 形式）を MLflow に保存するヘルパー関数を作成するとともに、特定のドリフト値を抽出して MLflow のメトリクスとして個別に保存します。これにより、異なる実行結果の比較や、特定の数値に基づいたアラートの生成が可能になります。

図 6: MLflow のアーティファクトとしての Evidently データドリフトレポート

図 7: MLflow のデータドリフトメトリクスとパラメータ

各データドリフト計算を MLflow の実行（run）として保存するもう一つの利点は、モデル名やトレーニングジョブの名前、データセットのサイズなど、重要な情報をパラメータとして追加できることです。

Evidently は、モデル品質の計算用のプリセットも提供しています。このノートブックでは ClassificationPreset を使用しており、これはモデルの予測結果と正解データセットに基づいて、精度（accuracy）、適合率（precision）、再現率（recall）、F1 スコアなどを算出します。ClassificationPreset が持つ広範な指標群の代わりに、Evidently では カスタムレポート もサポートしており、特定のユースケースに最も関連性の高い指標を計算対象として指定できます。例えば、サンプルノートブックで使用されているマーケティングデータセットには不均衡なラベルが存在するため、精度よりも適合率、再現率、AUC（曲線下面積）の方がより関連性の高い指標となります。

データモニタリングと同様に、モデル品質に関する Evidently のレポートはアーティファクトとして MLflow に追加され、個々の指標は MLflow メトリクスとして抽出可能です。ただし、Evidently はモデルドリフト（元のトレーニングジョブから得られた指標と、予測結果から計算された指標との差）を算出するものではありません。しかし、モニタリングコードにドリフト計算を追加することは可能であり、その例については