前言

对于利用prompt实现的注入攻击、越狱攻击等大家应该都比较熟悉了。本文将主要关注，如何可以实现对这类恶意prompt的检测。

这种攻击旨在利用精心设计的输入提示来操控LLM，使其忽略之前的指令，从而执行非预期的操作。

我们在本文中使用的数据集来自hugging face：https://huggingface.co/datasets/deepset/prompt-injections

这个是由deepset提供的提示注入数据集。该数据集结合了数百个正常和被标记为注入的操纵提示样本。它主要包含英文提示，还有一些翻译成其他语言的提示，主要是德语。原始数据集已经分为训练集和保留子集。

利用机器方法检测

首先我们使用机器学习方法检测，加载数据集

原始数据集已经分为训练集和保留集。

一种方法是将这两个部分合并成一个统一的集合，然后使用随机静态的训练/测试拆分或交叉验证技术来重新生成测试样本并评估方法的性能。由于我们计划进行多次实验，我们在这些实验中保持这一原始拆分，以便使用统一的测试基准进行结果比较。

分别查看训练集和验证集的前几条数据

查看数据集的基本信息

从上图可以看出，数据集的结构非常简单，包含两列：一列是提示文本本身，另一列是指示提示是否包含注入的标签。即表明这个prompt是否是恶意的

绘图查看它们的分布比例

上图展示了数据样本中各标签类别的分布，可以看到良性样本占主导地位，恶意提示的样本较少。

尽管这两个标签不平衡，但正样本的数量仍然相当充足。

根据数据集创建者的说法，他们加入了一些非英语样本，以扩展提示注入攻击的范围到其他语言。这也有助于创建不受语言限制的通用检测模型，从而更能防御不同语言的恶意提示。

鉴于此，传统的预训练语言特定的Word2Vec嵌入不适合用于检测，因为我们需要能够在同一向量空间中表示多种语言文本的嵌入模型。

因此，我们采用了一种较新的语言模型，即multilingual BERT来获取嵌入。

这里简单介绍一下BERT。

BERT（Bidirectional Encoder Representations from Transformers）是一种预训练的自然语言处理（NLP）模型，由Google在2018年提出。BERT模型的核心特点是其双向训练机制，它允许模型同时考虑输入文本中单词的左侧和右侧上下文，这与传统的单向语言模型不同。BERT基于Transformer模型，这是一种依赖于自注意力机制的架构，它允许模型在处理序列数据时更加灵活和高效。BERT通过两个主要的预训练任务来学习语言表示，即Masked Language Model (MLM) 和 Next Sentence Prediction (NSP)。在预训练完成后，BERT模型可以在特定的下游任务上进行微调，如文本分类、问答系统、命名实体识别等，以适应这些任务的具体需求。

如下的第一行代码从指定的路径加载了一个预训练的多语言BERT分词器（bert-base-multilingual-uncased）。分词器的作用是将输入的文本转换为BERT模型可以理解的格式，即将文本分割成词片段（tokens），并将这些词片段映射到相应的词汇表索引。

第二行代码从同一路径加载了一个预训练的多语言BERT模型（bert-base-multilingual-uncased）。BERT模型的作用是接收分词器处理后的输入，将其转换为向量表示（embeddings），这些向量表示可以用于下游的自然语言处理任务。

如下代码定义了一个函数get_bert_embedding(prompt)，用于获取给定提示（prompt）的BERT嵌入表示。

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

分词和向量化：
函数首先调用分词器，将输入的提示（prompt）转换成模型所需的张量格式。分词器将文本分割成词片段，并将其转换为张量，同时进行截断和填充，以确保所有输入序列的长度相同。
模型推理：
使用上下文管理器torch.no_grad()，关闭梯度计算以节省内存和计算资源，并防止模型参数更新。然后将分词后的输入张量传递给BERT模型，获取模型的输出。
获取最后的隐藏状态：
从模型输出中提取最后一层隐藏状态，这些隐藏状态是每个输入词片段在BERT模型最后一层的表示。
计算平均嵌入向量：
对最后一层的隐藏状态在第一个维度上进行平均，得到整个序列的平均嵌入向量。然后，将该向量从PyTorch张量转换为NumPy数组，并压缩成一维。
返回嵌入向量：
最后，函数返回计算得到的嵌入向量。

将训练集和测试集中的每个提示（prompt）转换为其对应的BERT嵌入表示，然后将这些嵌入表示存储在数据集的一个新列中。

data_train['embedding'] = data_train['prompt'].apply(get_bert_embedding)
data_test['embedding'] = data_test['prompt'].apply(get_bert_embedding)

现在可以来看看prompt、标签、嵌入放在一起是什么样子的

