ModelScan: Protección contra los ataques de serialización de modelos

Los modelos de Machine Learning (ML) se comparten públicamente a través de internet, dentro de los equipos y entre equipos. El auge de los modelos fundacionales ha hecho que los modelos de ML públicos se consuman cada vez más para entrenamiento adicional y ajuste fino. Los modelos de ML se utilizan cada vez más para tomar decisiones críticas y respaldar aplicaciones de misión crítica. A pesar de ello, los modelos aún no se escanean con el mismo rigor que un archivo PDF en tu bandeja de entrada.

Esto debe cambiar, y disponer de las herramientas adecuadas es el primer paso.

ModelScan es un proyecto de código abierto de Protect AI que escanea modelos para determinar si contienen código inseguro. Es la primera herramienta de escaneo de modelos compatible con múltiples formatos de modelos. Actualmente, ModelScan admite los formatos H5, Pickle y SavedModel. Esto te protege cuando usas PyTorch, TensorFlow, Keras, Sklearn, XGBoost, y hay más en camino.

TL;DR

Si estás listo para empezar a escanear tus modelos, es sencillo:

pip install modelscan

Con él instalado, escanea un modelo:

modelscan -p /path/to/model_file.pkl

Por qué deberías escanear los modelos

Los modelos suelen crearse a partir de pipelines automatizados; otros pueden provenir del portátil de un científico de datos. En cualquier caso, el modelo debe trasladarse de una máquina a otra antes de usarse. Ese proceso de guardar un modelo en disco se llama serialización.

Un ataque de serialización de modelos es aquel en el que se añade código malicioso al contenido de un modelo durante la serialización (guardado) antes de su distribución: una versión moderna del Caballo de Troya.

El ataque funciona explotando el proceso de guardado y carga de los modelos. Cuando cargas un modelo con model = torch.load(PATH), PyTorch abre el contenido del archivo y comienza a ejecutar el código que contiene. En el instante en que cargas el modelo, el exploit ya se ha ejecutado.

Un ataque de serialización de modelos puede utilizarse para ejecutar:

• Robo de credenciales (credenciales de la nube para escribir y leer datos en otros sistemas de tu entorno)
• Robo de datos (la solicitud enviada al modelo)
• Envenenamiento de datos (los datos enviados después de que el modelo haya realizado su tarea)
• Envenenamiento del modelo (alterar los resultados del propio modelo)

Estos ataques son increíblemente sencillos de ejecutar y puedes ver ejemplos prácticos en nuestra carpeta 📓notebooks.

Aplicar y automatizar la seguridad de los modelos

ModelScan ofrece un escaneo robusto de código abierto. Si necesitas seguridad integral de IA, considera Guardian. Es nuestro producto de escaneo de modelos de nivel empresarial.

Características de Guardian:

1. Escaneo de vanguardia: Accede a nuestros escáneres más recientes, mayor compatibilidad con modelos y detección automática del formato de modelo.
2. Seguridad proactiva: Define y aplica requisitos de seguridad para los modelos de Hugging Face antes de que entren en tu entorno, sin necesidad de cambios en el código.
3. Cobertura a nivel empresarial: Implementa una postura de seguridad cohesionada en toda tu organización, integrándose sin problemas con tus pipelines de CI/CD.
4. Registro de auditoría integral: Obtén visibilidad total de todos los escaneos y resultados, lo que te permite identificar y mitigar amenazas de manera efectiva.

Primeros pasos

Cómo funciona ModelScan

Si cargar un modelo con tu framework de machine learning ejecuta automáticamente el ataque, ¿cómo comprueba ModelScan el contenido sin cargar el código malicioso?

Sencillo: lee el contenido del archivo byte a byte, como si fuera una cadena, buscando firmas de código inseguras. Esto lo hace increíblemente rápido, escaneando los modelos en el tiempo que tarda tu ordenador en procesar el tamaño total del archivo desde el disco (segundos en la mayoría de los casos). También es seguro.

ModelScan clasifica el código inseguro como:

• CRITICAL
• HIGH
• MEDIUM
• LOW

Si se detecta un problema, ponte en contacto inmediatamente con los autores del modelo para determinar la causa.

En algunos casos, el código puede estar incrustado en el modelo para facilitar la reproducción a un científico de datos, pero te expone a ataques. Usa tu criterio para determinar si eso es apropiado para tus cargas de trabajo.

¿Qué modelos y frameworks son compatibles?

Esto se ampliará continuamente, así que presta atención a los cambios en nuestras notas de versión.

Actualmente, ModelScan admite cualquier formato derivado de Pickle y muchos otros:

Instalación

ModelScan se instala en tus sistemas como un paquete de Python (se admiten Python 3.9 a 3.12). Como se mostró antes, puedes instalarlo ejecutando esto en tu terminal:

