# NDSS-2026 專案實作計劃

> 本文件依 2026-07-08 對 `Codes/`、`Data/`、`models/`、git 狀態的實際盤點重寫,取代舊版(舊版描述的目錄結構已被整批重構,內容已過時)。同日稍後已重新核對一次(git log/status、models 大小、各階段檔案清單),內容與盤點結果一致,僅更新本文件自身的 git 狀態描述(見 §6)。

## 1. 論文與方法總覽

**論文標題:** *Turn a Blind Eye: Defending Indirect Prompt Injection via Oracle-Guided Self Attention Alignment*(ACL Rolling Review 投稿,匿名,作者疑似 JK Hu,原始檔案:`Original/_ARR_Jan_2026___JK_Hu__Copy_.pdf`)

**要解決的問題:** 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;目前程式碼實作是把 attention_mask 在 `I_a` 位置歸零,效果等價但作法不同,見第 3 節)
3. 用 KL 散度損失,強迫 student 在特定 attention head 上的 attention 分布去模仿 oracle 狀態下的分布 → 讓模型把資料裡混入的指令當成「看不見的 padding」

**兩大技術貢獻:**

| 貢獻 | 對應論文章節 | 對應程式碼階段 |
|---|---|---|
| Signal-Based Head Identification(用 ROC-AUC + LCB 選出「指令敏感」head) | §2.2 | `Codes/2-2_head_identification_scoring/` |
| 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 架構(目前實際目錄)

```
Codes/1_raw_dataset/                                   ✅ 已補齊(vendor 進三個第三方 repo)
    ├── FocalLoRA/                — 同作者前作,baseline 方法來源
    ├── Should-It-Be-Executed-Or-Processed/  — SEP 資料集(ICLR 2025, Zverev et al.)
    └── topicattack/              — TopicAttack 攻擊/資料(EMNLP 2025)
        ↓
Codes/2-1_head_identification_preprocess/               ✅ 已補齊
    將 SEP / inj-squad / FocalLoRA 原始資料轉成 <inst>/<data> 標記格式
        ↓
Codes/2-2_head_identification_scoring/                  ✅ 已對 Llama-3.1-8B-Instruct 跑出完整結果
    01_attn_sep.py        — 收集 attention 權重(存 pickle shard)
    02_score.py            — ROC-AUC / proportion 打分,產生 head_scoring_combined.json
    03_sep_pick_head.py     — LCB 排序,依多組 λ 產生 heads_sorted/*.json
    Ident_IH_FocalLora.py   — 用 FocalLoRA 同款演算法打分,產生對照組 heads_sorted/focallora.json
    04/05_visualize*.py     — attention 熱圖視覺化
        ↓
Codes/2-3_head_identification_thresholding/             ✅ 已補齊(對應論文 Eq.3 的 threshold sweep)
    Ident_H_00_remove_existing_knoweledge.py — 濾掉模型不需看 data 就能答對的樣本
    Ident_H_01_EvaluateInstructiveHead.py     — 掃過一串 topk% 遮蔽比例,產出 ASR/utility 曲線(最佳點仍是人工判讀,沒有自動選 r* 的程式碼)
        ↓
Codes/3-1_model_training_data_gen/                      ✅ 已補齊,已產出實際訓練資料
    EvaluateModel.py / generate_single_turn.py — 套用攻擊(none/naive/ignore)與分隔符設定(prompt-based / native-tool 空query / native-tool 純response),tokenize 並標記 <data>/<inst>
    generate_len_test.py — 產生固定長度的容量測試資料
        ↓
Codes/3-2_model_training/                               ✅ 有三代訓練腳本並存,見第 3 節
    train_attn_kl_clean.py  — 【目前最新主線】KL-to-teacher 對齊 + head-level LoRA mask + 內建 eval
    _tuning.fix.chunked.py   — 前一代(flash-attn 攔截 loss),chunked 記憶體優化版
    _tuning.fix.modified.py  — 前一代原始版(非 chunked)
        ↓
Codes/4-1_evaluation_single/                            ❌ 只有 README 目標說明,程式碼未寫
    規劃:訓練只看弱攻擊,評估要扛強攻擊,量測 攻擊強度 / ASR / utility 三維度
```