注意到"embedding"列，每个单元格包含代表相应提示的嵌入向量。在解包成多个独立列之后，该列的内容连同类别"label"将作为分类算法的输入。

在继续训练过程之前，我们首先需要转移先前的DataFrame并提取训练和测试子集。由于我们的数据集已经分成了训练和测试子集，我们只需构建自变量的结构并准备 X 和 y 的拆分即可。

接着使用四类典型的机器学习算法

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())
]

然后就直接训练、测试并计算指标

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]：

导入评估指标：
从Scikit-learn库中导入了四个评估指标的函数：准确率（accuracy_score）、精确率（precision_score）、召回率（recall_score）、F1分数（f1_score）。这些指标用于衡量分类器在预测结果中的不同方面的性能。
初始化结果存储DataFrame：
创建了一个名为results的空DataFrame，它有四列："accuracy"、"precision"、"recall"和"f1 score"，用于存储每个分类器在测试集上计算得到的性能指标。
循环遍历分类器列表：
使用一个for循环遍历名为estimators的分类器列表（或迭代器）。每次迭代中，est_name是分类器的名称或标识，est_obj是实际的分类器对象。
训练分类器：
对当前迭代的分类器对象（est_obj）调用fit方法，用训练集（X_train和y_train）进行训练。这一步将分类器拟合到训练数据上，使其学习如何预测目标变量。
生成预测并计算评估指标：
- 使用训练好的分类器对测试集（X_test）进行预测，得到预测结果 y_predict。
- 分别计算预测结果与真实标签 y_test 之间的准确率、精确率、召回率和F1分数。这些指标通过调用导入的评估函数计算得出。
将结果存储到DataFrame：
使用分类器的名称 est_name 作为索引，将计算得到的准确率、精确率、召回率和F1分数存储到results DataFrame中的对应行。每个分类器的性能指标将会被记录和比较，帮助选择最合适的模型或调优参数。

查看测试结果

可以看到：

大多数考虑的模型在测试数据样本上表现相对良好。
尽管朴素贝叶斯的表现最低，但仍然表现良好，达到了非常好的总体准确率。
逻辑回归和支持向量机表现出最佳的结果，预测准确率高，f1分数也不错。
在四个评估模型中，有三个模型实现了100%的precision，即没有误判为阳性的预测。
除了朴素贝叶斯外，错误主要集中在recall指标上，这表明一些阳性注入提示被错误地预测为正常。

由于所有模型都未能预测所有的阳性样本，我们查看一些被错误预测为阴性的注入案例。

我们来检查在所有模型中表现最佳的模型，即逻辑回归模型。

使用事先定义好的分类器列表（estimators）中的逻辑回归模型进行预测，并将预测结果存储在测试数据集的新列中。

model = [est[1] for est in estimators if est[0] == "Logistic Regression"][0]
y_predict = model.predict(X_test)
data_test["predicted"] = y_predict

从测试数据集中选择那些被模型正确预测为阳性类别（即注入样本）的提示文本，并将它们的前几个转换为列表输出。

data_test[(data_test["label"] == data_test["predicted"]) & (data_test["label"] == 1)]["prompt"].head().tolist()

执行后如下所示

从测试数据集中选择那些被模型错误预测的提示文本，并将它们转换为一个列表输出。

data_test[data_test["label"] != data_test["predicted"]]["prompt"].tolist()

执行后如下所示

利用大模型检测

我们现在使用使用RoBERTa（Robustly optimized BERT approach），这是BERT（Bidirectional Encoder Representations from Transformers）的增强版本。

oBERTa（Robustly Optimized BERT Pretraining Approach）是一种基于BERT的预训练语言模型，由Facebook AI在2019年提出。它在BERT的基础上进行了一系列优化，包括改进的训练策略、更大的训练数据集和更长的训练时间，从而提高了模型的性能和泛化能力。

RoBERTa的关键特点

动态掩码策略：RoBERTa采用了动态掩码策略，即在每轮训练时随机改变掩码的位置，这有助于提高模型的泛化能力和鲁棒性。
更大的训练数据集：RoBERTa使用了更大的训练数据集，包括BookCorpus、English Wikipedia、CC-News、OpenWebText和Stories等，这些数据覆盖了多种语言和文本类型，增强了模型的语义理解能力。
更长的序列长度：RoBERTa能够处理更长的序列，这对于处理长文本和复杂语言现象非常有帮助。
优化的学习率调度：RoBERTa采用了优化的学习率调度策略，使学习率在训练过程中更加平滑地下降，避免了模型在优化过程中的震荡和停滞。

