LLM Jailbreak之GCG攻击

AI安全
发布时间 :
字数:2.4k 阅读 :
  1. 1. 概述 (Overview)
  2. 2. 设定具体维度
  3. 3. 攻击全流程深度解析 (In-Depth Attack Walkthrough)
    1. 3.1 初始化与设置 (Initialization & Setup)
    2. 3.2 迭代循环开始 (Entering the Main Loop)
    3. 3.3 步骤 A: 计算梯度 (Step A: Gradient Calculation)
    4. 3.4 步骤 B: 寻找最佳替换 Token (Step B: Finding the Best Replacement Tokens)
    5. 3.5 步骤 C: 随机采样与贪婪更新 (Step C: Randomized Sampling & Greedy Update)
  4. 4. 总结
  5. 代码

文章中主体内容通过gemini code进行整理 本人尝试基于代码与gemini不断交流来了解这种攻击的原理，实际来看效果不错，将最终的文档整理并分享。个人的疑惑是如何反向传播来实现后缀更新，学习后发现其核心思想是 梯度代表更新方向，梯度与词表embedding的点积代表选用该token带来的贡献。

GCG 攻击是一种针对白盒大模型的攻击方式，论文https://arxiv.org/abs/2307.15043

本文档通过一组具体的维度数值，对 GCG 攻击代码的实现进行最终的深度技术拆解。这能让你像在调试器中观察一样，清晰地看到每个张量（Tensor）的形状变化和计算过程。

1. 概述 (Overview)

GCG 攻击通过在输入嵌入空间中进行梯度下降，来迭代式地优化一个对抗性后缀。其目标是找到一个能让模型以最高概率生成指定有害回复的后缀。本解析将带你走过从 token ID 到下一个更优后缀的完整计算过程。

2. 设定具体维度

为了方便理解，我们在整个流程中假设以下维度：

  • embed_dim (词嵌入维度): 768
  • vocab_size (词汇表大小): 32000
  • user_len (用户指令长度): 10 tokens
  • adv_len (对抗性后缀长度): 5 tokens
  • target_len (目标回复长度): 8 tokens
  • total_length (总长度): 10 + 5 + 8 = 23 tokens
  • topk (候选词数量): 20
  • batch_size (候选后缀评估批次): 64

3. 攻击全流程深度解析 (In-Depth Attack Walkthrough)

我们将模拟并分解攻击循环中的一次完整迭代。

3.1 初始化与设置 (Initialization & Setup)

准备工作与之前相同，但现在我们知道分词后各个部分的 token 数量。

3.2 迭代循环开始 (Entering the Main Loop)

代码进入主循环: for step in range(num_steps):

3.3 步骤 A: 计算梯度 (Step A: Gradient Calculation)

  • A.1. 拼接与嵌入 (Concatenation & Embedding)

    • 操作: 拼接 token ID。
      # input_ids shape: (1, 10 + 5 + 8) -> (1, 23)
    • 操作: 转换为词嵌入向量，并启用梯度计算。
      # input_embeds shape: (1, 23, 768)

  • A.2. 前向传播 (Forward Pass)

    • 操作: 将 input_embeds 输入 LLM，得到 logits。
      # logits shape: (1, 23, 32000)

  • A.3. 损失计算 (Loss Calculation)

    • 操作: 切片 logits，只取出用于预测 target 的部分。预测 target 的第一个 token，需要用到第 10+5-1=14 个位置的 logit。预测 target 的最后一个 token，需要用到第 23-1-1=21 个位置的 logit。
      # 切片范围: logits[:, 14:22, :]
# pred_logits shape: (1, 8, 32000)
    • 操作: 计算交叉熵损失。pred_logits 会被 reshape 成 (8, 32000)，target_prompt_tokens 会被 reshape 成 (8,)。

  • A.4. 反向传播 (Backward Pass)

    • 操作: loss.backward()
    • 操作: 从 input_embeds.grad 中提取出只属于对抗性后缀的梯度。对抗性后缀从第 10 个 token 开始，到第 14 个 token 结束。
      # 切片范围: input_embeds.grad[:, 10:15, :]
# adv_grad shape: (1, 5, 768)
    • 结果: 我们获得了 adv_grad，一个 (1, 5, 768) 的张量，包含了宝贵的梯度信息。

