nishitvp/better_opts_attacksFramework de ataque para quebrar defesas de injeção de prompt baseadas em fine-tuning (SecAlign, SecAlign++, StruQ) usando ataques adversariais cientes da arquitetura em LLMs.
Este repositório contém código para executar os ataques ASTRA e ASTRA++ que quebram SecAlign++, SecAlign e StruQ. Este repositório também contém alguns exemplos de ataques gerados e logs de ataques.
Isto clonará o repositório, bem como todos os submódulos, no local correto necessário.
git clone --recurse-submodules https://github.com/nishitvp/better_opts_attacks.git
A base de código não usa nenhum pacote Python especial além da coleção padrão de pacotes relacionados a LLM e Deep Learning. Para instalar os requisitos, execute
python3 -m pip install -r requirements.txt
e isso deve instalar os pacotes necessários no seu ambiente.
Primeiro, baixe os modelos base do HuggingFace (alguns podem ter acesso restrito, exigindo que você obtenha acesso pelo HuggingFace)
hf login --token ${HF_PASSWORD}
hf download meta-llama/Meta-Llama-3-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Meta-Llama-3-8B-Instruct
hf download mistralai/Mistral-7B-Instruct-v0.1 --local-dir secalign_refactored/secalign_models/mistralai/Mistral-7B-Instruct-v0.1
hf download meta-llama/Llama-3.1-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Llama-3.1-8B-Instruct
hf download facebook/Meta-SecAlign-8B --local-dir secalign_refactored/secalign_models/Meta-SecAlign-8B
Baixe os adaptadores SecAlign e os modelos StruQ. Há um script útil para baixar e extrair no lugar certo.
python3 setup_scripts/download_models.py
Quando executado a partir da raiz do projeto, este script deve baixar e extrair os modelos corretos no local correto. Para verificar, confira se secalign_refactored/secalign_models contém as três pastas huggyllama, meta-llama e mistralai, que devem conter cada uma o(s) próprio(s) modelo(s).
Para executar o ASTRA para um Adversário com Conhecimento Forte, o script principal é experiment.py. Este script pegará os 40 exemplos do conjunto de dados AlpacaFarm, injetará o payload Print exactly 'Hacked!' com a configuração especificada e executará o ASTRA, bem como o baseline (GCG), controlando a inicialização (inicialização aleatória), o número de passes forward, iterações etc.
O usuário pode controlar os seguintes parâmetros definindo-os como parâmetros de linha de comando ao chamar experiment.py -
--model-name: Os valores possíveis são meta-llama, meta-llama-instruct, mistralai, mistralai-instruct, huggyllama.
--defense: Os valores possíveis são undefended, struq e secalign. O valor padrão é secalign. O parâmetro --model-name, juntamente com o parâmetro --defense, especifica de forma única o modelo específico sendo usado, de acordo com a tabela abaixo
--model-name | --defense | Modelo carregado |
|---|---|---|
meta-llama-instruct | secalign | Adaptadores SecAlign com Meta-Llama-3-8B-Instruct pré-ajustado por instrução |
mistralai-instruct | secalign | Adaptadores SecAlign com Mistral-7B-Instruct-v0.1 pré-ajustado por instrução |
meta-llama | secalign | Meta-Llama-3-8B defendido com SecAlign (sem pré-ajuste por instrução) |
mistralai | secalign | Mistral-7B-Instruct-v0.1 defendido com SecAlign (sem pré-ajuste por instrução) |
huggyllama | secalign | Llama-2-7B defendido com SecAlign (sem pré-ajuste por instrução) |
meta-llama | struq | Meta-Llama-3-8B defendido com StruQ (sem pré-ajuste por instrução) |
mistralai | struq | Mistral-7B-v0.1 defendido com StruQ (sem pré-ajuste por instrução) |
huggyllama | struq | Llama-2-7B defendido com StruQ (sem pré-ajuste por instrução) |
meta-llama-instruct | undefended | Meta-Llama-3-8B-Instruct sem defesa (bruto) |
mistralai-instruct | undefended | Mistral-7B-Instruct-v0.1 sem defesa (bruto) |
meta-llama | undefended | Meta-Llama-3-8B sem defesa (bruto) (sem pré-ajuste por instrução) |
mistralai | undefended | Mistral-7B-Instruct-v0.1 sem defesa (bruto) (sem pré-ajuste por instrução) |
huggyllama | undefended | Llama-2-7B sem defesa (bruto) (sem pré-ajuste por instrução) |
Todos os outros pares de --model-name e --defense são inválidos.
--prefix-length: Comprimento do prefixo adversarial a ser usado em uma determinada execução de avaliação. O padrão é 5.
--suffix-length: Comprimento do sufixo adversarial a ser usado em uma determinada execução de avaliação. O padrão é 20.
--expt-folder-prefix: O caminho onde você deseja que os logs da execução do experimento sejam registrados.
Um comando completo pode parecer algo como -
python3 experiment.py --model-name meta-llama-instruct --defense secalign --prefix-length 5 --suffix-length 20 --expt-folder-prefix logs/astra_llama_secalign
Ao executar o comando acima, o script distribuirá automaticamente a carga de trabalho igualmente entre diferentes subprocessos, cada um dos quais usará uma GPU para calcular os ataques. Recomendamos executar os ataques em GPUs com pelo menos 48GB de memória.
Cada subprocesso registrará a transcrição completa de cada exemplo que atacar em uma subpasta separada do caminho especificado em --expt-folder-prefix. Os logs são armazenados em um formato semelhante a um banco de dados consultável contendo objetos Python (veja o arquivo utils/experiment_logger para mais detalhes). Cada transcrição registrada conterá as sequências de tokens, valores de perda e outros detalhes importantes, como a configuração de inicialização durante a otimização, que podem ser recuperados posteriormente.
Quando o experimento completo do ASTRA terminar, os logs gerados podem ser analisados posteriormente usando o arquivo analysis/analysis.ipynb. Basta copiar e colar o Jupyter notebook no caminho --expt-folder-prefix, substituir o caminho do modelo no notebook pelo caminho correto do modelo e executar o notebook. Isso deve imprimir o número de sucessos do GCG e o número de sucessos do ASTRA na última célula. (P.S. O script de análise demora um pouco para executar porque ele realmente gera as saídas em cada uma das sequências de tokens de entrada geradas durante a otimização para cada exemplo.)
O notebook analysis.ipynb também contém código auxiliar para plotar curvas de perda média e fazer explorações adicionais com os logs gerados.
Os scripts de análise foram mantidos separados fundamentalmente para garantir que a otimização em si não seja contaminada pelo código de análise post-facto.
Também compartilhamos alguns ataques gerados pelo ASTRA no arquivo zipado localizado em data/attacks_generated.zip. Concretamente, o arquivo zipado contém 4 arquivos, cada um com 40 exemplos de ataques gerados pelo ASTRA, dois para cada uma das versões defendidas com SecAlign do Mistral e do Llama-3, respectivamente. Os arquivos fornecem exemplos gerados para um orçamento de 20 e 25 tokens (em uma configuração (0, 20) e (5, 20)). A injeção já foi colocada no campo de entrada nas configurações.
Para executar o ASTRA++ para um Adversário com Conhecimento Fraco, o script principal é experiment_universal.py. Este script pegará 10 exemplos de treinamento do conjunto de dados AlpacaFarm, injetará o payload Print exactly 'Hacked!' com a configuração de prefixo e sufixo especificada e treinará para obter uma injeção de prompt universal ao final. O script também executa o baseline, versão universal do GCG, após o treinamento do ASTRA++.
Assim como o script experiment.py da seção anterior, experiment_universal.py também aceita os argumentos --model-name, --defense, --prefix-length, --suffix-length, --expt-folder-prefix.
Além de todos os modelos listados na seção anterior sobre a execução do ASTRA, o script experiment_universal.py também aceita mais um modelo possível, correspondente à defesa SecAlign++. Para treinar com a defesa SecAlign++, passe o parâmetro --model-name como secalign_refactored/secalign_models/Meta-SecAlign-8B e o parâmetro --defense como meta_secalign.
Há um argumento de linha de comando adicional que este script aceita, que é --training-run, que corresponde a qual conjunto de 10 exemplos a execução de treinamento ocorrerá.
Os demais parâmetros têm a mesma interpretação de antes.
Quando o experimento completo do ASTRA++ terminar, os logs gerados podem ser analisados posteriormente usando o arquivo analysis/analysis_universal.ipynb. Como antes, copie e cole o Jupyter notebook no caminho --expt-folder-prefix, carregue o modelo corretamente (há alguns detalhes complicados devido à forma como o SecAlign++ é carregado, mas isso será uniformizado no futuro). Isso deve imprimir o número de sucessos de teste do GCG e o número de sucessos do ASTRA, juntamente com as melhores injeções de prompt universais alcançadas.
Para o ASTRA++, enviamos as injeções de prompt universais geradas para o Llama-3.1-8B-Instruct defendido com SecAlign++ no arquivo data/universal_pis_secalign++.json, juntamente com os metadados e a configuração em que foram geradas. Além disso, também enviamos uma transcrição de todos os ataques em todas as execuções de treinamento de universalidade em data/final_result_logs.pkl (Prometemos, o pickle não contém código malicioso :).
Embora os scripts acima executem os ataques ASTRA e ASTRA++ prontos para uso, a base de código permite a configuração de praticamente todos os parâmetros que podem afetar o ataque, como tamanhos de conjuntos de dados de treinamento, limites escolhidos, sementes aleatórias, configurações de inicialização e assim por diante. Se você tiver alguma sugestão interessante, sinta-se à vontade para enviar um pull request.
Para quaisquer dúvidas, relatos de bugs, detalhes ou esclarecimentos, sinta-se à vontade para abrir uma issue no Github ou enviar um e-mail aos autores.