ModelScan: モデルシリアライゼーション攻撃からの保護

機械学習（ML）モデルは、インターネット上、チーム内、チーム間で公開共有されています。基盤モデル（Foundation Models）の台頭により、公開されているMLモデルは、さらなるトレーニングやファインチューニングのためにますます消費されるようになっています。MLモデルは、重要な意思決定やミッションクリティカルなアプリケーションの駆動にますます使用されています。 それにもかかわらず、モデルは受信トレイ内のPDFファイルほどの厳格さではまだスキャンされていません。

この状況を変える必要があり、適切なツールがその第一歩です。

ModelScanは、Protect AIによるオープンソースプロジェクトで、モデルに安全でないコードが含まれているかどうかを判定するためにモデルをスキャンします。複数のモデル形式をサポートする最初のモデルスキャンツールです。 ModelScanは現在、H5、Pickle、SavedModel形式をサポートしています。これにより、PyTorch、TensorFlow、Keras、Sklearn、XGBoostを使用する際に保護され、さらに多くの形式が追加される予定です。

TL;DR

モデルのスキャンをすぐに始めたい場合は、簡単です:

pip install modelscan

インストールしたら、モデルをスキャンします:

modelscan -p /path/to/model_file.pkl

モデルをスキャンすべき理由

モデルは自動化されたパイプラインから作成されることが多く、データサイエンティストのノートパソコンから来るものもあります。どちらの場合でも、モデルは使用される前に、あるマシンから別のマシンへ移動する必要があります。モデルをディスクに保存するこのプロセスは、シリアライゼーションと呼ばれます。

モデルシリアライゼーション攻撃（Model Serialization Attack） とは、配布前にシリアライゼーション（保存）中にモデルの内容に悪意のあるコードが追加される攻撃です。これは、トロイの木馬の現代版です。

この攻撃は、モデルの保存および読み込みプロセスを悪用して機能します。model = torch.load(PATH) でモデルを読み込むと、PyTorchはファイルの内容を開き、その中のコードの実行を開始します。モデルを読み込んだ瞬間に、エクスプロイトが実行されます。

モデルシリアライゼーション攻撃は、以下を実行するために使用される可能性があります:

• 認証情報の窃取（環境内の他のシステムへのデータの書き込みおよび読み取りのためのクラウド認証情報）
• データの窃取（モデルに送信されたリクエスト）
• データポイズニング（モデルがタスクを実行した後に送信されるデータ）
• モデルポイズニング（モデル自体の結果を改ざん）

これらの攻撃は非常に簡単に実行でき、実際の動作例は 📓notebooks フォルダで確認できます。

モデルセキュリティの強制と自動化

ModelScanは堅牢なオープンソースのスキャンを提供します。包括的なAIセキュリティが必要な場合は、Guardian を検討してください。これは、エンタープライズグレードのモデルスキャン製品です。

Guardianの特徴:

1. 最先端のスキャン: 最新のスキャナ、より広範なモデルサポート、自動モデル形式検出にアクセスできます。
2. プロアクティブなセキュリティ: コード変更を必要とせずに、Hugging Faceモデルが環境に入る前にセキュリティ要件を定義して適用します。
3. エンタープライズ全体のカバレッジ: CI/CDパイプラインとシームレスに統合し、組織全体で一貫したセキュリティ態勢を実装します。
4. 包括的な監査証跡: すべてのスキャンと結果の完全な可視性を獲得し、脅威を効果的に特定して軽減できます。

はじめに

ModelScanの仕組み

機械学習フレームワークでモデルを読み込むと攻撃が自動的に実行されるなら、ModelScanは悪意のあるコードを読み込まずにどのようにモデルの内容をチェックするのでしょうか？

簡単です。ファイルの内容を文字列のように1バイトずつ読み取り、安全でないコードシグネチャを探します。これにより非常に高速で、コンピュータがディスクからファイルサイズ全体を処理するのと同じ時間（ほとんどの場合数秒）でモデルをスキャンします。また、安全です。

ModelScanは安全でないコードを次のようにランク付けします:

• CRITICAL
• HIGH
• MEDIUM
• LOW

問題が検出された場合は、すぐにモデルの作者に連絡して原因を特定してください。

場合によっては、データサイエンティストが再現を容易にするためにコードがモデルに埋め込まれていることがありますが、それは攻撃に対して開かれた状態になります。それがワークロードに適切かどうかは、ご自身の判断で決定してください。

サポートされているモデルとフレームワーク

これは継続的に拡大される予定ですので、リリースノートの変更に注目してください。

現在、ModelScanはPickle派生のあらゆる形式とその他多くの形式をサポートしています:

インストール

ModelScanはPythonパッケージ（Python 3.9〜3.12対応）としてシステムにインストールされます。上記のように、ターミナルで次を実行してインストールできます:

pip install modelscan

プロジェクトの依存関係に含めて全員が利用できるようにするには、requirements.txt または pyproject.toml に次のように追加します:

