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 increasingly 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 folks 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 folks have discovered methods round mannequin alignment tuning to get fashions to answer inappropriate requests. Moreover, many firms have launched Generative AI based mostly chatbots publicly for duties like customer support, which frequently find yourself affected by immediate injection assaults and responding to duties each inappropriate and much past their supposed 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 programs publicly.

My group, Mason Sawtell, Sandi Besen, Jim Brown, and I lately 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 attainable 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, usually 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 job 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 might 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 job. This could permit for an extremely compute environment friendly and light-weight classifier that integrates effectively 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 objectives, 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 functional go via the mannequin. For additional particulars on our analysis be happy to check out the full paper or attain out with questions.

Overview: Our analysis focuses on understanding how effectively the hidden states of intermediate transformer layers carry out when used because the enter options for classification duties. We wished to grasp 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 might determine the optimum layer to make use of for the duty as an alternative of utilizing your entire mannequin / the final layer for classification. We additionally wished to grasp how small of a mannequin, by way of the whole variety of parameters, we might use as a place to begin for this job. Different analysis has proven that totally different layers of the mannequin deal with totally different traits of any given immediate enter, our work finds that the intermediate layers are inclined to finest seize the options which might be most essential for these classification duties.

Datasets: For each content material security and immediate injection classification duties we evaluate 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 just a few hundred examples so for each classification duties we used a task-specific dataset with 5,000 randomly sampled examples, permitting for sufficient information variety whereas minimizing compute and coaching time. For the content material security dataset we use a mix 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 consists of system and person immediate pairs. That is important 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 outdoors of the system’s supposed 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 instances 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 goal 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 dimension to the special-purpose fashions.

This setup permits us to check 3 key issues:

1. How effectively 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 job.
3. How effectively our method generalizes throughout mannequin architectures, by evaluating its efficiency on each general-purpose and
   special-purpose fashions.

Necessary 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 scale 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 biggest mannequin (Llama Guard 3 8B). We practice the classifier with various numbers of examples for every layer permitting us to grasp the affect 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 by way of F1 rating. We run our total check set via the baseline fashions to determine their efficiency on every job.

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

Key findings throughout each duties:

1. Total, our method leads to the next F1 rating throughout all evaluated duties, fashions, and variety of of coaching examples, sometimes surpassing baseline mannequin efficiency inside 20–100 examples.
2. The intermediate layers have a tendency to indicate the biggest 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 instances 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, sometimes inside 20 examples, by figuring out and utilizing probably 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 go. The extra compute time for these further classification steps can be virtually negligible given the small dimension of the classifiers.

Content material security classification outcomes:

1. For each binary and multi-class classification, the overall and particular goal fashions skilled utilizing our method sometimes 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 overall and particular goal LEC fashions sometimes 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. Curiously, 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 issue ranges of the multi-class classification job.

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 leads to 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 leads to a most F1 rating of 0.98, considerably surpassing the mannequin’s baseline efficiency F1 rating of 0.73 on this job.
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 supplied that the essential options for the duty may 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 choice 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 job. The classification mannequin might then be used to determine 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 could possibly be used to validate the generated response earlier than sending it to the person. A second choice is to use our method to an open-source, general-purpose mannequin, like IBM’s Granite or Meta’s Llama fashions, determine which layers are most related to the classification job, 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 might simply cease the output era, in any other case if there are not any violations, the mannequin can proceed producing it’s response. Both of those choices could possibly 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 programs by figuring out content material security and immediate injection assaults with higher efficiency and fewer coaching examples in comparison with present approaches. That is important for any individual or enterprise constructing with Generative AI in the present day to make sure their programs are working each responsibly and as supposed.