3.4 步骤 B: 寻找最佳替换 Token (Step B: Finding the Best Replacement Tokens)

  • B.1. 梯度与嵌入矩阵相乘 (Gradient-Embedding Product)

    • 目标: 将 768 维的梯度信号，转换为 32000 维的 token 分数。
    • 操作: adv_grad 左乘词汇表嵌入矩阵的转置。
      # model.get_input_embeddings().weight.T shape: (768, 32000)
# token_scores = adv_grad @ model.get_input_embeddings().weight.T
# shape: (1, 5, 768) @ (768, 32000) -> (1, 5, 32000)
    • 结果解读: token_scores 是一个 (1, 5, 32000) 的张量。它的 [0, i, j] 位置的值，代表了将后缀的第 i 个 token 换成词汇表的第 j 个 token 的“收益”。

  • B.2. Top-K 筛选 (Top-K Selection)

    • 操作: 对 token_scores 在最后一个维度（32000）上进行排序，选出分数最高的 20 个。
      # topk_indices = torch.topk(token_scores, k=20, dim=-1).indices
# topk_indices shape: (1, 5, 20)
    • 结果: topk_indices 保存了后缀每个位置最值得替换的 20 个候选 token 的 ID。

3.5 步骤 C: 随机采样与贪婪更新 (Step C: Randomized Sampling & Greedy Update)

  • C.1. 生成候选批次 (Generating a Candidate Batch)

    • 操作: 循环 64 次，每次都通过随机替换一个 token，生成一个新的后缀。最终得到一个包含 64 个新后缀的列表。

  • C.2. 并行评估 (Parallel Evaluation)

    • 操作: 将 64 个新后缀与 user_prompt 和 target_prompt 分别拼接，形成一个批次。
      # 伪代码
# batched_input_ids shape: (64, 23)
    • 操作: 对这个批次执行一次无梯度的前向传播，并计算每个样本的损失。
      # batched_logits shape: (64, 23, 32000)
# ... 经过切片和损失计算后 ...
# all_losses shape: (64,)

  • C.3. 贪婪选择 (Greedy Selection)

    • 操作: 在 all_losses (一个有64个损失值的向量) 中找到最小值的索引。
      # best_index = torch.argmin(all_losses) -> 一个从 0 到 63 的整数
    • 操作: 将这个 best_index 对应的新后缀，作为下一次迭代的起始后缀。

循环回到 3.2，用这个从64个候选中脱颖而出的、当前最优的后缀，重复上述所有步骤。

4. 总结

通过这组具体的数值，我们可以更清晰地追踪数据在算法中的流动和变换。整个过程是一个高度优化的搜索循环，它将连续的梯度信息巧妙地转化为了在离散 token 空间中进行选择的依据，从而高效地找到攻击向量。

代码

为了更深入的了解原理，这是Gemini生成代码，可以直接运行测试学习

import torch
import torch.nn.functional as F

# --- 1. 配置与维度设定 ---
# (与我们之前在 README 中讨论的维度保持一致)
print("--- 1. 配置与维度设定 ---")
embed_dim = 768       # 词嵌入维度
vocab_size = 32000    # 词汇表大小
user_len = 10         # 用户指令长度
adv_len = 5           # 对抗性后缀长度
target_len = 8        # 目标回复长度
total_len = user_len + adv_len + target_len
topk = 20             # 候选词数量

print(f"Embed Dim: {embed_dim}, Vocab Size: {vocab_size}")
print(f"User Len: {user_len}, Adv Len: {adv_len}, Target Len: {target_len}\n")


# --- 2. 随机初始化模拟数据 ---
# 我们不需要一个真实的LLM，只需要形状正确并能进行反向传播的模拟组件
print("--- 2. 随机初始化模拟数据 ---")

# 模拟模型的词嵌入层
mock_embedding_layer = torch.nn.Embedding(vocab_size, embed_dim)
print(f"模拟词嵌入矩阵 Shape: {mock_embedding_layer.weight.shape}")

# 模拟一个极简的“模型”，它只做一件事：将嵌入向量转换为logits
# 这足以让我们完成梯度的计算
mock_llm_head = torch.nn.Linear(embed_dim, vocab_size)
print(f"模拟LLM头部 Shape: W: {mock_llm_head.weight.shape}, B: {mock_llm_head.bias.shape}")

# 创建随机的输入 token IDs
user_prompt_ids = torch.randint(0, vocab_size, (1, user_len))
adv_prompt_ids = torch.randint(0, vocab_size, (1, adv_len))
target_ids = torch.randint(0, vocab_size, (1, target_len))

