CacheTrap

Репозиторий для статьи "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.

arXiv

CacheTrap ищет одиночный битовый флип в KV-кэше дообученного классификатора LLM и измеряет полученную долю успешных атак (ASR) для каждого целевого класса.

---

1. Установка

Python 3.10, CUDA 12.x.

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

Доступ к Hugging Face. meta-llama/Llama-2-7b-chat-hf и meta-llama/Llama-3.1-8B-Instruct требуют одобрения. Примите лицензии на Hub, затем:

huggingface-cli login

---

2. Обучение

Одна модель / один датасет:

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

Чекпоинт сохраняется в TrainedModels/merged_<model>_<dataset>/.

Все модели × все датасеты:

bash scripts/train.sh

Отредактируйте MODELS, DATASETS и NUM_GPUS в начале скрипта, чтобы запустить подмножество. Логи пишутся в logs/training/.

Примечание: На GPU с меньшим объёмом памяти используйте train_model_lowmem.py вместо train_model.py — как в команде выше, так и внутри scripts/train.sh.

---

3. Атака

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

--dataset — это датасет, на котором обучалась жертва; --calib_dataset — данные, доступные атакующему. Согласно статье, жертвы, обученные на OpenBookQA, TREC или ARC-Challenge, калибруются с помощью ARC-Easy, а жертвы, обученные на ARC-Easy, калибруются с помощью OpenBookQA. Два пакетных скрипта реализуют именно это разделение:

Скрипт	Датасеты жертвы	Калибровка
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

Каждый из них в конце запускает summarize_logs.py по своему каталогу логов. Запустите python attack.py --help для полного списка опций.

Выходные данные

• logs/<run>/<model>_<dataset>.log — полный вывод каждого запуска
• logs/<run>/summary.csv — одна строка на целевой класс (базовая точность, точность под атакой, ASR, местоположение флипа)
• logs/<run>/summary.txt — необработанные блоки сводки
• flip_locations/*.json — местоположения флипов, если был передан --save_flip_locations

Чтобы переоценить без повторного поиска:

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. Быстрый тест с опубликованным чекпоинтом

Чтобы опробовать атаку без какого-либо обучения, скачайте опубликованный чекпоинт merged_llama3_1_8b_arc_easy и поместите его в TrainedModels/:

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

Затем выполните следующую команду для атаки:

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

Либо выполните следующую команду для новой атаки

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

Это выведет сводку по каждому классу: базовую точность, точность под атакой, ASR и выбранное местоположение флипа.

---

Цитирование

Цитируйте препринтную версию следующим образом

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