ModelScan: защита от атак на сериализацию моделей

Модели машинного обучения (ML) публикуются в открытом доступе в интернете, используются внутри команд и между командами. Развитие фундаментальных моделей привело к тому, что публичные 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 предоставляет надёжное сканирование с открытым исходным кодом. Если вам нужна комплексная безопасность ИИ, рассмотрите Guardian. Это наш корпоративный продукт для сканирования моделей.

Возможности Guardian:

1. Передовое сканирование: доступ к нашим новейшим сканерам, более широкой поддержке моделей и автоматическому определению формата моделей.
2. Проактивная безопасность: определяйте и применяйте требования безопасности для моделей Hugging Face до их попадания в вашу среду — без изменений кода.
3. Корпоративное покрытие: реализуйте согласованную политику безопасности во всей организации, легко интегрируясь с вашими CI/CD-конвейерами.
4. Всеобъемлющий аудит-след: получите полную видимость всех сканирований и результатов, что позволит эффективно выявлять и устранять угрозы.

Начало работы

Как работает ModelScan

Если загрузка модели с помощью вашего фреймворка машинного обучения автоматически выполняет атаку, как ModelScan проверяет содержимое без загрузки вредоносного кода?

Всё просто: он читает содержимое файла по одному байту за раз, как строку, и ищет сигнатуры небезопасного кода. Это делает его невероятно быстрым: модели сканируются за то время, которое требуется вашему компьютеру для чтения файла с диска (в большинстве случаев — секунды). Кроме того, это безопасно.

ModelScan классифицирует небезопасный код по уровням:

• КРИТИЧЕСКИЙ
• ВЫСОКИЙ
• СРЕДНИЙ
• НИЗКИЙ

Если обнаружена проблема, немедленно свяжитесь с авторами модели, чтобы определить причину.

В некоторых случаях код может быть встроен в модель, чтобы специалисту по данным было проще воспроизвести результаты, но это открывает возможность для атаки. Используйте собственное суждение, чтобы решить, допустимо ли это для ваших рабочих нагрузок.

Какие модели и фреймворки поддерживаются?

Этот список будет постоянно расширяться, поэтому следите за изменениями в наших примечаниях к выпускам.

В настоящее время ModelScan поддерживает любой формат на основе Pickle и многие другие:

Установка

ModelScan устанавливается в вашу систему как пакет Python (поддерживаются Python 3.9–3.12). Как показано выше, установить его можно, выполнив в терминале:

pip install modelscan

Чтобы добавить его в зависимости вашего проекта и сделать доступным для всех, добавьте его в requirements.txt или pyproject.toml следующим образом:

modelscan = ">=0.1.1"

Для сканеров моделей в форматах Tensorflow или HD5 требуется установка с дополнительными зависимостями:

pip install 'modelscan[ tensorflow, h5py ]'

Использование ModelScan через CLI

ModelScan поддерживает следующие аргументы через CLI:

Помните, что модели — это такая же форма цифрового контента, как и любая другая. Перед использованием сканируйте содержимое из любых ненадёжных источников.

Коды состояния выхода CLI

Коды состояния выхода CLI:

• 0: Сканирование завершено успешно, уязвимости не найдены
• 1: Сканирование завершено успешно, найдены уязвимости
• 2: Сбой сканирования: modelscan завершился с ошибкой во время сканирования
• 3: Инструменту не были переданы поддерживаемые файлы
• 4: Ошибка использования: в CLI были переданы недопустимые или неполные параметры

Использование ModelScan программно в Python

Хотя 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 и записать их в другое место.

Это однозначное «нет» использованию.

Интеграция ModelScan в ваши ML-пайплайны и CI/CD-пайплайны

Разовое сканирование — отличный первый шаг. Пожалуйста, приучите себя, коллег и друзей делать это всякий раз, когда они загружают новую модель для изучения. Но этого недостаточно для повышения безопасности производственных MLOps-процессов.

Сканирование моделей необходимо выполнять более одного раза, чтобы добиться следующего:

1. Сканируйте все предварительно обученные модели перед загрузкой для дальнейшей работы, чтобы скомпрометированная модель не повлияла на вашу среду разработки моделей или среду Data Science.
2. Сканируйте все модели после обучения, чтобы обнаружить атаку на цепочку поставок, которая ставит под угрозу новые модели.
3. Сканируйте все модели перед развёртыванием на конечной точке, чтобы убедиться, что модель не была скомпрометирована после хранения.

Красные блоки ниже показывают это в традиционном ML-конвейере.

Те же процессы применимы к тонкой настройке (fine-tuning) или любым модификациям LLM, фундаментальных моделей и внешних моделей.

Встраивайте сканирование в процессы развёртывания в ваших CI/CD-системах, чтобы обеспечить безопасное использование и при развёртывании моделей, если это происходит за пределами ваших ML-пайплайнов.

Глубокое погружение

В папке 📓notebooks вы можете изучить ряд блокнотов, которые наглядно демонстрируют, как атаки на сериализацию моделей могут выполняться против различных ML-фреймворков, таких как TensorFlow и PyTorch.

Чтобы детальнее разобраться в том, как именно работают эти атаки, ознакомьтесь с 🖹 Model Serialization Attack Explainer.

Если вам известны другие подходы к оценке моделей в статическом контексте, пожалуйста, свяжитесь с нами — нам будет интересно узнать о них больше!

Лицензирование

Copyright 2024 Protect AI

Лицензировано на условиях Apache License, Version 2.0 («Лицензия»); вы не можете использовать этот файл иначе, как в соответствии с Лицензией. Копию Лицензии можно получить по адресу

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

Если иное не требуется применимым законодательством или не согласовано в письменной форме, программное обеспечение, распространяемое по Лицензии, распространяется на основе «AS IS» («КАК ЕСТЬ») БЕЗ КАКИХ-ЛИБО ГАРАНТИЙ ИЛИ УСЛОВИЙ, явных или подразумеваемых. См. Лицензию, чтобы ознакомиться с конкретными разрешениями и ограничениями.

Благодарности

Мы черпали вдохновение у Matthieu Maitre, создателя PickleScan. Мы ценим эту работу и значительно расширили её в ModelScan. ModelScan выпущен как open source в том же духе, что и PickleScan.

Внесение вклада

Мы будем рады вашему вкладу в наш проект с открытым исходным кодом ModelScan. Если вы хотите внести вклад, пожалуйста, ознакомьтесь с подробностями на странице Contribution.