# NDSS-2026 專案實作計劃

> 本文件整理自對 `Codes/` 程式碼與 `Original/_ARR_Jan_2026___JK_Hu__Copy_.pdf` 論文的閱讀與比對結果,作為後續實際跑通整條 pipeline 的執行依據。

## 1. 論文與方法總覽

**論文標題:** *Turn a Blind Eye: Defending Indirect Prompt Injection via Oracle-Guided Self Attention Alignment*(ACL Rolling Review 投稿,匿名,作者疑似 JK Hu)

**要解決的問題:** Agent 使用工具時容易受 Indirect Prompt Injection(IPI)攻擊——惡意指令藏在外部檢索資料(tool output)裡,劫持模型執行流程。

**核心方法:Oracle-Guided Self Attention Alignment**
不靠外部監督或大量對抗訓練,而是讓模型「用自己在乾淨狀態下的行為當老師」:

1. **Student View** `X_stu`:輸入包含完整攻擊指令 `I_a`
2. **Oracle View** `X_ora`:把 `I_a` 遮蔽成 padding token(保留序列長度與位置結構,只讓模型看不到攻擊指令的語意)
3. 用 KL 散度損失,強迫 student 在特定 attention head 上的 attention 分布去模仿 oracle 狀態下的分布 → 讓模型把資料裡混入的指令當成「看不見的 padding」

**兩大技術貢獻:**

| 貢獻 | 對應論文章節 | 對應程式碼階段 |
|---|---|---|
| Signal-Based Head Identification(用 ROC-AUC + LCB 選出「指令敏感」head) | §2.2 | `Codes/2-2_head_identification/` |
| Oracle-Guided Self Attention Alignment(對選中 head 的 Q/K 做 LoRA + KL loss) | §2.3 | `Codes/3-2_model_training/` |

**主要實驗結果(節錄):** 對比 SecAlign、StruQ、Sandwich、Spotlight,在最難的 TopicAttack 攻擊下,本方法在 prompt-based separator 設定下把 ASR 從 SecAlign 的 90.67% 降到 65.77%;若用 native tool token 當分隔符,可進一步降到 11.22%。

**論文承認的限制:** 依賴分隔符不被偽造、僅適用標準 attention 架構(不含 Mamba/RWKV)、僅處理文字模態、存在被濫用來抑制安全防護的 dual-use 疑慮(但需要乾淨 Oracle 狀態當參照,風險可控)。

---

## 2. Pipeline 架構

```
原始資料 (1_raw_dataset/)                      ❌ 這份 cleaned copy 缺失
        ↓
前處理與 tokenize 標記 (2-1_head_identification_preprocess/)  ❌ 缺失
        ↓
[2-2] 頭部識別 (head identification)
      • Ident_IH_01_attn_sep.py   — 收集 attention 權重
      • Ident_IH_02_score.py     — ROC-AUC + proportion 打分
      • Ident_IH_03_sep_pick_head.py — LCB 排序,產生多組 λ 版本清單
      • Ident_IH_04/05_visualize*.py — attention 熱圖視覺化
      • EvaluateInstructiveHead.py — 掃描不同 topk 比例的防禦效果(對應論文 Eq.3 的 threshold sweep,但沒有自動選 r* 的程式碼,需人工看圖決定)
        ↓
[3-1] 訓練資料生成 (model_training_preprocess)
      • EvaluateModel.py — 套用攻擊(naive/ignore/suffix_attack/conv_attack…)與防禦(sandwich/spotlight),tokenize 並標記 <data>/<inst>
      • inj-likechen/generate_training_dataset*.py — 產生訓練樣本
        ↓
[3-2] LoRA 訓練 (model_training)
      • _tuning.modified.py — 核心訓練腳本:Flash Attention 2 攔截 + 對選中 head 的 Q/K 做 LoRA + KL 對齊損失
      • 內建 quick_eval_mmlu / quick_eval_asr_util 做即時評估
        ↓
[4] 模型評估 (model_evaluation)
      • TestInstructiveHead.py — 測試 head masking / LoRA 後模型是否忽略注入指令
      • GetAS.py — 計算 Attack Success Rate
```

---

## 3. 目前程式碼狀態

### ✅ 已修復
- **`lib` symlink 損毀問題**:原本 `Codes/{2-2,3-1,3-2,4}/lib` 都是空資料夾(從 `Original/SortedCode.tgz` 用不支援 symlink 的工具解壓所致)。已用 Git Bash 的 GNU tar 重新解壓並整個取代 `Codes/`,四個 `lib` 目錄現在都正確包含 9 個檔案。

