利用AI检测基于prompt的AI攻击：全面教学指南

前言

随着大型语言模型(LLM)的广泛应用，基于prompt的AI攻击(如注入攻击、越狱攻击)日益增多。本文详细讲解如何利用AI技术检测这类恶意prompt，涵盖机器学习方法和大型语言模型方法。

数据集准备

我们使用来自Hugging Face的deepset/prompt-injections数据集：

• 包含数百个正常和被标记为注入的操纵提示样本
• 主要包含英文提示，部分翻译成其他语言(主要是德语)
• 已预先分为训练集和测试集

数据集结构：

• 两列：prompt文本本身和标签(0表示良性，1表示恶意)
• 标签分布不平衡，良性样本占主导地位

机器学习检测方法

1. 文本向量化

由于数据集包含多语言样本，我们使用multilingual BERT获取嵌入表示：

from transformers import BertTokenizer, BertModel
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-uncased')
model = BertModel.from_pretrained('bert-base-multilingual-uncased')

def get_bert_embedding(prompt):
    tokens = tokenizer(prompt, return_tensors='pt', truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**tokens)
    last_hidden_states = outputs.last_hidden_state
    embedding_vector = last_hidden_states.mean(dim=1).squeeze().numpy()
    return embedding_vector

2. 应用分类算法

我们测试四种经典机器学习算法：

from sklearn.naive_bayes import GaussianNB
from sklearn.linear_model import LogisticRegression
from sklearn import svm
from sklearn.ensemble import RandomForestClassifier

estimators = [
    ("Naive Bayes", GaussianNB()),
    ("Logistic Regression", LogisticRegression()),
    ("Support Vector Machine", svm.SVC()),
    ("Random Forest", RandomForestClassifier())
]

3. 训练与评估

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

results = pd.DataFrame(columns=["accuracy", "precision", "recall", "f1 score"])

for est_name, est_obj in estimators:
    est_obj.fit(X_train, y_train)
    y_predict = est_obj.predict(X_test)
    
    accuracy = accuracy_score(y_test, y_predict)
    precision = precision_score(y_test, y_predict)
    recall = recall_score(y_test, y_predict)
    f1 = f1_score(y_test, y_predict)
    
    results.loc[est_name] = [accuracy, precision, recall, f1]

4. 结果分析

测试结果对比：

关键发现：

• 逻辑回归和支持向量机表现最佳
• 三个模型实现了100%的精确率(无假阳性)
• 错误主要集中在召回率指标(漏检)

大型语言模型检测方法

1. 使用预训练XLM-RoBERTa

RoBERTa关键特点：

• 动态掩码策略
• 更大的训练数据集
• 更长的序列长度
• 优化的学习率调度

实现代码：

from transformers import pipeline

classifier = pipeline("zero-shot-classification", 
                     model="joeddav/xlm-roberta-large-xnli")

def classify_prompt(prompt):
    candidate_labels = ["Injection", "Normal"]
    output = classifier(prompt, candidate_labels)
    return 1 if output['labels'][0] == "Injection" else 0

2. 评估结果

直接使用预训练模型效果：

• 准确率: 0.78
• 精确率: 0.75
• 召回率: 0.67
• F1分数: 0.71

不如机器学习方法效果好。

微调大型语言模型

1. 准备数据集类

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

2. 设置训练参数

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="../output",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=5,
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    logging_dir="../output/logs",
)

3. 定义评估函数

from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def evaluate_model(trainer, epoch):
    predictions, labels = trainer.predictions.argmax(axis=1), trainer.label_ids
    accuracy = accuracy_score(labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average="binary")
    
    global results_df
    results_df.loc[len(results_df)] = [epoch, accuracy, precision, recall, f1]
    
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }

4. 训练与评估

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=lambda p: evaluate_model(p, trainer.state.epoch),
)

trainer.train()

5. 微调结果

微调后模型性能：

• 准确率: 0.84
• 比直接使用预训练模型效果好
• 但仍略逊于传统机器学习方法

结论与建议

1. 方法选择：

   • 传统机器学习方法(特别是逻辑回归和SVM)在本任务中表现最佳
   • 微调大模型效果次之，但可能具有更好的泛化能力
   • 直接使用预训练大模型效果最差

2. 实践建议：

   • 对于实时性要求高的场景，优先考虑传统机器学习方法
   • 对于多语言或复杂场景，可考虑微调大模型
   • 结合两种方法可能获得更好的检测效果

3. 改进方向：

   • 收集更多恶意prompt样本改善数据不平衡
   • 尝试集成学习方法结合多种模型优势
   • 探索prompt-specific的特征工程方法

参考资源

1. deepset官方
2. Prompt注入数据集
3. BERT原始论文
4. RoBERTa介绍
5. RoBERTa优化细节