ml-security-research-lab/cachetrapИсследовательский код для атаки типа «троян» в модели серого ящика, которая инвертирует один бит KV-кэша в дообученных классификаторах LLM и измеряет долю успешных атак по каждому классу.
Репозиторий для статьи "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap ищет единичный битовый флип в KV-кэше дообученного классификатора LLM и измеряет полученную долю успешных атак (ASR) для каждого целевого класса.
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 являются gated-моделями. Примите лицензии на Hub, затем:
huggingface-cli login
Одна модель / один датасет:
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/.
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. Два batch-скрипта реализуют именно это разделение:
| Скрипт | Датасеты жертвы | Калибровка |
|---|---|---|
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
Чтобы попробовать атаку без обучения, скачайте опубликованный
чекпоинт 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}
}