OGAAA/IMPLEMENTATION_PLAN.md

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 環境(已確認存在且套件版本一致)

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. 待確認事項(建議直接問組員/論文作者)

• 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