modelscan = ">=0.1.1"

TensorflowまたはHD5形式のモデル用スキャナは、extrasを指定したインストールが必要です:

pip install 'modelscan[ tensorflow, h5py ]'

CLIでのModelScanの使用

ModelScanはCLI経由で次の引数をサポートしています:

モデルは他の形式のデジタルメディアと同じように、信頼できないソースからのコンテンツは使用前にスキャンしてください。

CLI終了コード

CLIの終了ステータスコードは次のとおりです:

• 0: スキャンが正常に完了し、脆弱性が見つかりませんでした
• 1: スキャンが正常に完了し、脆弱性が見つかりました
• 2: スキャンに失敗しました。スキャン中にmodelscanでエラーが発生しました
• 3: サポートされているファイルがツールに渡されませんでした
• 4: 使用法エラー。CLIに無効または不完全なオプションが渡されました

Pythonでプログラム的にModelScanを使用する

ModelScanはCLIで簡単に使用できますが、Pythonアプリケーションやワークフローに直接統合することもできます。

from modelscan.modelscan import ModelScan
from modelscan.settings import DEFAULT_SETTINGS

# Initialize ModelScan with default settings
scanner = ModelScan(settings=DEFAULT_SETTINGS)

# Scan a model file or directory 
results = scanner.scan("/path/to/model_file.pkl")

# Check if issues were found
if scanner.issues.all_issues:
    print(f"Found {len(scanner.issues.all_issues)} issues!")
    
    # Access issues by severity
    issues_by_severity = scanner.issues.group_by_severity()
    for severity, issues in issues_by_severity.items():
        print(f"{severity}: {len(issues)} issues")
        
# Generate a report (default is console output)
scanner.generate_report()

独自の設定でスキャン動作をカスタマイズできます:

# Start with default settings and customize
custom_settings = DEFAULT_SETTINGS.copy()

# Update settings as needed
custom_settings["reporting"]["module"] = "modelscan.reporting.json_report.JSONReport"
custom_settings["reporting"]["settings"]["output_file"] = "scan_results.json"

# Initialize with custom settings
scanner = ModelScan(settings=custom_settings)

結果の理解

スキャンが完了すると、問題が見つかった場合に次のような出力が表示されます:

ここでは、モデルに ReadFile と WriteFile の両方で安全でない演算子があります。 明らかに、モデルがファイルを任意に読み書きすることは望ましくありません。ここで、モデルの作成者に連絡して、これが何を意図していたのかを確認します。この特定のケースでは、攻撃者がAWS認証情報を読み取り、それらを別の場所に書き込むことを可能にします。

これは使用に関して明確にNOです。

MLパイプラインとCI/CDパイプラインへのModelScanの統合

アドホックなスキャンは素晴らしい第一歩です。新しいモデルを探索のためにダウンロードするたびにこれを行うように、自分自身、同僚、友人に徹底させてください。しかし、本番のMLOpsプロセスにおけるセキュリティを向上させるには十分ではありません。

次のことを達成するには、モデルスキャンを複数回実行する必要があります:

1. さらなる作業のために読み込む前に、すべての事前学習済みモデルをスキャンして、侵害されたモデルがモデル構築環境やデータサイエンス環境に影響を与えるのを防ぎます。
2. トレーニング後、すべてのモデルをスキャンして、新しいモデルを侵害するサプライチェーン攻撃を検出します。
3. エンドポイントにデプロイする前にすべてのモデルをスキャンして、保存後にモデルが侵害されていないことを確認します。

以下の赤いブロックは、従来のMLパイプラインでのこれを強調しています。

このプロセスは、LLM、基盤モデル、または外部モデルのファインチューニングやあらゆる変更でも同様です。

モデルのデプロイがMLパイプラインの外部で行われる場合も、CI/CDシステムのデプロイプロセスにスキャンを組み込んで、使用を保護してください。

さらに詳しく

📓notebooks フォルダ内では、TensorFlowやPyTorchなどのさまざまなMLフレームワークに対してモデルシリアライゼーション攻撃がどのように実行されるかを示す多数のノートブックを探索できます。

これらの攻撃が正確にどのように機能するかについてもっと深く掘り下げるには、🖹 Model Serialization Attack Explainer を確認してください。

静的コンテキストでモデルを評価する他のアプローチに遭遇した場合は、ぜひご連絡ください。私たちはもっと学びたいと思っています!

ライセンス

Copyright 2024 Protect AI

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

謝辞

私たちは、PickleScan を構築した Matthieu Maitre から多大な影響を受けました。 その成果に感謝し、ModelScanで大幅に拡張しました。ModelScanは、PickleScanと同じ精神でOSSとして公開されています。

貢献

私たちは、オープンソースのModelScanプロジェクトへのあなたの貢献を歓迎します。 貢献したい場合は、Contribution page の詳細に従ってください。