【MoE】Buffer Overflow in Mixture of Experts

From:https://arxiv.org/abs/2402.05526v1

威胁模型和攻击方法

漏洞所在

该攻击利用了某些 MoE 实现中的两个关键特性：

• 有限的专家缓冲区：每个专家子网络在处理一个批次（batch）的数据时，只能处理固定数量的输入令牌（token）。例如，一个专家的缓冲区容量可能只有 38 个令牌。一旦满了，就无法再接收新的令牌。
• 依赖于批次顺序的路由策略（Batch-Dependent Routing）：论文中提到的“传统路由策略”（vanilla routing strategy）会按照令牌在批次中的顺序，依次将它们分配给专家。这意味着，排在批次前面的输入会优先获得它们最想去的专家。而排在后面的输入，可能会发现它想去的专家已经被“占满”了。
  
  将一个批次中的所有令牌（总计B*T个，其中B是批次大小，T是每个样本的令牌数）分
  配给n个专家中的一个或多个（具体是k个)
  

那么问题就是出现了，假设 k=1，一个批次中有攻击者数据和受害者数据。

• 攻击者的令牌排在前面，它们首先被处理。它们计算出门控得分，并被分配到它们的首选专家，比如 专家A，并开始填满 专家A 的缓冲区。
• 轮到处理受害者的令牌时，它也计算出门控得分，发现自己的首选专家也是 专家A。
• 但此时，专家A 的缓冲区可能已经被攻击者的令牌占满了。根据系统设置，受害者的这个令牌要么被丢弃，要么被强制分配给一个次优的专家，从而导致计算错误。

攻击的实现步骤 Threat model and attack method

论文中展示了一种黑盒随机搜索攻击 Algorithm 2 Random search attack。攻击者无需知道模型的内部权重，只需要观察模型的最终输出即可。这是一种相对简单但可能非常耗时的攻击方式。

让我们以论文中的例子来具体说明：受害者的输入是 $X^{\ast }$= “Solve the following equation: 1+1=”。模型本应预测出的下一个词是 “2”。攻击者的目标是让模型输出一个错误答案，比如 “1”。

以下是攻击的完整流程：

第一步：控制批次（Batch）

攻击者会构造自己的输入数据，称为“对抗性数据” $X^A$。然后，他们设法让自己的数据$X^A$和受害者的数据$X^{\ast }$ 被放入同一个批次中进行处理。最关键的一点是，攻击者会将受害者的输入$X^{\ast }$ 放在批次的最后一个位置。

• 批次结构：[对抗性数据1, 对抗性数据2, …, 受害者数据 ($X^{\ast }$ )]

第二步：识别并锁定“最佳”专家

模型的门控网络会将令牌路由到最适合处理它的专家那里。对于 “1+1=” 这样的算术问题，模型自然会将其中的令牌发送给擅长数学计算的专家。攻击者的目标就是阻止受害者的令牌访问这个 算术专家。

第三步：构造恶意数据以填满缓冲区

这是攻击的核心。攻击者需要找到一种特殊的对抗性数据 $X^A$，它的特性是：它所包含的令牌也会被门控网络大量路由到 算术专家 那里。

由于攻击者的数据位于批次的前部，它的令牌会首先被处理。这些令牌被一个个地分配给 算术专家，逐渐填满其有限的缓冲区。

第四步：“溢出”与强制重路由

当系统处理到批次末尾的受害者输入$X^{\ast }$ 时，那些本应被送到 算术专家 的令牌发现，这个专家的缓冲区已经被攻击者的令牌完全占满了。

此时，系统有两种选择（论文的实验假设是第一种）：

1. 丢弃令牌（Drop the token）：该令牌不被任何专家处理。
2. 重路由令牌（Reroute the token）：将该令牌发送到另一个缓冲区未满的、非最优的专家那里。

第五步：诱发错误

这种强制的重路由或令牌丢失，破坏了模型的正常计算流程。由于 “1+1=” 的关键令牌没有被专业的 算术专家 处理，模型的最终计算结果被干扰，从而输出一个错误的答案（例如，“1” 而不是 “2”）。

攻击算法分析

算法输入参数说明

• l: 损失函数（Loss function）。这是攻击的核心评估标准。它衡量了攻击的“成功程度”。在论文中，这个损失函数被定义为即正确类别 y 的输出概率与所有其他错误类别中的最高概率之差。攻击者的目标是最小化这个损失值。当这个值变为负数时，意味着某个错误答案的概率超过了正确答案，攻击成功。

• M: 迭代总次数，即攻击者愿意尝试多少次来优化其对抗性数据。

• r: 每次迭代中要改变的令牌数量。

• x*: 攻击的目标输入样本。

• y: x* 对应的正确输出。

• V: 词汇表（Vocabulary），即所有可用令牌的集合。

算法伪代码

当循环 M 次结束后，算法返回在整个搜索过程中找到的、能够产生最低损失值（即攻击效果最好）的对抗性数据集。

这个算法的本质就是一种非常朴素的爬山算法。这种方法的优点是实现简单，且适用于黑盒场景。缺点是效率极低，可能会陷入局部最优解，并且需要大量的模型查询（M 次），在实际应用中成本很高。

模拟实验 Attack demonstration

为了进行演示，作者们精心设计了实验环境和参数：

模拟参数设置

