MBTSAD: Mitigating Backdoors in Language Models Based on Token Splitting and Attention Distillation

📄 arXiv: 2501.02754v1 📥 PDF

作者: Yidong Ding, Jiafei Niu, Ping Yi

分类: cs.CR, cs.CL

发布日期: 2025-01-06

备注: Accepted by ICTAI 2024

💡 一句话要点

提出MBTSAD，一种无需预训练权重即可缓解语言模型后门攻击的方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 后门攻击防御 语言模型安全 注意力蒸馏 Token Splitting Out-of-Distribution数据

📋 核心要点

  1. 现有后门防御方法依赖预训练权重，但在预训练权重不可用时失效，限制了应用场景。
  2. MBTSAD通过token splitting生成数据并进行重训练，再利用注意力蒸馏缓解后门攻击，无需预训练权重。
  3. 实验表明，MBTSAD在缓解后门攻击的同时，保持了在干净数据上的性能，且不依赖预训练权重。

📝 摘要（中文）

近年来，基于注意力机制的模型在各个领域表现出色，但仍然容易受到后门攻击，通常源于下载或在中毒数据集上进行微调。许多现有的NLP模型后门缓解方法依赖于预训练（未微调）的权重，但这些方法在预训练权重不可用的情况下会失效。本文提出了MBTSAD，它仅利用少量干净数据即可缓解语言模型中的后门，并且不需要预训练权重。具体来说，MBTSAD在token splitting生成的数据集上重新训练后门模型。然后，MBTSAD利用注意力蒸馏，将重新训练的模型作为教师模型，原始后门模型作为学生模型。实验结果表明，MBTSAD实现了与基于预训练权重的方法相当的后门缓解性能，同时保持了在干净数据上的性能。MBTSAD不依赖于预训练权重，增强了其在预训练权重不可访问的场景中的实用性。此外，我们简化了对抗训练的min-max问题，并可视化文本表示，发现MBTSAD第一步中的token splitting方法生成了Out-of-Distribution (OOD)数据，从而使模型学习更通用的特征并消除后门模式。

🔬 方法详解

问题定义：现有NLP模型的后门防御方法通常依赖于预训练权重，这在预训练权重不可用或难以获取的情况下（例如，模型私有化部署或使用特定领域的预训练模型）构成了严重的限制。因此，如何在没有预训练权重的情况下有效地缓解语言模型中的后门攻击是一个关键问题。

核心思路：MBTSAD的核心思路是，首先通过token splitting生成具有一定扰动的数据，迫使模型学习更鲁棒的特征，从而削弱后门触发器的影响。然后，利用注意力蒸馏，将重训练后的模型作为教师模型，指导原始后门模型学习，进一步消除后门模式。这种方法旨在使模型忘记后门，同时保留其在干净数据上的性能。

技术框架：MBTSAD主要包含两个阶段：1) Token Splitting重训练：使用token splitting方法对少量干净数据进行处理，生成新的训练数据集。然后，使用该数据集重新训练原始的后门模型。2) 注意力蒸馏：将重训练后的模型作为教师模型，原始后门模型作为学生模型，利用注意力蒸馏损失函数，使学生模型学习教师模型的注意力分布，从而消除后门模式。

关键创新：MBTSAD的关键创新在于它提出了一种无需预训练权重的后门缓解方法。通过token splitting生成OOD数据，迫使模型学习更通用的特征，从而削弱后门触发器的影响。此外，利用注意力蒸馏，进一步提升了模型的鲁棒性。这种方法在预训练权重不可用的场景下具有重要的实用价值。

关键设计：Token splitting的具体方法是将token分割成多个子token，例如将“apple”分割成“ap”, “pl”, “e”。注意力蒸馏损失函数通常采用KL散度或MSE等方法，衡量教师模型和学生模型注意力分布的差异。实验中，需要选择合适的token splitting比例和蒸馏温度等超参数，以达到最佳的缓解效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MBTSAD在不依赖预训练权重的情况下，实现了与依赖预训练权重的方法相当的后门缓解性能。实验结果表明，MBTSAD在保持干净数据性能的同时，显著降低了后门攻击的成功率。此外，论文还通过可视化文本表示，验证了token splitting方法能够生成OOD数据，从而使模型学习更通用的特征。

🎯 应用场景

MBTSAD可应用于各种NLP任务中，尤其适用于预训练权重不可用的场景，例如私有化部署的语言模型、特定领域的定制模型等。该方法能够有效缓解后门攻击，提高模型的安全性，保障用户数据的隐私和安全。此外，该方法还可以作为一种通用的模型鲁棒性提升技术，应用于其他对抗攻击的防御。

📄 摘要（原文）

In recent years, attention-based models have excelled across various domains but remain vulnerable to backdoor attacks, often from downloading or fine-tuning on poisoned datasets. Many current methods to mitigate backdoors in NLP models rely on the pre-trained (unfine-tuned) weights, but these methods fail in scenarios where the pre-trained weights are not available. In this work, we propose MBTSAD, which can mitigate backdoors in the language model by utilizing only a small subset of clean data and does not require pre-trained weights. Specifically, MBTSAD retrains the backdoored model on a dataset generated by token splitting. Then MBTSAD leverages attention distillation, the retrained model is the teacher model, and the original backdoored model is the student model. Experimental results demonstrate that MBTSAD achieves comparable backdoor mitigation performance as the methods based on pre-trained weights while maintaining the performance on clean data. MBTSAD does not rely on pre-trained weights, enhancing its utility in scenarios where pre-trained weights are inaccessible. In addition, we simplify the min-max problem of adversarial training and visualize text representations to discover that the token splitting method in MBTSAD's first step generates Out-of-Distribution (OOD) data, leading the model to learn more generalized features and eliminate backdoor patterns.