ModelScan: Proteção Contra Ataques de Serialização de Modelos

Modelos de Machine Learning (ML) são compartilhados publicamente pela internet, dentro de equipes e entre equipes. O aumento dos Modelos de Fundação (Foundation Models) fez com que modelos públicos de ML fossem cada vez mais consumidos para treinamento adicional/ajuste fino. Modelos de ML são cada vez mais usados para tomar decisões críticas e alimentar aplicações de missão crítica. Apesar disso, os modelos ainda não são verificados com o mesmo rigor dado a um arquivo PDF na sua caixa de entrada.

Isso precisa mudar, e as ferramentas adequadas são o primeiro passo.

O ModelScan é um projeto de código aberto da Protect AI que verifica modelos para determinar se eles contêm código inseguro. É a primeira ferramenta de verificação de modelos compatível com múltiplos formatos de modelo. Atualmente, o ModelScan suporta os formatos: H5, Pickle e SavedModel. Isso protege você ao usar PyTorch, TensorFlow, Keras, Sklearn, XGBoost, com mais formatos a caminho.

TL;DR

Se você está pronto para começar a verificar seus modelos, é simples:

pip install modelscan

Com ele instalado, verifique um modelo:

modelscan -p /path/to/model_file.pkl

Por Que Você Deve Verificar Modelos

Os modelos geralmente são criados a partir de pipelines automatizados; outros podem vir do laptop de um cientista de dados. Em qualquer um dos casos, o modelo precisa ser movido de uma máquina para outra antes de ser usado. Esse processo de salvar um modelo em disco é chamado de serialização.

Um Ataque de Serialização de Modelo ocorre quando código malicioso é adicionado ao conteúdo de um modelo durante a serialização (salvamento) antes da distribuição — uma versão moderna do Cavalo de Troia.

O ataque funciona explorando o processo de salvar e carregar modelos. Quando você carrega um modelo com model = torch.load(PATH), o PyTorch abre o conteúdo do arquivo e começa a executar o código dentro dele. No instante em que você carrega o modelo, o exploit já foi executado.

Um Ataque de Serialização de Modelo pode ser usado para executar:

• Roubo de credenciais (credenciais de nuvem para gravar e ler dados em outros sistemas do seu ambiente)
• Roubo de dados (a solicitação enviada ao modelo)
• Envenenamento de dados (os dados enviados depois que o modelo executa sua tarefa)
• Envenenamento do modelo (alteração dos resultados do próprio modelo)

Esses ataques são incrivelmente simples de executar, e você pode ver exemplos funcionais em nossa pasta 📓notebooks.

Aplicando e Automatizando a Segurança de Modelos

O ModelScan oferece verificação robusta de código aberto. Se você precisa de segurança abrangente para IA, considere o Guardian. É o nosso produto de verificação de modelos de nível enterprise.

Recursos do Guardian:

1. Verificação de Ponta: Acesse nossos scanners mais recentes, suporte mais amplo a modelos e detecção automática de formato de modelo.
2. Segurança Proativa: Defina e aplique requisitos de segurança para modelos Hugging Face antes que eles entrem no seu ambiente — sem necessidade de alterações de código.
3. Cobertura em Toda a Empresa: Implemente uma postura de segurança coesa em toda a organização, integrando-se perfeitamente aos seus pipelines de CI/CD.
4. Trilha de Auditoria Abrangente: Obtenha visibilidade total de todas as verificações e resultados, capacitando você a identificar e mitigar ameaças de forma eficaz.

Começando

Como o ModelScan Funciona

Se carregar um modelo com o seu framework de machine learning executa automaticamente o ataque, como o ModelScan verifica o conteúdo sem carregar o código malicioso?

Simples: ele lê o conteúdo do arquivo byte a byte, como se fosse uma string, procurando assinaturas de código inseguras. Isso o torna incrivelmente rápido, verificando modelos no tempo que o seu computador leva para processar o tamanho total do arquivo a partir do disco (segundos na maioria dos casos). Também é seguro.

O ModelScan classifica o código inseguro como:

• CRITICAL (Crítico)
• HIGH (Alto)
• MEDIUM (Médio)
• LOW (Baixo)

Se um problema for detectado, entre em contato imediatamente com os autores do modelo para determinar a causa.

Em alguns casos, código pode ser incorporado ao modelo para facilitar a reprodução por um cientista de dados, mas isso abre espaço para ataques. Use seu critério para determinar se isso é apropriado para suas cargas de trabalho.

Quais Modelos e Frameworks São Suportados?

Isso será expandido continuamente, então fique atento às mudanças em nossas notas de versão.

Atualmente, o ModelScan suporta qualquer formato derivado de Pickle e muitos outros:

Instalação