1. 目标模型 (Target Model)

   • Mixtral-8x7B：这是一个非常著名的大型 MoE 模型。它拥有 8 个专家（per layer），每个token最多可以被路由到 2 个专家。词汇表大小为 32,000。
   • 关键修改：原版的 Mixtral-8x7B 为了性能，默认没有设置专家缓冲区的容量上限。这意味着它本身对这种攻击是免疫的。为了模拟存在漏洞的环境，作者们手动为模型增加了一个缓冲区容量限制。他们将每个专家的缓冲区容量设置为 38 个token。

2. 目标输入 $x^{\ast }$
   $x^{\ast }$= “Solve the following equation: 1+1=”
   其序列长度为 12 个 token，在没有攻击的情况下，模型预测的下一个最可能的令牌是 “2”，其 softmax 概率为 25.9%。

   # 关于12 个 token如何计算的代码
   from transformers import AutoTokenizertokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
   text = "Solve the following equation: 1+1="# 进行分词
   tokens = tokenizer.tokenize(text)
   token_ids = tokenizer.encode(text)print(f"原始文本: {text}")
   print(f"分词结果: {tokens}")
   print(f"令牌数量: {len(tokens)}")# encode 会加上起始符 <s>
   print(f"\n编码后的ID: {token_ids}")
   print(f"编码后的ID数量 (包含起始符): {len(token_ids)}")
   

   原始文本: Solve the following equation: 1+1=
   分词结果: ['Sol', 've', ' the', ' following', ' equation', ':', ' ', '1', '+', '1', '=']
   令牌数量: 11
   编码后的ID: [1, 2034, 323, 272, 5974, 4218, 28747, 28705, 28723, 28748, 28723, 28781]
   编码后的ID数量 (包含起始符): 12
   

3. 批次设置 (Batch Setting)：
   批次大小 B = 8。这意味着一个批次中总共有 8 个样本。 其中 1 个是受害者的输入 $x\ast$， 其余 7 个是攻击者控制的对抗性数据 $X^A$。总令牌数：8 个样本 * 12 令牌/样本 = 96 个令牌。

4. 攻击参数 (Attack Parameters)：
   迭代次数 M = 1000 次。 每次迭代修改 r = 5 个令牌（在每个对抗性样本中）。总共修改 7 * 5 = 35 个令牌。

实验过程与结果分析

• 蓝色实线 (Correct token)：代表正确答案 “2” 的概率，初始值为约 24%。
• 橙色虚线 (Incorrect token)：代表当时第二可能的答案 “1” 的概率，初始值约为 20%。

缓解措施

在讨论缓解措施之前，我们首先要回顾一下攻击能够成功的两个根本原因：

1. 有限的专家缓冲区容量（Finite expert buffer capacity）：这是“溢出”得以发生的前提。
2. 依赖于批次顺序的、确定性的专家路由（Batch-dependent, deterministic expert routing）：这使得攻击者可以通过控制批次顺序和内容来精确地操纵路由结果。

因此，所有的缓解措施都旨在打破这两个条件之一或两者。

1. 移除缓冲区容量限制 (The Ultimate Fix)

• 方法：最直接、最彻底的解决方法就是完全取消专家缓冲区的容量上限，或者将其设置为批次中的总令牌数。
• 效果：攻击将完全失效。因为缓冲区永远不会“溢出”，攻击者也就无法通过填充缓冲区来阻止其他令牌的访问。
• 缺点：这种方法违背了使用缓冲区限制的初衷。它可能导致硬件资源利用不均和计算浪费，因为某些专家可能会接收到远超平均数量的令牌，而其他专家则处于空闲状态。这在性能和成本上是不可取的。

2. 随机化批次顺序 (Randomize batch order)

• 方法：在将一批请求送入模型处理之前，先将批次内各个请求的顺序随机打乱。
• 原理：这种方法直接攻击了漏洞的第二个根源。攻击者无法再保证他们的恶意数据 X^A 一定会排在受害者数据 x* 的前面。如果 x* 随机排到了前面，它的令牌就能优先访问专家，攻击就失败了。
• 效果：极大地降低了攻击的成功率和效率。攻击者可能需要进行海量尝试，才能“幸运地”排在目标前面。
• 额外建议：对于特别敏感的查询，可以进行两次推理，每次使用不同的随机顺序。如果两次的输出结果差异很大，就说明可能受到了攻击，需要进行人工检查。

3. 使用大的容量松弛因子 (Use a large capacity slack, C)

• 方法：这是在性能和安全之间的一个折中方案。不完全取消限制，而是将缓冲区容量设置得足够大。
• 原理：如第 6 节所示，当缓冲区大到一定程度时，攻击者即使控制了批次中的大部分数据，也难以将其完全填满。
• 效果：能够有效防御攻击。当 C 值足够大时（例如，在实验中 C>=2），攻击就会失效。

4. 采样门控权重而非选择 Top-k (Sampling from gate weights instead of selecting the top-k)

• 方法：改变专家路由的根本机制。原始方法是确定性的：总是选择门控网络输出概率最高的 k 个专家。新方法是随机性的：将门控网络的输出（一个概率分布）视为一个抽奖盘，从中随机采样 k 个专家。
• 原理：这种方法为路由过程引入了随机性。攻击者即使能够预测门控网络的输出概率，也无法 100% 确定令牌最终会被路由到哪个专家。
• 效果：使得攻击变得非常困难。攻击者不知道应该去“堵塞”哪个专家的缓冲区，因为目标是随机变化的。他们可能需要同时堵塞多个专家，这大大增加了攻击成本和难度。