pip install modelscan

Para incluirlo en las dependencias de tu proyecto y que esté disponible para todos, añádelo a tu requirements.txt o pyproject.toml de esta manera:

modelscan = ">=0.1.1"

Los escáneres para modelos con formato Tensorflow o HD5 requieren instalación con extras:

pip install 'modelscan[ tensorflow, h5py ]'

Uso de ModelScan mediante la CLI

ModelScan admite los siguientes argumentos a través de la CLI:

Recuerda que los modelos son como cualquier otra forma de medio digital: debes escanear el contenido de cualquier fuente no confiable antes de usarlo.

Códigos de salida de la CLI

Los códigos de estado de salida de la CLI son:

• 0: El escaneo se completó correctamente, no se encontraron vulnerabilidades
• 1: El escaneo se completó correctamente, se encontraron vulnerabilidades
• 2: El escaneo falló, modelscan lanzó un error durante el escaneo
• 3: No se pasaron archivos compatibles a la herramienta
• 4: Error de uso, se pasaron a la CLI opciones no válidas o incompletas

Uso de ModelScan programáticamente en Python

Aunque ModelScan se puede usar fácilmente mediante CLI, también puedes integrarlo directamente en tus aplicaciones o flujos de trabajo de 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()

Puedes personalizar el comportamiento del escaneo con tus propios ajustes:

# 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)

Comprender los resultados

Una vez que se ha completado un escaneo, verás una salida como esta si se encuentra un problema:

Aquí tenemos un modelo que tiene un operador inseguro tanto para ReadFile como para WriteFile en el modelo. Claramente no queremos que nuestros modelos lean y escriban archivos de forma arbitraria. Ahora nos pondríamos en contacto con el creador de este modelo para determinar qué esperaba que hiciera. En este caso particular, permite que un atacante lea nuestras credenciales de AWS y las escriba en otro lugar.

Eso es un NO rotundo para su uso.

Integración de ModelScan en tus pipelines de ML y pipelines de CI/CD

El escaneo ad hoc es un excelente primer paso; por favor, inculca en ti mismo, en tus colegas y amigos la práctica de hacer esto cada vez que descarguen un modelo nuevo para explorarlo. No es suficiente para mejorar la seguridad de los procesos de MLOps en producción.

El escaneo de modelos debe realizarse más de una vez para lograr lo siguiente:

1. Escanear todos los modelos preentrenados antes de cargarlos para trabajos posteriores, a fin de evitar que un modelo comprometido afecte a tus entornos de construcción de modelos o de ciencia de datos.
2. Escanear todos los modelos después del entrenamiento para detectar un ataque a la cadena de suministro que comprometa nuevos modelos.
3. Escanear todos los modelos antes de implementarlos en un endpoint para asegurarse de que el modelo no haya sido comprometido después del almacenamiento.

Los bloques rojos de abajo resaltan esto en un pipeline de ML tradicional.

Los procesos serían los mismos para el ajuste fino o cualquier modificación de LLM, modelos fundacionales o modelos externos.

Incorpora escaneos en los procesos de implementación de tus sistemas de CI/CD para asegurar el uso a medida que los modelos se implementan, también si esto se hace fuera de tus pipelines de ML.

Para profundizar

Dentro de la carpeta 📓notebooks puedes explorar varios notebooks que muestran exactamente cómo se pueden realizar ataques de serialización de modelos contra varios frameworks de ML como TensorFlow y PyTorch.

Para profundizar en el meollo de cómo funcionan exactamente estos ataques, consulta 🖹 Explicación de los ataques de serialización de modelos.

Si encuentras otros enfoques para evaluar modelos en un contexto estático, por favor contáctanos, nos encantaría aprender más.

Licencia

Copyright 2024 Protect AI

Licenciado bajo la Licencia Apache, Versión 2.0 (la «Licencia»); no puedes utilizar este archivo salvo que cumplas con la Licencia. Puedes obtener una copia de la Licencia en

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

Salvo que lo exija la ley aplicable o se acuerde por escrito, el software distribuido bajo la Licencia se distribuye sobre una base «TAL CUAL», SIN GARANTÍAS NI CONDICIONES DE NINGÚN TIPO, ya sean expresas o implícitas. Consulta la Licencia para conocer el texto específico que rige los permisos y las limitaciones bajo la Licencia.

Agradecimientos

Nos inspiró enormemente Matthieu Maitre, quien creó PickleScan. Apreciamos su trabajo y lo hemos ampliado significativamente con ModelScan. ModelScan se publica como OSS con el mismo espíritu que PickleScan.

Contribuciones

Nos encantaría que contribuyeras a nuestro proyecto de código abierto ModelScan. Si deseas contribuir, sigue los detalles en la página de contribución.