ml-security-research-lab/cachetrapCodice di ricerca per un attacco trojan gray-box che inverte un singolo bit della KV-cache in classificatori LLM fine-tuned e misura il tasso di successo dell'attacco per classe.
Repository per "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap cerca un singolo bit flip nella KV cache di un classificatore LLM fine-tuned e misura il conseguente attack success rate (ASR) per classe target.
Python 3.10, CUDA 12.x.
conda create -n cachetrap python=3.10
conda activate cachetrap
pip install -r requirements.txt
Accesso a Hugging Face. meta-llama/Llama-2-7b-chat-hf e
meta-llama/Llama-3.1-8B-Instruct sono gated. Accetta le licenze sull'Hub, poi:
huggingface-cli login
Un modello / un dataset:
python train_model.py --model llama3_1_8b --dataset arc_easy
Il checkpoint viene scritto in
TrainedModels/merged_<model>_<dataset>/.
Tutti i modelli × tutti i dataset:
bash scripts/train.sh
Modifica MODELS, DATASETS e NUM_GPUS all'inizio dello script per eseguire un sottoinsieme.
I log vanno in logs/training/.
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa
--dataset è il dataset su cui è stata addestrata la vittima; --calib_dataset sono i dati
disponibili all'attaccante. Seguendo il paper, le vittime addestrate su OpenBookQA, TREC
o ARC-Challenge vengono calibrate con ARC-Easy, e le vittime addestrate su ARC-Easy vengono
calibrate con OpenBookQA. I due batch script codificano esattamente questa suddivisione:
| Script | Dataset vittima | Calibrazione |
|---|---|---|
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
Ciascuno termina eseguendo summarize_logs.py sulla propria directory di log.
Esegui python attack.py --help per l'elenco completo delle opzioni.
logs/<run>/<model>_<dataset>.log — output completo per esecuzionelogs/<run>/summary.csv — una riga per classe target (accuratezza baseline, accuratezza
sotto attacco, ASR, posizione del flip)logs/<run>/summary.txt — i blocchi di riepilogo grezziflip_locations/*.json — posizioni dei flip, se è stato passato --save_flip_locationsPer rivalutare senza ripetere la ricerca:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
Per provare l'attacco senza addestrare nulla, scarica il checkpoint rilasciato
merged_llama3_1_8b_arc_easy e posizionalo sotto TrainedModels/:
TrainedModels/
└── merged_llama3_1_8b_arc_easy/
├── config.json
├── model-*.safetensors
├── tokenizer.json
└── ...
Poi esegui il seguente comando per attaccare:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
In alternativa, esegui il seguente comando per un attacco da zero
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
Questo stampa un riepilogo per classe di accuratezza baseline, accuratezza sotto attacco, ASR, e la posizione del flip selezionata.
Cita la versione pre-print come
@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}
}