ml-security-research-lab/cachetrapCódigo de investigación para un ataque troyano de caja gris que invierte un único bit de la caché KV en clasificadores LLM ajustados y mide la tasa de éxito del ataque por clase.
Repositorio para "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap busca un único bit flip en la KV cache de un clasificador LLM fine-tuned y mide la tasa de éxito del ataque (ASR) resultante por clase objetivo.
Python 3.10, CUDA 12.x.
conda create -n cachetrap python=3.10
conda activate cachetrap
pip install -r requirements.txt
Acceso a Hugging Face. meta-llama/Llama-2-7b-chat-hf y
meta-llama/Llama-3.1-8B-Instruct están restringidos. Acepta las licencias en el Hub y luego:
huggingface-cli login
Un modelo / un dataset:
python train_model.py --model llama3_1_8b --dataset arc_easy
El checkpoint se escribe en
TrainedModels/merged_<model>_<dataset>/.
Todos los modelos × todos los datasets:
bash scripts/train.sh
Edita MODELS, DATASETS y NUM_GPUS al inicio del script para ejecutar un subconjunto.
Los logs van a logs/training/.
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa
--dataset es el dataset sobre el que se entrenó la víctima; --calib_dataset son los datos
disponibles para el atacante. Siguiendo el artículo, las víctimas entrenadas con OpenBookQA, TREC
o ARC-Challenge se calibran con ARC-Easy, y las víctimas entrenadas con ARC-Easy se
calibran con OpenBookQA. Los dos scripts por lotes codifican exactamente esa división:
| Script | Datasets de la víctima | Calibración |
|---|---|---|
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 uno termina ejecutando summarize_logs.py sobre su propio directorio de logs.
Ejecuta python attack.py --help para ver la lista completa de opciones.
logs/<run>/<model>_<dataset>.log — salida completa por ejecuciónlogs/<run>/summary.csv — una fila por clase objetivo (precisión baseline, precisión
bajo ataque, ASR, ubicación del flip)logs/<run>/summary.txt — los bloques de resumen sin procesarflip_locations/*.json — ubicaciones de flip, si se pasó --save_flip_locationsPara reevaluar sin repetir la búsqueda:
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 probar el ataque sin entrenar nada, descarga el checkpoint publicado
merged_llama3_1_8b_arc_easy y colócalo bajo TrainedModels/:
TrainedModels/
└── merged_llama3_1_8b_arc_easy/
├── config.json
├── model-*.safetensors
├── tokenizer.json
└── ...
Luego ejecuta el siguiente 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, ejecuta lo siguiente para un ataque nuevo
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
Esto imprime un resumen por clase de la precisión baseline, la precisión bajo ataque, el ASR y la ubicación de flip seleccionada.
Cita la versión pre-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}
}