`Codes/lib_code/` 是橫跨 2-1~3-2 各階段共用的函式庫(`tokenize_data_mask.py`、`head_mask_inference.py`、`attack_defense_tools.py` 等),取代了舊版每個階段各自一份 `lib/` 的重複結構。

---

## 3. 目前程式碼狀態

### `Codes/3-2_model_training/` 三代訓練腳本的差異(依檔案內容時間排序,非目錄 mtime)

1. **`_tuning.fix.modified.py`**(607 行,2026-05-28 09:20)—「Head-LoRA Finetune (Refactored v2)」。monkeypatch `ALL_ATTENTION_FUNCTIONS["flash_attention_2"]` 攔截並重算最後一個 token 的 attention row,loss 由 `head_attention_loss`(把被選中 head 的 attention 推離 data token)+ `head_preservation_loss`(用 KL 讓同 KV group 裡沒被選中的 head 保持接近原模型)組成。
2. **`_tuning.fix.chunked.py`**(832 行,2026-05-28 11:25)— 同架構的 chunked 版,新增 `--attn-chunk-size` 等參數避免長序列時把整張 attention map 攤開造成 OOM。
3. **`train_attn_kl_clean.py`**(822 行,2026-05-28 12:11,**目前主線**)— 拿掉 flash-attn monkeypatch,改用明確的 **KL-to-teacher** 寫法:`collect_teacher_attn` 在關掉 LoRA adapter、把攻擊指令位置的 attention_mask 歸零的情況下跑一次 forward 取得 teacher 分布,`tuned_kl_loss`/`kl_to_teacher` 對選中 head 算 chunked KL 散度。LoRA 只掛在選中 head 所在層的 `q_proj`/`k_proj`,並用 `apply_head_mask_to_lora` 對輸出做 head-level 遮罩(只放行被選中 head 對應維度,比舊版「整層共用」更接近論文講的單一 head 級隔離)。內建 `eval_topicattack` / `eval_mmlu`,不再依賴外部 `evallib.py`。

`Codes/3-2_model_training/outputs_lora/` 底下已有 `train_attn_kl_clean.py` 在 Qwen 與 Llama 上跑出的實際 smoke test checkpoint(`attn_kl_clean_smoke/`、`attn_kl_clean_qwen_smoke/` 等,各含 `final/adapter_model.safetensors`),並附 `eval_mmlu_smoke.json` / `eval_topicattack_smoke.json`,證明整條「載入模型 → LoRA 訓練 → 內建評估」流程在本機能跑通(目前只驗證過 1-sample 規模,尚未跑滿整份資料集)。

### 已知與論文的實作差異

1. Oracle View 的實作方式與論文文字描述不同:論文 Eq.4 寫的是把 `I_a` 的 token 換成 padding token(`∅`);目前 `train_attn_kl_clean.py` 是保留原 token,但把該區段的 `attention_mask` 歸零並在 teacher pass 關閉 LoRA adapter。兩者對 teacher 分布的效果理論上等價(該位置一律拿不到 attention 分數),但語意上不完全等同論文寫法,值得跟作者確認是否為刻意簡化。
2. 論文 Eq.3 的「自動求最佳遮蔽比例 r*」在 `Ident_H_01_EvaluateInstructiveHead.py` 裡沒有自動化程式碼,仍是掃過固定一串 topk% 後人工判讀。
3. 論文 Figure 2 對比 FocalLoRA 用的排序設定是 `all_roc_inst_0.1`,`2-2_head_identification_scoring/model_score/sep_Llama-3.1-8B-Instruct/heads_sorted/` 底下確實存在這個檔案,但 `train_attn_kl_clean.py` 的 `--topk` 預設值是 `18.75p`、`--head-path` 預設指向 `all_roc_inst_0.1.json`——需確認論文最終數字是否就是用這組參數跑出來的。
4. `_tuning.fix.modified.py` / `_tuning.fix.chunked.py` 的 loss 除了 KL 項外還有額外的 preservation loss,`train_attn_kl_clean.py` 已經拿掉、改成純粹對選中 head 算 KL——如果论文 Eq.5 描述的是後者,那麼 `train_attn_kl_clean.py` 才是跟論文公式一致的版本。