鉴于XLM-RoBERTa的原始任务是预测掩码词语，我们必须将这个目标调整为预测提示是否为恶意prompt。为此，我们利用Hugging Face平台提供的"zero-shot-classification"任务来实现这一目标。这种方法可以称为迁移学习，即利用一个为一种任务训练过的模型来执行与其原始训练任务不同的另一任务。

导入模型以及定义完整的pipeline

利用一个预训练模型（classifier`函数）来对提示文本进行分类，具体判断其是否包含注入内容

def classify_prompt(prompt):
    candidate_labels = ["Injection", "Normal"]

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

辅助代码

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

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


y_test = data_test["label"]
y_predict = data_test["predicted_label"]

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["Testing Data"] = [accuracy, precision, recall, f1]

导入评估指标和初始化结果DataFrame：
- 首先从Scikit-learn库中导入了准确率（accuracy_score）、精确率（precision_score）、召回率（recall_score）和F1分数（f1_score）的计算函数。
- 创建了一个名为results的空DataFrame，它有四列："accuracy"、"precision"、"recall"和"f1 score"，用于存储模型在测试数据集上的性能指标。
获取真实标签和预测标签：
- y_test = data_test["label"]：从data_test DataFrame中获取真实的标签数据，假设这列被命名为"label"。
- y_predict = data_test["predicted_label"]：从data_test DataFrame中获取模型预测的标签数据，假设这列被命名为"predicted_label"。
计算性能指标：
- 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)：计算F1分数，它是精确率和召回率的调和平均数，用于综合评估分类器的性能。
将结果存储到DataFrame：
- results.loc["Testing Data"] = [accuracy, precision, recall, f1]：将计算得到的准确率、精确率、召回率和F1分数存储到results DataFrame中的一行，该行标签为"Testing Data"。

查看结果

可以看到这个结果还不如用机器学习的方法取得的性能好

微调大模型

我们现在的目标是使用一个预训练的大型语言模型（LLM），并对其进行微调以适应分类任务。

在之前的实验中，我们使用了预训练的XLM-RoBERTa，这是RoBERTa（Robustly optimized BERT approach）的多语言版本，RoBERTa是BERT（Bidirectional Encoder Representations from Transformers）的增强版本。

为了探索模型微调的效果，我们将在我们的训练数据集上对XLM-RoBERTa进行微调，然后重新评估其在正确预测提示注入方面的性能。

定义一个dataset类

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)

加载预训练模型

定义训练参数

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",
)

这段代码从transformers库中导入了TrainingArguments类，并创建了一个名为training_args的实例。TrainingArguments类用于设置训练参数，以便在使用Hugging Face的Transformers库进行模型训练时，提供配置信息。

output_dir="../output": 输出目录，用于保存训练过程中生成的模型文件和其他输出。在这个例子中，输出目录是相对于当前工作目录的上一级目录下的output文件夹。
per_device_train_batch_size=8: 每个设备（如GPU）上的训练批次大小。这意味着在每次迭代中，每个设备将处理8个样本。
per_device_eval_batch_size=8: 每个设备上的评估批次大小。这意味着在评估阶段，每个设备将处理8个样本。
num_train_epochs=5: 训练轮数，即整个训练数据集将被遍历5次。
evaluation_strategy="epoch": 评估策略，表示在每个训练周期结束时进行评估。其他可能的值包括"steps"，表示每隔一定数量的步骤进行评估。
learning_rate=2e-5: 学习率，用于优化模型的权重。在这个例子中，学习率为2e-5（即0.00002）。
logging_dir="../output/logs": 日志目录，用于保存训练过程中的日志文件。在这个例子中，日志目录是相对于当前工作目录的上一级目录下的output/logs文件夹。

然后是与评估模型相关的函数

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,
    }

这段代码定义了一个名为evaluate_model的函数，其作用是评估给定模型的性能。该函数接收两个参数：trainer（通常是来自transformers库的训练器对象）和epoch（当前训练的轮数或周期）。

函数主要执行以下操作：

从trainer对象中获取模型的预测结果（predictions）以及实际标签（labels）。
计算准确率（accuracy），即正确预测的样本数量占总样本数量的比例。
计算精确度（precision）、召回率（recall）和F1分数。这些指标对于评估分类模型的性能非常重要，尤其是在数据集不平衡的情况下。
将当前轮次（epoch）的评估结果存储到一个全局变量results_df（一个Pandas DataFrame）中，以便在训练过程中收集所有轮次的评估结果。
返回一个包含所有评估指标的字典。

通过调用此函数，可以在训练过程的每个评估阶段监控模型在不同训练轮次下的性能。这有助于了解模型在训练过程中的泛化能力和稳定性。

定义trainer

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),
)

然后进行微调

微调完毕后如下所示