print(f"原始对抗性后缀 (Token IDs): {adv_prompt_ids.tolist()[0]}\n")


# --- 3. 执行单步 GCG 攻击流程 ---
print("--- 3. 执行单步 GCG 攻击流程 ---")

# A.1. 拼接与嵌入
input_ids = torch.cat([user_prompt_ids, adv_prompt_ids, target_ids], dim=1)
input_embeds = mock_embedding_layer(input_ids)
# 关键！input_embeds是一个中间变量（非叶子节点），默认不会保存梯度。
# 我们需要使用 retain_grad() 来显式地保留它的梯度。
input_embeds.retain_grad()
print(f"A.1. 拼接后输入 Shape: {input_ids.shape}")
print(f"     嵌入后输入 Shape: {input_embeds.shape}")

# A.2. 前向传播
logits = mock_llm_head(input_embeds)
print(f"A.2. 模型输出 Logits Shape: {logits.shape}")

# A.3. 损失计算 (只关心 target 部分)
pred_logits = logits[:, user_len + adv_len - 1 : -1, :]
loss = F.cross_entropy(pred_logits.view(-1, vocab_size), target_ids.view(-1))
print(f"A.3. 用于计算损失的 Logits Shape: {pred_logits.shape}")
print(f"     计算出的 Loss: {loss.item():.4f}")

# A.4. 反向传播 & 提取梯度
loss.backward()
adv_grad = input_embeds.grad[:, user_len : user_len + adv_len, :]
print(f"A.4. 提取出的对抗性后缀梯度 Shape: {adv_grad.shape}\n")


# --- 4. 确定下一个 Token ---
print("--- 4. 确定下一个 Token ---")

# B.1. 梯度与嵌入矩阵相乘
token_scores = adv_grad @ mock_embedding_layer.weight.T
print(f"B.1. 梯度与嵌入矩阵转置相乘...")
print(f"     ({adv_grad.shape}) @ ({mock_embedding_layer.weight.T.shape}) -> {token_scores.shape}")

# B.2. Top-K 筛选
# 梯度指向损失增大的方向，所以我们要找分数最小（最负）的token
# 这等价于对分数的负数取 top-k
neg_token_scores = -token_scores
_, topk_indices = torch.topk(neg_token_scores, k=topk, dim=-1)
print(f"B.2. Top-K 筛选结果 Shape: {topk_indices.shape}")
print(f"     这代表后缀的5个位置，每个位置都有了{topk}个最佳候选Token。\n")


# --- 5. 贪婪地选择并更新一个Token ---
# 为了简化Demo，我们不实现批处理评估，而是直接选择全局最优的那个替换
print("--- 5. 贪婪地选择并更新一个Token ---")

# 将 (5, 32000) 的分数展平为 (5 * 32000)
flat_scores = neg_token_scores.view(-1)

# 找到分数最高的那个候选词的索引
best_candidate_flat_idx = torch.argmax(flat_scores)

# 将一维索引转换回 (位置, token_id)
# .item() 返回的是一个通用的Python Number，静态检查器无法确定其为整数。
# 我们通过int()进行显式转换，以确保类型正确。
best_pos = int((best_candidate_flat_idx // vocab_size).item())
best_token_id = int((best_candidate_flat_idx % vocab_size).item())

print(f"全局最优替换: 在后缀的第 {best_pos} 个位置，将原Token替换为新Token ID {best_token_id}")

# 创建新的对抗性后缀
new_adv_prompt_ids = adv_prompt_ids.clone()
new_adv_prompt_ids[0, best_pos] = best_token_id

print(f"原始对抗性后缀: {adv_prompt_ids.tolist()[0]}")
print(f"更新后对抗性后缀: {new_adv_prompt_ids.tolist()[0]}\n")

print("--- Demo结束 ---")
print("在真实的攻击中，这个新的后缀将会作为下一次迭代的输入，重复上述过程。")

文章参考:

博客地址: qwrdxer.github.io

欢迎交流, QQ: 1944270374. WX: qwrdxer

转载请注明来源，欢迎对文章中的引用来源进行考证，欢迎指出任何有错误或不够清晰的表达。可以在下面评论区评论，也可以邮件至 1944270374@qq.com

©2020-2023 qwrdxer

Built with Hexo and 3-hexo theme