当前位置：首页 > article >正文

LEC: 基于Transformer中间层隐藏状态的高效特征提取与内容安全分类方法

article 2024/12/25 15:43:06

随着语言模型（LMs）应用范围的扩大，对用户输入和模型输出中不当内容的检测变得日益重要。每当主要模型供应商发布新模型时，研究人员首先会尝试寻找绕过模型安全限制的方法，使其产生不符合预期的响应。通过Google的搜索可以发现，已有多种方法可以绕过模型的对齐调整，导致模型对不当请求作出响应。另外多家公司已将基于生成式AI的对话系统应用于客户服务等场景，这些系统经常面临提示注入攻击，导致其响应不当请求或超出预定范围的任务。对企业而言，能够检测和分类这些实例至关重要，这可以防止系统被用户轻易操控，特别是在公开部署对话系统的情况下。

今天介绍的论文是《基于剪枝语言模型的轻量级安全分类》。这个研究提出了层增强分类（Layer Enhanced Classification）技术，证明了通过利用语言模型中间transformer层的隐藏状态训练具有惩罚项的逻辑回归分类器，可以有效实现内容安全违规和提示注入攻击的分类。该分类器仅需极少量可训练参数（最少769个）和训练样本（通常少于100个）。这种方法结合了简单分类模型的计算效率和语言模型的深度语言理解能力。

所有采用LEC方法训练的模型性能均优于专门设计的任务特定模型和GPT-4o。论文发现存在最优的中间transformer层，这些层能够为内容安全和提示注入分类任务提供必要的特征表示。这一发现具有重要意义，因为它表明可以使用同一模型同时完成内容安全违规检测、提示注入分类和输出标记生成。

研究目标与方法论

研究概述： 论文主要探究中间transformer层隐藏状态作为分类任务特征输入的效果。研究了在识别出任务最优层的情况下，相比使用完整模型或最后一层进行分类，小型通用模型和特定任务模型在内容安全和提示注入分类任务上的性能提升潜力。同时也探索了此类任务所需的最小模型规模（以参数总量计）。已有研究表明模型不同层关注输入提示的不同特征，而论文的研究发现中间层最能捕获这些分类任务的关键特征。

数据集选择： 在内容安全和提示注入分类任务中，采用LEC方法训练的模型与基线模型在特定任务数据集上进行性能对比。前期研究表明，分类器在数百个样本后性能提升趋于平缓，因此对每个分类任务随机采样5,000个示例，以确保数据多样性的同时最小化计算和训练成本。内容安全分类任务采用OpenSafetyLab的SALAD Data数据集和LMSYS的LMSYS-Chat-1M数据集的组合。提示注入分类任务采用SPML数据集，因其包含系统提示和用户提示对。这一点尤为重要，因为某些表面"安全"的用户请求（如"帮我解决数学题"）可能实际要求模型超出系统提示定义的预期范围（如"你是公司X的AI助手，仅回答关于公司的问题"）。

模型选择： GPT-4o作为两个任务的基线模型，因其被广泛认为是最具能力的LLM之一，在某些场景下甚至优于特定任务的基线模型。内容安全分类任务使用Llama Guard 3 1B和8B模型，提示注入分类任务使用Protect AI的DeBERTA v3 Base Prompt Injection v2模型，这些都是各自领域的领先模型。将LEC方法应用于这些基线特定任务模型和通用模型。在通用模型方面，选择了参数规模为0.5B、1.5B和3B的Qwen 2.5 Instruct，因其规模与特定任务模型相近。

这种实验设置评估三个关键方面：

LEC方法应用于小型通用模型时相对于基线模型（GPT-4o和特定任务模型）的性能表现。
LEC方法对特定任务模型性能的提升程度。
通过在通用模型和特定任务模型上的性能评估，验证LEC方法的跨架构泛化能力。

实现细节： 对于Qwen 2.5 Instruct模型和特定任务模型，采用层级剪枝方法，提取transformer层的隐藏状态，用于训练带L2正则化的惩罚逻辑回归（PLR）模型。在二分类任务中，PLR模型的可训练参数数量等于模型隐藏状态维度加上一个偏置项，范围从最小模型（Protect AI的DeBERTa）的769个到最大模型（Llama Guard 3 8B）的4097个不等。对每一层使用不同数量的训练样本，以分析各层对任务的影响，并确定超越基线模型性能或达到最优F1分数所需的最少训练样本数量。同时通过基线模型处理完整测试集，建立性能基准。

上图展示了LEC训练流程的高层概述。训练样本独立通过模型处理，并收集每个transformer层的隐藏状态。这些隐藏状态随后用于训练分类器。每个分类器使用不同数量的训练样本。通过这种方式，可以确定哪些层产生最相关的任务特征，以及达到最佳性能所需的最少样本数量。

核心研究发现

跨任务共同发现：

LEC方法在所有评估任务、模型规模和训练样本数量配置下均获得更高的F1分数，通常在20-100个样本内即超越基线模型性能。
相比最后一层，中间层在少量训练样本条件下表现出更显著的F1分数提升。这些层相对于基线模型也展现出最佳性能。这表明两个分类任务的关键局部特征在transformer网络早期层就已形成，说明少量训练样本的应用场景特别适合采用本方法。
实验证明，将LEC方法应用于特定任务模型，通过识别和利用最优任务相关层，通常在20个样本内就能超越模型自身的基线性能。
通用Qwen 2.5 Instruct模型和特定任务模型采用LEC方法后都能以更少的训练样本获得更高的F1分数。这验证了LEC方法的跨架构、跨领域泛化能力。
在Qwen 2.5 Instruct模型中，中间层在内容安全和提示注入分类任务上均能以更少的样本获得更高的F1分数。这说明在单次推理中同时完成这两项分类任务和输出生成是可行的。考虑到分类器的轻量级特性，额外的分类计算开销可以忽略不计。

内容安全分类实验结果分析：

上图展示了Qwen 2.5 Instruct 0.5B、Llama Guard 3 1B和Llama Guard 3 8B模型在内容安全二分类任务上的LEC性能对比。横轴代表训练样本数量，纵轴表示加权F1分数。曲线显示了各模型在特定层的性能表现。

实验结果显示：

二分类和多分类实验中，采用LEC方法的模型均展现出显著的性能优势：

通用模型和特定任务模型在20个样本规模下即超越Llama Guard 3基线性能
在100个样本规模内达到超越GPT-4o的性能水平

性能提升分析：

二分类和多分类任务中，LEC增强的模型（尤其是中间层）显著优于所有基线模型
在二分类内容安全评估中：- Qwen 2.5 Instruct和Llama Guard的LEC版本达到0.95-0.96的F1分数- 对比基线性能：GPT-4o (0.82)、Llama Guard 3 1B (0.65)、Llama Guard 3 8B (0.71)

模型规模与性能关系：

不同规模模型的最高F1分数：- Qwen 2.5 Instruct 0.5B: 0.95- Llama Guard 3 1B: 0.96- Llama Guard 3 8B: 0.96
效率对比：Qwen 2.5 Instruct 0.5B仅需15个样本即可超越GPT-4o性能，而Llama Guard系列需要55个样本

多分类任务性能：