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

A novel strategy 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 vital 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 use 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 reply to inappropriate requests. Moreover, many corporations 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 cases is extraordinarily vital 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 revealed our paper Light-weight Security Classification utilizing Pruned Language Fashions as an ArXiv preprint. Our work introduces a brand new strategy, Layer Enhanced Classification (LEC), and demonstrates that utilizing LEC it’s potential to successfully classify each content material security violations and immediate injection assaults by utilizing 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 strategy combines the computational effectivity of a easy classification mannequin with the strong language understanding of a Language Mannequin.

All the fashions skilled utilizing our strategy, LEC, outperform special-purpose fashions designed for every process in addition to GPT-4o. We discover that there are optimum intermediate transformer layers that produce the mandatory options for each content material security and immediate injection classification duties. That is vital as a result of it suggests you should utilize the identical mannequin to concurrently classify content material security violations, immediate injections, and generate the output tokens. Alternatively, you would 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 process. This is able to enable for an extremely compute environment friendly and light-weight classifier that integrates nicely 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, strategy, key outcomes, and implications of our analysis. In a future article, I plan to share how we utilized our strategy 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 move via the mannequin. For additional particulars on our analysis be at liberty to take a look at the complete paper or attain out with questions.

Overview: Our analysis focuses on understanding how nicely the hidden states of intermediate transformer layers carry out when used because the enter options for classification duties. We wished 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 might determine the optimum layer to make use of for the duty as a substitute of utilizing your complete mannequin / the final layer for classification. We additionally wished to know how small of a mannequin, when it comes to the whole variety of parameters, we might use as a place to begin for this process. Different analysis has proven that completely different layers of the mannequin deal with completely different traits of any given immediate enter, our work finds that the intermediate layers are likely to finest seize the options which can be most vital for these classification duties.

Datasets: For each content material security and immediate injection classification duties we evaluate the efficiency of fashions skilled utilizing our strategy to baseline fashions on task-specific datasets. Earlier 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 information variety whereas minimizing compute and coaching time. For the content material security dataset we use a mix of the SALAD Information 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 crucial as a result of some person requests might sound “protected” (e.g., “assist me resolve this math drawback”) however they ask the mannequin to reply exterior 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 broadly thought of some 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 Shield AI’s DeBERTA v3 Base Immediate Injection v2 mannequin since these fashions are thought of leaders of their respective areas. We apply our strategy, 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 match 3 key issues:

1. How nicely our strategy 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 strategy improves the efficiency of the special-purpose mannequin relative to its personal baseline efficiency on that process.
3. How nicely our strategy generalizes throughout mannequin architectures, by evaluating its efficiency on each general-purpose and
   special-purpose fashions.

Essential 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 (Shield 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 know the influence 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 whole check set via the baseline fashions to determine their efficiency on every process.

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

Key findings throughout each duties:

1. Total, our strategy ends in 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 point out the biggest enchancment in F1 rating in comparison with the ultimate layer when skilled on fewer examples. These layers additionally are likely to have the very best efficiency relative to the baseline fashions. This means that native options vital 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 strategy.
3. Moreover, we discovered that making use of our strategy to the special-purpose fashions outperforms the fashions personal baseline efficiency, sometimes 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 larger F1 scores inside fewer examples with our strategy. This implies that our strategy generalizes throughout architectures and domains.
5. Within the Qwen 2.5 Instruct fashions, we discover that the intermediate mannequin layers attain larger 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 move. The extra compute time for these additional classification steps could 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 strategy 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. As 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 strategy 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 process.

Immediate injection classification outcomes:

1. Making use of our strategy 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 strategy 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 strategy 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 process.
5. The intermediate layers obtain the best weighted F1 scores for each the DeBERTa mannequin and throughout Qwen 2.5 Instruct mannequin sizes.

In our analysis we centered on two accountable AI associated classification duties however anticipate this strategy to work for different classification duties supplied that the vital options for the duty might be detected by the intermediate layers of the mannequin.

We demonstrated that our strategy 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 strategy improves the efficiency of present special-purpose fashions in comparison with their very own baseline outcomes.

Our outcomes recommend 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 process. 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 might be used to validate the generated response earlier than sending it to the person. A second choice is to use our strategy 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 process, 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 would simply cease the output technology, in any other case if there are not any violations, the mannequin can proceed producing it’s response. Both of those choices might be prolonged to use to AI-agent based mostly eventualities relying on the mannequin used for every agent.

In abstract, LEC supplies 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 crucial for any particular person or enterprise constructing with Generative AI right now to make sure their techniques are working each responsibly and as supposed.