ModelScan: Protezione dagli attacchi di serializzazione dei modelli

I modelli di Machine Learning (ML) vengono condivisi pubblicamente su internet, all'interno dei team e tra team diversi. L'ascesa dei Foundation Model ha portato a un consumo sempre maggiore di modelli ML pubblici per ulteriore addestramento/fine-tuning. I modelli ML sono sempre più utilizzati per prendere decisioni critiche e alimentare applicazioni mission-critical. Nonostante ciò, i modelli non vengono ancora scansionati con il rigore riservato a un file PDF nella tua casella di posta.

Questo deve cambiare, e gli strumenti adeguati sono il primo passo.

ModelScan è un progetto open source di Protect AI che scansiona i modelli per determinare se contengono codice non sicuro. È il primo strumento di scansione di modelli a supportare più formati di modelli. ModelScan attualmente supporta: formati H5, Pickle e SavedModel. Questo ti protegge quando usi PyTorch, TensorFlow, Keras, Sklearn, XGBoost, con altri in arrivo.

In breve

Se sei pronto per iniziare a scansionare i tuoi modelli, è semplice:

pip install modelscan

Con lo strumento installato, scansiona un modello:

modelscan -p /path/to/model_file.pkl

Perché dovresti scansionare i modelli

I modelli sono spesso creati da pipeline automatizzate, altri possono provenire dal laptop di un data scientist. In entrambi i casi, il modello deve essere spostato da una macchina all'altra prima di essere utilizzato. Questo processo di salvataggio del modello su disco è chiamato serializzazione.

Un attacco di serializzazione del modello è un attacco in cui codice dannoso viene aggiunto al contenuto di un modello durante la serializzazione (salvataggio) prima della distribuzione — una versione moderna del Cavallo di Troia.

L'attacco funziona sfruttando il processo di salvataggio e caricamento dei modelli. Quando carichi un modello con model = torch.load(PATH), PyTorch apre il contenuto del file e inizia a eseguire il codice al suo interno. Nell'istante in cui carichi il modello, l'exploit è già stato eseguito.

Un attacco di serializzazione del modello può essere utilizzato per eseguire:

• Furto di credenziali (credenziali cloud per scrivere e leggere dati su altri sistemi del tuo ambiente)
• Furto di dati (la richiesta inviata al modello)
• Data Poisoning (i dati inviati dopo che il modello ha svolto il suo compito)
• Model Poisoning (alterazione dei risultati del modello stesso)

Questi attacchi sono incredibilmente semplici da eseguire e puoi vedere esempi funzionanti nella nostra cartella 📓notebooks.

Applicare e automatizzare la sicurezza dei modelli

ModelScan offre una solida scansione open source. Se hai bisogno di una sicurezza AI completa, considera Guardian. È il nostro prodotto di scansione di modelli di livello enterprise.

Funzionalità di Guardian:

1. Scansione all'avanguardia: Accedi ai nostri scanner più recenti, a un supporto più ampio per i modelli e al rilevamento automatico del formato dei modelli.
2. Sicurezza proattiva: Definisci e applica requisiti di sicurezza per i modelli Hugging Face prima che entrino nel tuo ambiente, senza bisogno di modifiche al codice.
3. Copertura a livello aziendale: Implementa una postura di sicurezza coerente in tutta l'organizzazione, integrandoti perfettamente con le pipeline CI/CD.
4. Audit trail completo: Ottieni piena visibilità su tutte le scansioni e i risultati, consentendoti di identificare e mitigare efficacemente le minacce.

Per iniziare

Come funziona ModelScan

Se caricare un modello con il tuo framework di machine learning esegue automaticamente l'attacco, come fa ModelScan a controllare il contenuto senza caricare il codice dannoso?

Semplice: legge il contenuto del file un byte alla volta, proprio come una stringa, cercando firme di codice non sicure. Questo lo rende estremamente veloce: scansiona i modelli nel tempo impiegato dal tuo computer per elaborare la dimensione totale del file dal disco (secondi nella maggior parte dei casi). È anche sicuro.

ModelScan classifica il codice non sicuro come:

• CRITICAL
• HIGH
• MEDIUM
• LOW

Se viene rilevato un problema, contatta immediatamente gli autori del modello per determinare la causa.

In alcuni casi, il codice può essere incorporato nel modello per rendere più facile la riproduzione per un data scientist, ma ti espone ad attacchi. Usa la tua discrezione per determinare se è appropriato per i tuoi carichi di lavoro.

Quali modelli e framework sono supportati?

Questo sarà in continua espansione, quindi tieni d'occhio le modifiche nelle nostre note di rilascio.

