CacheTrap

Dépôt pour « CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs », IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.

arXiv

CacheTrap recherche un unique bit flip dans le cache KV d'un classifieur LLM affiné et mesure le taux de réussite de l'attaque (ASR) résultant par classe cible.

---

1. Installation

Python 3.10, CUDA 12.x.

conda create -n cachetrap python=3.10
conda activate cachetrap
pip install -r requirements.txt

Accès Hugging Face. meta-llama/Llama-2-7b-chat-hf et meta-llama/Llama-3.1-8B-Instruct sont soumis à autorisation. Acceptez les licences sur le Hub, puis :

huggingface-cli login

---

2. Entraînement

Un modèle / un jeu de données :

python train_model.py --model llama3_1_8b --dataset arc_easy

Le checkpoint est écrit dans TrainedModels/merged_<model>_<dataset>/.

Tous les modèles × tous les jeux de données :

bash scripts/train.sh

Modifiez MODELS, DATASETS et NUM_GPUS en haut du script pour n'exécuter qu'un sous-ensemble. Les journaux sont écrits dans logs/training/.

Remarque : sur les GPU disposant de moins de mémoire, utilisez train_model_lowmem.py à la place de train_model.py, à la fois dans la commande ci-dessus et dans scripts/train.sh.

---

3. Attaque

python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa

--dataset est le jeu de données sur lequel la victime a été entraînée ; --calib_dataset est la donnée disponible pour l'attaquant. Conformément à l'article, les victimes entraînées sur OpenBookQA, TREC ou ARC-Challenge sont calibrées avec ARC-Easy, et les victimes entraînées sur ARC-Easy sont calibrées avec OpenBookQA. Les deux scripts par lots encodent exactement cette répartition :

Script	Jeux de données de la victime	Calibration
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

Chacun se termine par l'exécution de summarize_logs.py sur son propre répertoire de journaux. Exécutez python attack.py --help pour la liste complète des options.

Sorties

• logs/<run>/<model>_<dataset>.log — sortie complète par exécution
• logs/<run>/summary.csv — une ligne par classe cible (précision de référence, précision sous attaque, ASR, emplacement du flip)
• logs/<run>/summary.txt — les blocs de résumé bruts
• flip_locations/*.json — emplacements des flips, si --save_flip_locations a été passé

Pour réévaluer sans répéter la recherche :

python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
    --load_flip_locations flip_locations/llama3_1_8b_arc_easy.json

---

4. Test rapide avec un checkpoint publié

Pour essayer l'attaque sans rien entraîner, téléchargez le checkpoint publié merged_llama3_1_8b_arc_easy et placez-le sous TrainedModels/ :

TrainedModels/
└── merged_llama3_1_8b_arc_easy/
    ├── config.json
    ├── model-*.safetensors
    ├── tokenizer.json
    └── ...

Exécutez ensuite la commande suivante pour lancer l'attaque :

python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
    --load_flip_locations flip_locations/llama3_1_8b_arc_easy.json

Vous pouvez également exécuter ce qui suit pour une nouvelle attaque

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

Cela affiche un résumé par classe de la précision de référence, de la précision sous attaque, de l'ASR, et de l'emplacement du flip sélectionné.

---

Citation

Citez la version pré-print comme suit

@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}
}