---

## 4. 環境準備(2026-07-08 確認仍然有效)

### 4.1 硬體
- 本機(`/home/chou/JunKai`,hostname `ee6032lab-Precision-7960-Tower`):原生 Linux 6.8.0-111,**4x NVIDIA RTX A6000(49GB)**,`/` 剩餘約 950G。
- 目前 GPU 0/1/3 閒置,GPU 2 有其他 process 占用(~2GB,30% util)——訓練/下載建議避開 GPU 2。

### 4.2 conda 環境(已確認存在且套件版本一致)
```bash
conda activate focallora4
```
實際安裝版本(已用 `conda run -n focallora4 python -c "..."` 驗證):
```
torch==2.5.1+cu121
transformers==4.57.6
peft==0.19.1
flash-attn==2.7.4.post1
```
`transformers` 特意鎖在 4.x 系列最後一版,因為部分舊版訓練腳本(`_tuning.fix.*`)會 monkeypatch `transformers` 內部 API(`ALL_ATTENTION_FUNCTIONS`),跨大版本升級風險較高。`train_attn_kl_clean.py` 已不依賴這個 monkeypatch,但環境尚未變動,仍鎖在同一版本。

---

## 5. 資料與模型 Checklist

| 項目 | 狀態 | 說明 |
|---|---|---|
| `Codes/1_raw_dataset/`(FocalLoRA / SEP / topicattack) | ✅ 已補齊 | 已攤平成一般追蹤檔案(見 commit `01bb07d`) |
| `Codes/2-1_head_identification_preprocess/` | ✅ 已補齊 | SEP / inj_squid / FocalLora 轉檔腳本與資料都在 |
| `head_scoring` 結果(`2-2_.../model_score/`) | ✅ 已有 **Llama-3.1-8B-Instruct** 一組完整結果(6.9GB,含 ~30 組 λ 排序 json + 熱圖) | ❌ 尚未對任何 Qwen 模型跑過 |
| `focallora.json`(baseline 對照) | ✅ 已存在(`heads_sorted/focallora.json`) | 不再是待補項目 |
| `Data/`(訓練/評估用 jsonl/json) | ✅ 已補齊(`Indirect PIA detection/`、`InjecAgent/`) | |
| Qwen2-7B-Instruct | ✅ 已下載(15G) | `models/Qwen2-7B-Instruct/` |
| Qwen3-8B | ✅ 已下載(16G) | `models/Qwen3-8B/` |
| Qwen3-4B | ✅ 已下載(7.6G) | `models/Qwen3-4B/` |
| **Llama-3.1-8B-Instruct** | ❌ **本機未下載** | 但幾乎所有 stage-2/3 腳本的預設路徑(`DEFAULT_MODEL_PATH` 等)都寫死指向作者原機路徑 `/data/local/hujk/models/Llama-3.1-8B-Instruct`,且唯一存在的 head-scoring 結果就是針對這顆模型算的 → **目前最大的落差:有 head 排序結果的模型,本機沒有權重可以直接拿來訓練** |

---

## 6. Git 狀態與待處理事項(新增)

- Repo 只有 2 個 commit(`6edf7da` first commit、`01bb07d` 攤平 `1_raw_dataset` submodule),remote 為 `github.com/HenryChou020514/JunKai.git`,只有 `main` 分支,無 stash。
- **本文件所描述的整個新版目錄結構(`2-2_..._scoring/`、`2-3_..._thresholding/`、`3-1_..._data_gen/`、`4-1_.../`、`lib_code/`,以及 `train_attn_kl_clean.py` 系列腳本)目前都還是 git untracked**,而舊版扁平結構(`2-2_head_identification/`、`3-1_model_training_preprocess/`、`4_model_evaluation/`、`sharedlibs/`)雖已從硬碟刪除,但刪除動作也還沒 commit。換句話說 **git 歷史目前完全沒反映這次重構**,建議找時間整批 commit 一次,避免新結構與已跑出來的中間產物(head scoring 結果、smoke test checkpoint)在意外情況下遺失。
- 本文件(`IMPLEMENTATION_PLAN.md`)先前一度從工作目錄消失且未 commit;已於 2026-07-08 重寫回來,目前 `git status` 顯示為 `M`(相對 HEAD 的修改,而非刪除),尚未 commit。