Al momento, ModelScan supporta qualsiasi formato derivato da Pickle e molti altri:

Installazione

ModelScan viene installato sui tuoi sistemi come pacchetto Python (supporta da Python 3.9 a 3.12). Come mostrato sopra, puoi installarlo eseguendo questo comando nel terminale:

pip install modelscan

Per includerlo nelle dipendenze del tuo progetto e renderlo disponibile a tutti, aggiungilo al tuo requirements.txt o pyproject.toml in questo modo:

modelscan = ">=0.1.1"

Gli scanner per modelli in formato Tensorflow o HD5 richiedono l'installazione con gli extra:

pip install 'modelscan[ tensorflow, h5py ]'

Utilizzare ModelScan tramite CLI

ModelScan supporta i seguenti argomenti tramite CLI:

Ricorda che i modelli sono come qualsiasi altra forma di supporto digitale: dovresti scansionare il contenuto proveniente da qualsiasi fonte non attendibile prima dell'uso.

Codici di uscita della CLI

I codici di stato di uscita della CLI sono:

• 0: Scansione completata con successo, nessuna vulnerabilità trovata
• 1: Scansione completata con successo, vulnerabilità trovate
• 2: Scansione fallita, modelscan ha generato un errore durante la scansione
• 3: Nessun file supportato è stato passato allo strumento
• 4: Errore di utilizzo, alla CLI sono state passate opzioni non valide o incomplete

Utilizzare ModelScan programmaticamente in Python

Sebbene ModelScan possa essere facilmente utilizzato tramite CLI, puoi anche integrarlo direttamente nelle tue applicazioni o nei tuoi workflow 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()

Puoi personalizzare il comportamento della scansione con le tue impostazioni:

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

Comprendere i risultati

Una volta completata una scansione, vedrai un output come questo se viene trovato un problema:

Qui abbiamo un modello che ha un operatore non sicuro sia per ReadFile che per WriteFile. Chiaramente non vogliamo che i nostri modelli leggano e scrivano file arbitrariamente. A questo punto contatteremmo il creatore di questo modello per determinare cosa si aspettava che facesse. In questo caso specifico permette a un attaccante di leggere le nostre credenziali AWS e scriverle in un altro posto.

Questo è un chiaro NO all'utilizzo.

Integrare ModelScan nelle tue pipeline ML e CI/CD

La scansione ad hoc è un ottimo primo passo; per favore, imprimila in te stesso, nei colleghi e negli amici, da eseguire ogni volta che scaricano un nuovo modello da esplorare. Non è sufficiente per migliorare la sicurezza dei processi MLOps di produzione.

La scansione dei modelli deve essere eseguita più di una volta per ottenere quanto segue:

1. Scansiona tutti i modelli pre-addestrati prima di caricarli per ulteriori lavori, per evitare che un modello compromesso influisca sui tuoi ambienti di creazione di modelli o di data science.
2. Scansiona tutti i modelli dopo l'addestramento per rilevare un attacco alla supply chain che comprometta nuovi modelli.
3. Scansiona tutti i modelli prima di distribuirli su un endpoint per assicurarti che il modello non sia stato compromesso dopo l'archiviazione.

I blocchi rossi qui sotto evidenziano questo aspetto in una pipeline ML tradizionale.

I processi sarebbero gli stessi per il fine-tuning o qualsiasi modifica di LLM, modelli foundation o modelli esterni.

Incorpora le scansioni nei processi di distribuzione dei tuoi sistemi CI/CD per proteggere l'utilizzo quando i modelli vengono distribuiti, anche se questo avviene al di fuori delle tue pipeline ML.

Approfondimento

All'interno della cartella 📓notebooks puoi esplorare una serie di notebook che mostrano esattamente come gli attacchi di serializzazione dei modelli possono essere eseguiti contro vari framework ML come TensorFlow e PyTorch.

Per addentrarti più a fondo nei meccanismi esatti di questi attacchi, dai un'occhiata a 🖹 Spiegazione degli attacchi di serializzazione dei modelli.

Se incontri altri approcci per valutare i modelli in un contesto statico, contattaci, ci piacerebbe saperne di più!

Licenza

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.

Ringraziamenti

Siamo stati fortemente ispirati da Matthieu Maitre, che ha creato PickleScan. Apprezziamo il suo lavoro e lo abbiamo esteso in modo significativo con ModelScan. ModelScan è rilasciato come open source nello stesso spirito di PickleScan.

Contribuire

Ci piacerebbe che tu contribuissi al nostro progetto open source ModelScan. Se desideri contribuire, segui le indicazioni nella pagina Contributi.