hujk/OGAAA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
01bb07dba80f1023c0c2c6a473fc33e47dce6a32
OGAAA/IMPLEMENTATION_PLAN.md
HenryChou020514 6edf7da2b7 first commit
2026-07-07 19:03:00 +08:00

10 KiB
Raw Blame History

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)

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 條款):

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」)