---

## 7. 執行步驟(依目前程式碼實際可跑的路徑)

> Qwen 系列模型本機已就緒,但目前只有 Llama 的 head 排序結果。若要走「完全用 Qwen 跑通全流程」,第 1 步需要重跑;若要「先驗證訓練/評估邏輯」,可直接沿用現有 Llama head 排序結果,但需另外下載 Llama-3.1-8B-Instruct 權重。

1. **(視需要)針對 Qwen 模型補跑頭部識別**
   - `2-2_head_identification_scoring/Ident_IH_ALL_1-4_qwen2.sh` / `_qwen3.sh` / `_qwen3-4b.sh`
   - 內部依序執行 01(收集 attention)→ 02(打分)→ 03(排序產生 `heads_sorted/*.json`)→ 04(視覺化)

2. **(可選)頭部遮蔽效果掃描**
   - `2-3_head_identification_thresholding/Ident_H_01_EvaluateInstructiveHead_gpu0.sh`,依模型/資料集調整參數

3. **訓練資料生成**
   - `3-1_model_training_data_gen/generate_single_turn.py`(或 `EvaluateModel_gendataset_squad.sh`),依 `--config` 選擇 prompt-based / native-tool 兩類分隔符設定

4. **LoRA 訓練(主線)**
   - `train_attn_kl_clean.sh`(呼叫 `train_attn_kl_clean.py`),**執行前務必用以下參數覆蓋預設值**,否則會指向作者原機不存在的路徑:
     ```
     --model-path   models/Qwen3-8B (或其他本機已下載模型 / Llama 權重)
     --data-path    Codes/3-1_model_training_data_gen/single_turn/xxx.json
     --head-path    Codes/2-2_head_identification_scoring/model_score/<model>/heads_sorted/xxx.json
     --output-dir   自訂輸出路徑
     ```
   - 想先驗證流程可用 `--repeat-single-sample`、`--max-train-steps` 跑小規模 smoke test(`outputs_lora/` 下已有前例可參考)

5. **模型評估**
   - `train_attn_kl_clean.py --eval-only --lora-path <訓練產出的 adapter> --eval-topicattack --eval-mmlu`
   - `4-1_evaluation_single/` 規劃的「弱攻擊訓練、強攻擊評估」比較邏輯尚未實作,需要另外補程式碼

---

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

- [x] `IMPLEMENTATION_PLAN.md` 先前從工作目錄消失是否為刻意行為 → 已重寫回來,待 commit(見 §6)
- [ ] Oracle View 的實作(attention_mask 歸零 + 關閉 adapter)與論文 Eq.4(token 替換成 padding)之間的差異是否影響最終結果對應關係
- [ ] `train_attn_kl_clean.py` 的預設 `--topk 18.75p` / `--head-path .../all_roc_inst_0.1.json` 是否就是論文 Table 1/2、Figure 2 報告數字所用的設定
- [ ] 是否要正式下載 Llama-3.1-8B-Instruct(需 Meta gated repo 授權),或改為全面轉向以 Qwen 系列重跑論文對照實驗
- [ ] `4-1_evaluation_single/` 的評估程式碼由誰接手撰寫、預計時程
- [ ] 大量 untracked 的新目錄結構何時整批 commit,是否需要先清掉 `2-2_head_identification_scoring/model_score/` 這類體積龐大(6.9GB)的中間產物再進版控(或改用 git-lfs / .gitignore)

---

## 9. 參考資料
- 論文原檔:`Original/_ARR_Jan_2026___JK_Hu__Copy_.pdf`
- 論文簡報:`Original/OGSAA.pptx`
- 原始封存檔:`Original/SortedCode.tgz`
- `Paper/` 底下兩篇參考文獻:InjecAgent(`2403.02691v3`)、Can Indirect Prompt Injection Attacks Be Detected and Removed?(`2502.16580v5`,inj-TrivalQA/inj-SQuAD 資料集出處)
- Baseline 對照方法程式碼(已 vendor 進 `Codes/1_raw_dataset/FocalLoRA/`):https://github.com/HuJK/FocalLoRA