### ⚠️ 已知與論文的實作差異(不影響「能不能跑」,但影響「結果是否對得上論文」)
1. LoRA 實際上是「整層」粒度(`q_proj`/`k_proj` 整層),不是嚴格的「單一 head 級」參數隔離;loss 只從被選中的 head 算,梯度理論上集中在對應輸出維度,但同層其他 head 共用同一組 LoRA 低秩矩陣。
2. `_tuning.modified.py` 的 `head_attention_loss` 除了論文 Eq.5 的 KL 項,還多了一個 `lambda_data * data_mass` 懲罰項,論文公式沒寫這一項。
3. 論文 Figure 2 用來對比 FocalLoRA 的 setup 是 `all_roc_inst_0.1`,但 `_tuning_llama.sh` 實際執行用的是 `user_prop_inst_0.1.json` ——兩者不是同一組 head 排序設定,需跟作者/組員確認論文圖表對應哪一次訓練。
4. Eq.3 的「自動求最佳遮蔽比例 r\*」沒有看到自動化程式碼,`EvaluateInstructiveHead_gpu*.sh` 是手動掃過一串固定的 topk 比例(`0 3.125p 6.25p ... 100p`),最佳點是事後人工判讀。

### ❌ 待補齊(詳見第 5 節 checklist)

---

## 4. 環境準備

### 4.1 作業系統 / 硬體(已於 2026-07-07 在本機完成設置)
- 論文原始訓練環境:**Nvidia RTX Pro 6000**(單卡)。
- 本機(`/home/chou/JunKai`)實際環境:原生 Linux 6.8.0,**3x NVIDIA RTX A6000(49GB)**,約 987G 可用硬碟空間。約定訓練/下載時鎖定 **GPU 0、1**(GPU 2 上有其他人/殘留 process 佔用 ~1.8GB 顯存,避免干擾)。
- Flash Attention 2 已成功編譯安裝,不再有 Windows/WSL2 的相容性問題(該問題僅存在於先前的 Windows 開發機上)。

### 4.2 Python 套件(已安裝,實際鎖定版本)
`Codes/requirements.txt` 是反推整理的建議下限版本。實際安裝時,**transformers 特意鎖在 4.x 系列的最後一版(4.57.6)**,而非當時最新的 5.13.0 —— 因為 `_tuning.modified.py` 會直接 monkeypatch `transformers.models.llama.modeling_llama.ALL_ATTENTION_FUNCTIONS` 這類內部 API,跨大版本(4→5)升級的相容性風險較高,4.57.6 仍完整支援 Qwen3(需求 ≥4.51)。

實際安裝版本:
```
torch==2.5.1+cu121
transformers==4.57.6
peft==0.19.1
bitsandbytes==0.49.2
flash-attn==2.7.4.post1
spacy==3.8.14 (+ en_core_web_trf==3.8.0)
```
```bash
python -m spacy download en_core_web_trf
```
(`tokenize_data_mask.py` 寫死用這個 transformer 版英文模型做指令偵測,**沒有 mock/fallback**,沒裝會直接壞掉。已確認安裝成功。)

### 4.3 conda 環境
腳本裡 hardcode 的環境名稱原本筆誤為 `foacllora4`,**已於 2026-07-07 全域修正為 `focallora4`**(修正了 21 個檔案,含所有 `.sh`、`_tune.ipynb`、本文件)。環境已建立於本機 `/home/chou/miniconda3`(用 `conda tos accept` 接受了 Anaconda 預設 channel 條款):
```bash
conda create -n focallora4 python=3.10
conda activate focallora4
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install transformers==4.57.6 accelerate peft bitsandbytes huggingface_hub tokenizers sentencepiece spacy scikit-learn numpy tqdm matplotlib hf_transfer datasets
pip install flash-attn==2.7.4.post1 --no-build-isolation
python -m spacy download en_core_web_trf
```

---

## 5. 資料與模型準備 Checklist

