今天我們來延伸討論 PoisonedRAG 和 AgentPoison 怎麼實作的。
目標是讓某個問題被檢索到攻擊者寫的文本、並照著錯誤答案回答。做法是把一筆汙染文本拆成兩段：S 負責被檢索、I 負責誘導答案。以下三個階段就是它實際的實作。
# src/attack.py（節錄）
question = target_queries[i]['query']
adv_texts_b = self.all_adv_texts[id]['adv_texts'][:self.adv_per_query] # I：GPT-4 事先生成
adv_text_a = question + "." # S
adv_texts = [adv_text_a + i for i in adv_texts_b] # P = S + I
adv_text_a 就是 S (被檢索的內容) ，直接把目標問題原封不動當開頭。檢索運作原理是靠著向量相似度（把文字轉成一串數值，數值越接近代表語意越像），而跟一個問題最相似的文字內容其實就是問題本身，所以這樣組合出來的 P (汙染文本)天生就跟目標問題極度近似，檢索時當然就毫不費力就能被優先查詢到。
adv_texts_b 是 I (誘導執行的內容) ，研究團隊已在事先請 GPT-4 另外生成好 I 並存在 json 裡；生成時給 GPT-4 關於目標問題 Q 和想要的錯誤答案 R的資料，要它寫一小段「讀起來像有憑有據、但會把答案誘導成 R 形式」的短文，負責生成需要的條件及內容。
這階段做出的東西：一筆同時貼近問題、又夾帶錯誤答案的黑箱文本。
黑箱靠「S = 問題本身」取巧；白箱仰賴已知的檢索器參數，就能直接把相似度往上優化。
# src/attack.py（節錄）
for it_ in range(self.num_iter):
p_emb = self.get_emb(self.c_model, p_sent) # 汙染文本的向量
sim = torch.mm(p_emb, q_emb.T) # 與目標問題向量的相似度（內積）
loss = sim.mean()
loss.backward() # 對每個 token 求梯度
grad = embedding_gradient.get().sum(dim=0)
token_to_flip = random.randrange(len(adv_a)) # 隨機挑一個要替換的位置
candidates = hotflip_attack(grad[token_to_flip], embeddings.weight,
increase_loss=True, num_candidates=self.num_cand)
def hotflip_attack(averaged_grad, embedding_matrix, increase_loss=False, num_candidates=1, ...):
gradient_dot_embedding_matrix = torch.matmul(embedding_matrix, averaged_grad)
if not increase_loss:
gradient_dot_embedding_matrix *= -1
_, top_k_ids = gradient_dot_embedding_matrix.topk(num_candidates)
return top_k_ids
sim 是相似度，loss = sim.mean()，注意這裡是要相似度的數值變高——攻擊者要的就是相似度越高越好。loss.backward() 會算出每個 token 對相似度的影響方向，也就是所謂的梯度。
hotflip_attack 是這段程式碼的核心。HotFlip 是一種一次只取代一個 token、用梯度找出「換哪個字最能提升目標分數」的方法。 它把梯度和整個詞表每個字的向量做內積（torch.matmul），內積越大代表換上去越能拉高相似度，取前幾名當候選。之所以要用梯度，是因為詞表有好幾萬個字，逐一試不可行，梯度先幫你篩掉絕大多數沒參考性的 Token 。
這階段做出的東西：白箱測試能更容易的把相似度推高，比黑箱文本的方法 S=Q 更穩定，也更難用「相似度異常」去偵測。
汙染文本做好後注入知識庫，被測試的受害端會以下文形式讀取資料：
# src/prompts.py
MULTIPLE_PROMPT = 'You are a helpful assistant... Answer the question given the
information in those contexts. ... \n\nContexts: [context] \n\nQuery: [question] \n\nAnswer:'
使用者問 Q 時，檢索器會把最相近的幾筆資料（此時已包含那 5 筆 P）填進 [context]欄位，模型只依 context 回答。此階段的評估標準如以下：P 有沒有被檢索進前幾序列（檢索命中率），以及模型最後有沒有成功被誘導講出錯誤答案 R（ASR）。
這階段的結果：在 268 萬筆正常文本的 NQ 知識庫裡，每個問題只注入 5 筆，最後文本的 ASR 最高到 97%，過去的基準方法在同條件下不到 70%；測試資料與工具共計為三個資料集、八個模型。
AgentPoison 把 PoisonedRAG 那種偏人腦產出的 S 內容，升級成用最佳化函式自動化「算」出來的觸發字串後門，攻擊形式也從誘導錯誤的文本答案升級成誘導 agent 的偏差行為（例如逼自駕車急煞）。核心是一個跑 1000 輪的優化迴圈，我們分四階段分析。
汙染範例的組成是：一個正常指示詞 ＋ 觸發字串（10 個 token）＋ 目標行動。整段話裡只有觸發字串要被優化，所以梯度只從那 10 個特殊 token 取值並計算是否需要替換成更佳數值。
class GradientStorage:
def __init__(self, module, num_adv_passage_tokens):
self.num_adv_passage_tokens = num_adv_passage_tokens
module.register_full_backward_hook(self.hook)
def hook(self, module, grad_in, grad_out):
# 只留末端 N 個 token（＝觸發字串）的梯度，其餘不動
self._stored_gradient = grad_out[0][:, -self.num_adv_passage_tokens:]
register_full_backward_hook 的目的是在回傳過程設置攔截點，grad_out[0][:, -N:] 只截止至最後 N 個 token 的梯度。
這階段做的事：把可被優化 Token 鎖定在小部份且可植入的觸發字串上，不讓優化過程影響到過多使用者原本的指示詞內容。
def compute_avg_cluster_distance(query_embedding, cluster_centers):
distances = torch.norm(query_embedding.unsqueeze(1) - cluster_centers, dim=2)
overall_avg_distance = torch.mean(torch.mean(distances, dim=1))
variance = compute_variance(query_embedding)
score = overall_avg_distance - 0.1 * variance # 要最大化
return score
def compute_variance(embeddings):
mean_embedding = torch.mean(embeddings, dim=0, keepdim=True)
return torch.mean(torch.norm(embeddings - mean_embedding, dim=1))
cluster_centers 是先用正常指示詞在向量空間中算出來的幾個資料相對群集的「向量核心」，代表正常區域。這個函式關於：
overall_avg_distance：帶觸發字串的指示詞，離這些正常核心的平均距離。要最大化它，也就是把觸發指示詞推離正常區域，距離越遠越好。- 0.1 * variance：variance 是不同的觸發指示詞的發散程度（compute_variance 算的是每個點到自身的平均距離）。扣掉這段發散的距離，想辦法提高觸發指示詞的密集程度 ，讓檢索到這塊區域時能更高機率的選擇使用觸發詞。這階段的結果：這個 loss 保證觸發字串一出現，Agent 就會迅速檢索到已經被汙染的範例密集區域，並導致 Agent 的行為偏差，並且在未進行觸發時因為觸發字串的獨特性，攻擊不會被誤用與偵查到。
挑選 token 用的 hotflip_attack 跟 PoisonedRAG 白箱是同一個。差別是研究(二)多挑了十倍數量的候選Token，並有不同的篩選機制：
def candidate_filter(candidates, num_candidates, token_to_flip, adv_passage_ids, ppl_model):
ppl_scores = []
for candidate in candidates:
temp = adv_passage_ids.clone(); temp[:, token_to_flip] = candidate
ppl_scores.append(compute_perplexity(temp, ppl_model, device) * -1)
_, top_k = torch.tensor(ppl_scores).topk(num_candidates)
return candidates[top_k]
def compute_perplexity(input_ids, model, device):
outputs = model(input_ids, labels=input_ids)
loss = outputs[0]
return torch.exp(loss) # 困惑度
困惑度（perplexity）是指衡量一段文字對語言模型理解而言有多不通順，計算出的困惑值越低表示這段文字越像自然語言(也就是人話)。 這裡把困惑度乘上 -1 再取 topk，等於留下困惑度最低的候選，讓算出來的觸發字串不像一串亂碼。
這階段做的事：把候選 Token 從「能提高檢索概率」進階篩選成「能提高檢索、而且讀起來也很自然通順」的形式，這樣一來更方便繞過「看到亂碼或未知數據就判定攔截」的防禦機制。
for it_ in range(args.num_iter): # 1000 輪
for _ in range(args.num_grad_iter): # 累積 30 批梯度
loss = compute_avg_cluster_distance(query_embeddings, cluster_centers)
loss.backward()
grad += embedding_gradient.get().sum(dim=0) / args.num_grad_iter
token_to_flip = random.randrange(args.num_adv_passage_tokens)
candidates = hotflip_attack(grad[token_to_flip], embeddings.weight,
increase_loss=True, num_candidates=args.num_cand*10)
candidates = candidate_filter(candidates, args.num_cand, ...) # 困惑度濾
for candidate in candidates: # 逐一實測檢索分數
candidate_scores[i] += compute_avg_cluster_distance(candidate_query_embeddings, ...)
if (candidate_scores > current_score).any():
if target_gradient_guidance: # 在「檢索有提升」的候選裡
target_loss = target_word_prob(data, ..., "STOP", ...) # 再要求真的觸發目標動作
adv_passage_ids[:, token_to_flip] = candidates[best_candidate_idx]
這樣綜整上述內容，跑完一輪測試的完整動作是：
先累積 30 批梯度、用 hotflip_attack 選出一批候選 Token、用困惑度濾一遍、再逐一代進去實測檢索分數、只留有提升的 Token；接著 target_word_prob 再檢查這個候選 Token 會不會讓 agent 真的輸出偏差行為。當檢索和執行的兩階段都測試通過時，才會將符合標準的 token 替換上去。跑滿 1000 輪時，最終的觸發字串就成形了。
這階段的結果：帶觸發字串的指示詞有 80% 以上能成功誘導 Agent 檢索到汙染範例（ASR-retrieval），而真的能讓 Agent 進一步做出目標行動的提示詞約佔 63%（ASR-target）；記憶汙染比例不到 0.1%，甚至單一 token、單筆範例都可能奏效。
從「針對單一問題以及只影響回覆時的答案」，再到「針對觸發條件、演算法生成提示詞、影響 Agent 實際行為、還能跨檢索器轉移威脅內容」。這些討論與研究路徑說明，記憶汙染正從我們人為攻擊，變成一套可自動化、可最佳化的攻擊形式。
技術核心是同一個 HotFlip，差別只在目標函數的複雜度。
研究本質都是在操弄向量空間裡的排序、距離、密集程度，把汙染內容盡可能的擺在優先檢索區塊。
感謝大家今日份的閱讀，我們明天見。
參考資料
iThome鐵人賽
看影片追技術