ml-security-research-lab/cachetrap파인튜닝된 LLM 분류기에서 단일 KV-cache 비트를 뒤집고 클래스별 공격 성공률을 측정하는 그레이박스 트로이 목마 공격을 위한 연구 코드.
"CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs" 논문을 위한 저장소, IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap은 파인튜닝된 LLM 분류기의 KV 캐시에서 단일 비트 플립을 탐색하고, 그로 인해 발생하는 대상 클래스별 공격 성공률(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는 게이트가 설정되어 있다. 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로 캘리브레이션한다. 두 배치 스크립트는 정확히 이 분할을 구현한다:
| 스크립트 | 희생 모델 데이터셋 | 캘리브레이션 |
|---|---|---|
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}
}