ml-security-research-lab/cachetrapCódigo de pesquisa para um ataque trojan de caixa cinzenta que inverte um único bit da KV-cache em classificadores LLM ajustados e mede a taxa de sucesso do ataque por classe.
Repositório para "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap busca um único bit flip no KV cache de um classificador LLM ajustado e mede a taxa de sucesso do ataque (ASR) resultante por classe alvo.
Python 3.10, CUDA 12.x.
conda create -n cachetrap python=3.10
conda activate cachetrap
pip install -r requirements.txt
Acesso ao Hugging Face. meta-llama/Llama-2-7b-chat-hf e
meta-llama/Llama-3.1-8B-Instruct são restritos. Aceite as licenças no Hub e, em seguida:
huggingface-cli login
Um modelo / um dataset:
python train_model.py --model llama3_1_8b --dataset arc_easy
O checkpoint é gravado em
TrainedModels/merged_<model>_<dataset>/.
Todos os modelos × todos os datasets:
bash scripts/train.sh
Edite MODELS, DATASETS e NUM_GPUS no topo do script para executar um subconjunto.
Os logs vão para logs/training/.
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa
--dataset é o dataset no qual a vítima foi treinada; --calib_dataset são os dados
disponíveis para o atacante. Seguindo o artigo, vítimas treinadas em OpenBookQA, TREC
ou ARC-Challenge são calibradas com ARC-Easy, e vítimas treinadas em ARC-Easy são
calibradas com OpenBookQA. Os dois scripts em lote codificam exatamente essa divisão:
| Script | Datasets da vítima | Calibração |
|---|---|---|
scripts/attack_bundle1.sh | openbookqa, trec, arc_challenge | arc_easy |
scripts/attack_bundle2.sh | arc_easy | openbookqa |
bash scripts/attack_bundle1.sh
bash scripts/attack_bundle2.sh
Cada um termina executando summarize_logs.py sobre seu próprio diretório de logs.
Execute python attack.py --help para a lista completa de opções.
logs/<run>/<model>_<dataset>.log — saída completa por execuçãologs/<run>/summary.csv — uma linha por classe alvo (acurácia baseline, acurácia
sob ataque, ASR, localização do flip)logs/<run>/summary.txt — os blocos de resumo brutosflip_locations/*.json — localizações dos flips, se --save_flip_locations foi passadoPara reavaliar sem repetir a busca:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
Para testar o ataque sem treinar nada, baixe o checkpoint
merged_llama3_1_8b_arc_easy disponibilizado e coloque-o em TrainedModels/:
TrainedModels/
└── merged_llama3_1_8b_arc_easy/
├── config.json
├── model-*.safetensors
├── tokenizer.json
└── ...
Em seguida, execute o seguinte comando para atacar:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
Alternativamente, execute o seguinte para um novo ataque
python attack.py \
--model llama3_1_8b \
--dataset arc_easy \
--calib_dataset openbookqa \
--threat_model graybox \
--save_flip_locations flip_locations/llama3_1_8b_arc_easy_fresh.json
Isso imprime um resumo por classe da acurácia baseline, acurácia sob ataque, ASR, e a localização do flip selecionada.
Cite a versão de pré-print como
@article{nahian2025cachetrap,
title={CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs},
author={Nahian, Mohaiminul Al and Almalky, Abeer Matar A and Aragonda, Gamana and Zhou, Ranyang and Ahmed, Sabbir and Ponomarev, Dmitry and Yang, Li and Angizi, Shaahin and Rakin, Adnan Siraj},
journal={arXiv preprint arXiv:2511.22681},
year={2025}
}