O ModelScan é instalado nos seus sistemas como um pacote Python (suporta Python 3.9 a 3.12). Como mostrado acima, você pode instalá-lo executando isto no seu terminal:

pip install modelscan

Para incluí-lo nas dependências do seu projeto e disponibilizá-lo para todos, adicione-o ao seu requirements.txt ou pyproject.toml assim:

modelscan = ">=0.1.1"

Os scanners para modelos nos formatos Tensorflow ou HD5 exigem instalação com extras:

pip install 'modelscan[ tensorflow, h5py ]'

Usando o ModelScan via CLI

O ModelScan suporta os seguintes argumentos via CLI:

Lembre-se: modelos são como qualquer outra forma de mídia digital — você deve verificar o conteúdo de qualquer fonte não confiável antes de usá-lo.

Códigos de Saída da CLI

Os códigos de status de saída da CLI são:

• 0: Verificação concluída com sucesso, nenhuma vulnerabilidade encontrada
• 1: Verificação concluída com sucesso, vulnerabilidades encontradas
• 2: Falha na verificação, o modelscan lançou um erro durante a verificação
• 3: Nenhum arquivo suportado foi passado para a ferramenta
• 4: Erro de uso, a CLI recebeu opções inválidas ou incompletas

Usando o ModelScan Programaticamente em Python

Embora o ModelScan possa ser facilmente usado via CLI, você também pode integrá-lo diretamente em seus aplicativos ou fluxos de trabalho 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()

Você pode personalizar o comportamento da verificação com suas próprias configurações:

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

Entendendo os Resultados

Depois que uma verificação é concluída, você verá uma saída como esta se um problema for encontrado:

Aqui temos um modelo que possui um operador inseguro para ReadFile e WriteFile no modelo. Claramente, não queremos que nossos modelos leiam e gravem arquivos arbitrariamente. Agora devemos entrar em contato com o criador deste modelo para determinar o que ele esperava que isso fizesse. Neste caso específico, isso permite que um atacante leia nossas credenciais AWS e as grave em outro lugar.

Isso é um NÃO definitivo para uso.

Integrando o ModelScan em Seus Pipelines de ML e Pipelines de CI/CD

A verificação ad-hoc é um ótimo primeiro passo. Por favor, internalize isso e incentive colegas e amigos a fazerem o mesmo sempre que baixarem um novo modelo para explorar. Isso não é suficiente para melhorar a segurança dos processos de MLOps em produção.

A verificação de modelos precisa ser realizada mais de uma vez para alcançar o seguinte:

1. Verificar todos os modelos pré-treinados antes de carregá-los para trabalho adicional, a fim de evitar que um modelo comprometido impacte seus ambientes de construção de modelos ou de ciência de dados.
2. Verificar todos os modelos após o treinamento para detectar um ataque à cadeia de suprimentos que comprometa novos modelos.
3. Verificar todos os modelos antes de implantá-los em um endpoint para garantir que o modelo não tenha sido comprometido após o armazenamento.

Os blocos vermelhos abaixo destacam isso em um pipeline de ML tradicional.

Os processos seriam os mesmos para ajuste fino ou quaisquer modificações de LLMs, modelos de fundação ou modelos externos.

Incorpore verificações nos processos de implantação dos seus sistemas de CI/CD para garantir o uso seguro dos modelos conforme eles são implantados, caso isso seja feito fora dos seus pipelines de ML.

Aprofundando

Dentro da pasta 📓notebooks, você pode explorar diversos notebooks que mostram exatamente como Ataques de Serialização de Modelo podem ser executados contra vários frameworks de ML, como TensorFlow e PyTorch.

Para se aprofundar em como exatamente esses ataques funcionam, confira 🖹 Explicador de Ataques de Serialização de Modelo.

Se você encontrar outras abordagens para avaliar modelos em um contexto estático, entre em contato. Adoraríamos aprender mais!

Licenciamento

Copyright 2024 Protect AI

Licenciado sob a Apache License, Versão 2.0 (a "Licença"); você não pode usar este arquivo exceto em conformidade com a Licença. Você pode obter uma cópia da Licença em

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

A menos que seja exigido pela lei aplicável ou acordado por escrito, o software distribuído sob a Licença é distribuído "COMO ESTÁ", SEM GARANTIAS OU CONDIÇÕES DE QUALQUER TIPO, expressas ou implícitas. Consulte a Licença para obter as permissões específicas e as limitações sob a Licença.

Agradecimentos

Fomos fortemente inspirados por Matthieu Maitre, que criou o PickleScan. Agradecemos o trabalho e o estendemos significativamente com o ModelScan. O ModelScan é um projeto de código aberto no mesmo espírito do PickleScan.

Contribuindo

Adoraríamos que você contribuísse com nosso projeto de código aberto ModelScan. Se quiser contribuir, siga os detalhes na página de Contribuição.