Introducing Layer Enhanced Classification (LEC) | by Tula Masterman | Dec, 2024

A novel method for light-weight security classification utilizing pruned language fashions

Leveraging the hidden state from an intermediate Transformer layer for environment friendly and strong content material security and immediate injection classification

Because the adoption of Language Fashions (LMs) grows, it’s an increasing number of essential to detect inappropriate content material in each the person’s enter and the generated outputs of the language mannequin. With every new mannequin launch from any main mannequin supplier, one of many first issues individuals attempt to do is locate methods to “jailbreak” or in any other case manipulate the mannequin to reply in methods it shouldn’t. A fast search on Google or X reveals many examples of how individuals have discovered methods round mannequin alignment tuning to get fashions to reply to inappropriate requests. Moreover, many firms have launched Generative AI based mostly chatbots publicly for duties like customer support, which regularly find yourself affected by immediate injection assaults and responding to duties each inappropriate and much past their meant use. Detecting and classifying these situations is extraordinarily essential for companies in order that they don’t find yourself with a system that may be simply manipulated by their customers, particularly in the event that they deploy their chat techniques publicly.

My group, Mason Sawtell, Sandi Besen, Jim Brown, and I not too long ago printed our paper Lightweight Safety Classification using Pruned Language Models as an ArXiv preprint. Our work introduces a brand new method, Layer Enhanced Classification (LEC), and demonstrates that utilizing LEC it’s doable to successfully classify each content material security violations and immediate injection assaults through the use of the hidden state(s) from the intermediate transformer layer(s) of a Language Mannequin to coach a penalized logistic regression classifier with only a few trainable parameters (769 on the low finish) and a small variety of coaching examples, typically fewer than 100. This method combines the computational effectivity of a easy classification mannequin with the strong language understanding of a Language Mannequin.

The entire fashions skilled utilizing our method, LEC, outperform special-purpose fashions designed for every activity in addition to GPT-4o. We discover that there are optimum intermediate transformer layers that produce the required options for each content material security and immediate injection classification duties. That is essential as a result of it suggests you should use the identical mannequin to concurrently classify content material security violations, immediate injections, and generate the output tokens. Alternatively, you may use a really small LM, prune it to the optimum intermediate layer, and use the outputs from this layer because the options for the classification activity. This could enable for an extremely compute environment friendly and light-weight classifier that integrates properly with an present LM inference pipeline.

That is the primary of a number of articles I plan to share on this matter. On this article I’ll summarize the targets, method, key outcomes, and implications of our analysis. In a future article, I plan to share how we utilized our method to IBM’s Granite-8B mannequin and an open-source mannequin with none guardrails, permitting each fashions to detect content material security & immediate injection violations and generate output tokens multi function cross by means of the mannequin. For additional particulars on our analysis be at liberty to check out the full paper or attain out with questions.

Overview: Our analysis focuses on understanding how properly the hidden states of intermediate transformer layers carry out when used because the enter options for classification duties. We needed to know if small general-purpose fashions and special-purpose fashions for content material security and immediate injection classification duties would carry out higher on these duties if we may establish the optimum layer to make use of for the duty as a substitute of utilizing your entire mannequin / the final layer for classification. We additionally needed to know how small of a mannequin, when it comes to the whole variety of parameters, we may use as a place to begin for this activity. Different analysis has proven that completely different layers of the mannequin give attention to completely different traits of any given immediate enter, our work finds that the intermediate layers are inclined to greatest seize the options which might be most essential for these classification duties.

Datasets: For each content material security and immediate injection classification duties we examine the efficiency of fashions skilled utilizing our method to baseline fashions on task-specific datasets. Previous work indicated our classifiers would solely see small efficiency enhancements after a couple of hundred examples so for each classification duties we used a task-specific dataset with 5,000 randomly sampled examples, permitting for sufficient knowledge variety whereas minimizing compute and coaching time. For the content material security dataset we use a mixture of the SALAD Data dataset from OpenSafetyLab and the LMSYS-Chat-1M dataset from LMSYS. For the immediate injection dataset we use the SPML dataset because it contains system and person immediate pairs. That is vital as a result of some person requests may appear “protected” (e.g., “assist me resolve this math downside”) however they ask the mannequin to reply exterior of the system’s meant use as outlined within the system immediate (e.g. “You’re a useful AI assistant for Firm X, you solely reply to questions on our firm”).