| 項目 | 狀態 | 說明 |
|---|---|---|
| `1_raw_dataset/`(topicattack 資料、prompts) | ❌ 缺失 | tgz 裡本來就沒有,需跟組員/作者要 |
| `2-1_head_identification_preprocess/`(SEP 資料集 jsonl) | ❌ 缺失 | 同上 |
| `focallora.json`(FocalLoRA baseline 的 head 排序,論文 Fig.2/Table 3 對比用) | ❌ 缺失 | 需跑 `HuJK/FocalLoRA` 自己的 head-scoring 流程產生,或直接要檔案 |
| `EvaluateInstructiveHead_gpu5.sh` | ❌ 空檔案(0 bytes) | 需比照 gpu0~4/6 補上對應的模型/資料集組合 |
| Qwen2-7B-Instruct | ✅ 2026-07-07 已完成(15G,4 shard 驗證齊全) | `hf download Qwen/Qwen2-7B-Instruct --local-dir models/Qwen2-7B-Instruct` |
| Qwen3-8B | 🔄 2026-07-07 下載中(接續 Qwen2-7B 之後排隊執行) | `hf download Qwen/Qwen3-8B --local-dir models/Qwen3-8B` |
| Qwen3-4B | 🔄 2026-07-07 下載中(接續 Qwen3-8B 之後排隊執行) | `hf download Qwen/Qwen3-4B --local-dir models/Qwen3-4B` |
| Llama-3.1-8B-Instruct | 🚫 使用者決定不下載 | 若之後要跟論文 Table 1/2 的 Llama 結果對照,仍需申請 Meta gated repo 授權 |

模型放置路徑(相對 `Codes/` 上一層,本機實際路徑):
```
/home/chou/JunKai/models/Qwen2-7B-Instruct/
/home/chou/JunKai/models/Qwen3-8B/
/home/chou/JunKai/models/Qwen3-4B/
```

---

## 6. 執行步驟(Qwen-only 路線)

> 以下依 stage 順序列出,括號內為對應腳本。實際參數(資料路徑、輸出路徑)請依你補齊資料後的實際位置調整。

1. **頭部識別**
   - `2-2_head_identification/Ident_IH_01-04_qwen2.sh`(Qwen2-7B-Instruct)
   - `2-2_head_identification/Ident_IH_01-04_qwen3.sh`(Qwen3-8B)
   - `2-2_head_identification/Ident_IH_01-04_qwen3-4b.sh`(Qwen3-4B)
   - 依序內部會跑 01(收集 attention)→ 02(打分)→ 03(排序產生 `heads_sorted/*.json`)→ 04(視覺化)

2. **(可選)頭部識別效果掃描**
   - `2-2_head_identification/EvaluateInstructiveHead_gpu1.sh`(Qwen3-8B)
   - `2-2_head_identification/EvaluateInstructiveHead_gpu3.sh`(Qwen2-7B)
   - `2-2_head_identification/EvaluateInstructiveHead_gpu4.sh`(Qwen3-8B)

3. **訓練資料生成**
   - `3-1_model_training_preprocess/EvaluateModel_gendataset*.sh`(需先把 `MODEL=` 改成 Qwen 路徑,目前預設寫的是 Llama)

4. **LoRA 訓練**
   - `3-2_model_training/_tuning_qwen2.sh`
   - `3-2_model_training/_tuning_qwen3.sh`

5. **模型評估**
   - `4_model_evaluation/EvaluateModel2lora.sh`(預設 `MODEL=../../models/Qwen3-8B`,可直接用)
   - `4_model_evaluation/EvaluateModel1base.sh`(⚠️ 裡面同時列了 Llama 與 Qwen3-8B 兩行,執行前要確認哪行沒被註解掉)

---

## 7. 待確認事項(建議直接問組員/論文作者)

- [ ] `focallora.json` 能否直接取得,或需要自己跑 `HuJK/FocalLoRA` 產生
- [ ] `1_raw_dataset/`、`2-1_head_identification_preprocess/` 的實際內容與取得方式
- [ ] `EvaluateInstructiveHead_gpu5.sh` 應該對應哪個模型/資料集組合
- [ ] 論文 Figure 2 的 `all_roc_inst_0.1` 與訓練腳本預設的 `user_prop_inst_0.1` 是否為同一批結果
- [ ] `_tuning.modified.py` / `_tuning.all_l.modified.py` / `_tuning.fix.modified.py` 三個訓練腳本,論文最終報告的數字是用哪一版跑出來的

---

## 8. 參考資料
- 論文原檔:`Original/_ARR_Jan_2026___JK_Hu__Copy_.pdf`
- 原始封存檔:`Original/SortedCode.tgz`
- Baseline 對照方法程式碼:https://github.com/HuJK/FocalLoRA(同作者,「Don't forget the enjoin: FocalLoRA for instruction hierarchical alignment in large language models」)