Lightweight Safety Classification Using Pruned Language Models (2412.13435v1)

Abstract: In this paper, we introduce a novel technique for content safety and prompt injection classification for LLMs. Our technique, Layer Enhanced Classification (LEC), trains a Penalized Logistic Regression (PLR) classifier on the hidden state of an LLM's optimal intermediate transformer layer. By combining the computational efficiency of a streamlined PLR classifier with the sophisticated language understanding of an LLM, our approach delivers superior performance surpassing GPT-4o and special-purpose models fine-tuned for each task. We find that small general-purpose models (Qwen 2.5 sizes 0.5B, 1.5B, and 3B) and other transformer-based architectures like DeBERTa v3 are robust feature extractors allowing simple classifiers to be effectively trained on fewer than 100 high-quality examples. Importantly, the intermediate transformer layers of these models typically outperform the final layer across both classification tasks. Our results indicate that a single general-purpose LLM can be used to classify content safety, detect prompt injections, and simultaneously generate output tokens. Alternatively, these relatively small LLMs can be pruned to the optimal intermediate layer and used exclusively as robust feature extractors. Since our results are consistent on different transformer architectures, we infer that robust feature extraction is an inherent capability of most, if not all, LLMs.

• The paper demonstrates that intermediate hidden states can act as robust features for efficient safety and prompt injection classification.
• It introduces Layer Enhanced Classification (LEC) with penalized logistic regression, outperforming both general-purpose and specialized models.
• The approach minimizes training data requirements and computational costs while achieving high F1-scores with pruned LLMs.

Lightweight Safety Classification Using Pruned LLMs: An Expert Overview

The paper "Lightweight Safety Classification Using Pruned LLMs" proposes a novel methodology for content safety and prompt injection classification using LLMs. It introduces a computationally efficient approach that leverages the intermediate hidden states of LLMs, utilizing them as robust feature extractors for training lightweight classifiers. This framework, termed Layer Enhanced Classification (LEC), capitalizes on the strength of penalized logistic regression (PLR) classifiers and the deep representation capabilities of LLMs, achieving superior performance when compared to both general-purpose models like GPT-4o and specialized models finetuned for task-specific classification.

Key Contributions and Methodological Advances

1. Robust Feature Extraction: The authors establish that the intermediate hidden states of transformer-based architectures are inherently capable of serving as robust feature extractors. This allows a PLR classifier, with parameters equal to the size of the hidden state, to deliver state-of-the-art performance. The paper demonstrates that precise feature extraction occurs not at the final layer but optimally within one of the model's intermediate layers.
2. Reduction in Training Data Requirements: Noteworthy is the demonstration that a minimal set of high-quality training examples can lead to effective classifier training, with significant generalization to unknown data. A promising implication of this is for scenarios with limited availability of labeled training samples, where traditional training demands might be prohibitive.
3. Consistent Cross-Architecture Performance: The experimental results highlight the generality of the proposed approach. The tested LLMs, irrespective of being general-purpose or specially tailored for safety tasks, consistently showed improved classification performance when applying LEC. This suggests the potential for broad applicability across different model types and domains.
4. Superior Performance with Pruned Models: Pruned models, when utilizing only the optimal intermediate layer as feature extractors, surpass the performance of their unpruned counterparts. This advances a sustainable approach to machine learning by significantly reducing the computational footprint while maintaining high classification accuracy.

Empirical Results and Validation

The paper provides robust empirical validation through a series of experiments on both content safety and prompt injection classification tasks. Performance is measured using weighted average F1-scores across models of varying sizes, ranging from 184 million to 8 billion parameters. Across both tasks, the pruned models trained with LEC consistently outperform the baseline performances, achieving leading-edge F1-scores with fewer than 100 training examples.

For content safety, intermediate layers in both Qwen 2.5 Instruct models and Llama Guard models attained high F1-scores, often surpassing the benchmarks set by both GPT-4o and their larger, specialized counterparts. In prompt injection detection, Qwen 0.5B and DeBERTa variants similarly showcased competitive performance, further affirming the efficacy of leveraging intermediate representations.

Theoretical and Practical Implications

The findings underscore an inherent feature extraction capability in the intermediate layers of transformer architectures, promoting a shift away from sole reliance on final layer embeddings. They suggest a model agnostic utility, potentially transforming any transformer-based LLM into a multi-faceted tool capable of efficient classification with minimal computational resources.

Practically, this approach could integrate seamlessly into existing pipelines, enhancing content moderation systems while keeping computational costs low. It aligns with the goals of responsible AI by efficiently establishing robust guardrails and enhancing system integrity.

Future Developments

The paper alludes to exciting future directions, notably integrating LEC directly into the forward pass of LLMs, enabling real-time monitoring and classification as part of token generation processes. This integration could further diminish computational overhead while providing instantaneous assessments of content safety and potential ethic violations. There is also a potential for expanding LEC's application to other classification tasks, further reinforcing its versatility and efficiency within the AI ecosystem.

In conclusion, the research presents a compelling case for Layer Enhanced Classification as a strategic advancement in the field of LLMs, offering valuable insights for ongoing and future research in the optimization and deployment of machine learning models in real-world applications.