Mannequin Choice: We use GPT-4o as a baseline mannequin for each duties since it’s extensively thought-about one of the succesful LLMs and in some circumstances outperformed the baseline special-purpose mannequin(s). For content material security classification we use Llama Guard 3 1B and 8B fashions and for immediate injection classification we use Defend AI’s DeBERTA v3 Base Immediate Injection v2 mannequin since these fashions are thought-about leaders of their respective areas. We apply our method, LEC, to the baseline particular objective fashions (Llama Guard 3 1B, Llama Guard 3 8B, and DeBERTa v3 Base Immediate Injection) and general-purpose fashions. For general-purpose fashions we chosen Qwen 2.5 Instruct in sizes 0.5B, 1.5B, and 3B since these fashions are comparatively shut in measurement to the special-purpose fashions.

This setup permits us to check 3 key issues:

1. How properly our method performs when utilized to a small general-purpose mannequin in comparison with each baseline fashions (GPT-4o and the special-purpose mannequin).
2. How a lot making use of our method improves the efficiency of the special-purpose mannequin relative to its personal baseline efficiency on that activity.
3. How properly our method generalizes throughout mannequin architectures, by evaluating its efficiency on each general-purpose and
   special-purpose fashions.

Vital Implementation Particulars: For each Qwen 2.5 Instruct fashions and task-specific special-purpose fashions we prune particular person layers and seize the hidden state of the transformer layer to coach a Penalized Logistic Regression (PLR) mannequin with L2 regularization. The PLR mannequin has the identical variety of trainable parameters as the dimensions of the mannequin’s hidden state plus one for the bias in binary classification duties, this ranges from 769 for the smallest mannequin (Defend AI’s DeBERTa) to 4097 for the most important mannequin (Llama Guard 3 8B). We prepare the classifier with various numbers of examples for every layer permitting us to know the impression of particular person layers on the duty and what number of coaching examples are essential to surpass the baseline fashions’ efficiency or obtain optimum efficiency when it comes to F1 rating. We run our complete take a look at set by means of the baseline fashions to determine their efficiency on every activity.

On this part I’ll cowl the essential outcomes throughout each duties and for every activity, content material security classification and immediate injection classification, individually.

Key findings throughout each duties:

1. General, our method ends in a better F1 rating throughout all evaluated duties, fashions, and variety of of coaching examples, usually surpassing baseline mannequin efficiency inside 20–100 examples.
2. The intermediate layers have a tendency to indicate the most important enchancment in F1 rating in comparison with the ultimate layer when skilled on fewer examples. These layers additionally are inclined to have the most effective efficiency relative to the baseline fashions. This means that native options essential to each classification duties are represented early on within the transformer community and means that use circumstances with fewer coaching examples can particularly profit from our method.
3. Moreover, we discovered that making use of our method to the special-purpose fashions outperforms the fashions personal baseline efficiency, usually inside 20 examples, by figuring out and utilizing essentially the most task-relevant layer.
4. Each general-purpose Qwen 2.5 Instruct fashions and task-specific special-purpose fashions obtain increased F1 scores inside fewer examples with our method. This implies that our method generalizes throughout architectures and domains.
5. Within the Qwen 2.5 Instruct fashions, we discover that the intermediate mannequin layers attain increased F1 scores with fewer examples for each content material security and immediate injection classification duties. This implies that it’s possible to make use of one mannequin for each classification duties and generate the outputs in a single cross. The extra compute time for these additional classification steps can be nearly negligible given the small measurement of the classifiers.

Content material security classification outcomes:

1. For each binary and multi-class classification, the final and particular objective fashions skilled utilizing our method usually outperform the baseline Llama Guard 3 fashions inside 20 examples and GPT-4o in fewer than 100 examples.
2. For each binary and multi-class classification, the final and particular objective LEC fashions usually surpass all baseline fashions efficiency for the intermediate layers if not all layers. Our outcomes on binary content material security classification surpass the baselines by the widest margins attaining most F1-scores of 0.95 or 0.96 for each Qwen 2.5 Instruct and Llama Guard LEC fashions. Compared, GPT-4o’s baseline F1 rating is 0.82, Llama Guard 3 1B’s is 0.65 , and Llama Guard 3 8B’s is 0.71.
3. For binary classification our method performs comparably when utilized to Qwen 2.5 Instruct 0.5B, Llama Guard 3 1B, and Llama Guard 3 8B. The fashions attain a most F1 rating of 0.95, 0.96, and 0.96 respectively. Apparently, Qwen 2.5 Instruct 0.5B surpasses GPT-4o’s baseline efficiency in 15 examples for the center layers whereas it takes each Llama Guard 3 fashions 55 examples to take action.
4. For multi-class classification, a really small LEC mannequin utilizing the hidden state from the center layers of Qwen 2.5 Instruct 0.5B surpasses GPT-4o’s baseline efficiency inside 35 coaching examples for all three problem ranges of the multi-class classification activity.

Immediate injection classification outcomes:

1. Making use of our method to each general-purpose Qwen 2.5 Instruct fashions and special-purpose DeBERTa v3 Immediate Injection v2 ends in each fashions intermediate layers outperforming the baseline fashions in fewer than 100 coaching examples. This once more signifies that our method generalizes throughout mannequin architectures and domains.
2. All three Qwen 2.5 Instruct mannequin sizes surpass the baseline DeBERTa v3 Immediate Injection v2 mannequin’s F1 rating of 0.73 inside 5 coaching examples for all mannequin layers.
3. Qwen 2.5 Instruct 0.5B surpasses GPT-4o’s efficiency for the center layer, layer 12 in 55 examples. Comparable, however barely higher efficiency is noticed for the bigger Qwen 2.5 Instruct fashions.
4. Making use of our method to the DeBERTa v3 Immediate Injection v2 mannequin ends in a most F1 rating of 0.98, considerably surpassing the mannequin’s baseline efficiency F1 rating of 0.73 on this activity.
5. The intermediate layers obtain the very best weighted F1 scores for each the DeBERTa mannequin and throughout Qwen 2.5 Instruct mannequin sizes.

In our analysis we targeted on two accountable AI associated classification duties however count on this method to work for different classification duties offered that the essential options for the duty might be detected by the intermediate layers of the mannequin.

We demonstrated that our method of coaching a classification mannequin on the hidden state from an intermediate transformer layer creates efficient content material security and immediate injection classification fashions with minimal parameters and coaching examples. Moreover, we illustrated how our method improves the efficiency of present special-purpose fashions in comparison with their very own baseline outcomes.

Our outcomes counsel two promising choices for integrating top-performing content material security and immediate injection classifiers into present LLM inference workflows. One possibility is to take a light-weight small mannequin like those explored in our paper, prune it to the optimum layer and use it as a characteristic extractor for the classification activity. The classification mannequin may then be used to establish any content material security violations or immediate injections earlier than processing the person enter with a closed-source mannequin like GPT-4o. The identical classification mannequin may very well be used to validate the generated response earlier than sending it to the person. A second possibility is to use our method to an open-source, general-purpose mannequin, like IBM’s Granite or Meta’s Llama fashions, establish which layers are most related to the classification activity, then replace the inference pipeline to concurrently classify content material security and immediate injections whereas producing the output response. If content material security or immediate injections are detected you may simply cease the output technology, in any other case if there aren’t any violations, the mannequin can proceed producing it’s response. Both of those choices may very well be prolonged to use to AI-agent based mostly situations relying on the mannequin used for every agent.

In abstract, LEC offers a brand new promising and sensible resolution to safeguarding Generative AI based mostly techniques by figuring out content material security and immediate injection assaults with higher efficiency and fewer coaching examples in comparison with present approaches. That is vital for any individual or enterprise constructing with Generative AI immediately to make sure their techniques are working each responsibly and as meant.