# Lightweight Safety Classification Using Pruned Language Models

arXiv: https://arxiv.org/abs/2412.13435 · HTML: https://mlsanity.com/paper/2412.13435/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Safety alignment, refusal and guardrails, one of 12 directions in LLM safety and security.
- Builds on 28 of its 34 references that are in the map.
- First posted on arXiv 18 Dec 2024.

## Placed in

- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 28 of 28)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [Explainable Artificial Intelligence (XAI): Concepts, Taxonomies, Opportunities and Challenges toward Responsible AI](https://mlsanity.com/paper/1910.10045/) (2019)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations](https://mlsanity.com/paper/2312.06674/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset](https://mlsanity.com/paper/2309.11998/) (2023)
- [SALAD-Bench: A Hierarchical and Comprehensive Safety Benchmark for Large Language Models](https://mlsanity.com/paper/2402.05044/) (2024)
- [The Unreasonable Ineffectiveness of the Deeper Layers](https://mlsanity.com/paper/2403.17887/) (2024)
- [Deep k-Nearest Neighbors: Towards Confident, Interpretable and Robust Deep Learning](https://mlsanity.com/paper/1803.04765/) (2018)
- [AEGIS: Online Adaptive AI Content Safety Moderation with Ensemble of LLM Experts](https://mlsanity.com/paper/2404.05993/) (2024)
- [The geometry of hidden representations of large transformer models](https://mlsanity.com/paper/2302.00294/) (2023)
- [Attention Tracker: Detecting Prompt Injection Attacks in LLMs](https://mlsanity.com/paper/2411.00348/) (2024)
- [Seq2Seq-Vis: A Visual Debugging Tool for Sequence-to-Sequence Models](https://mlsanity.com/paper/1804.09299/) (2018)
- [From Understanding to Utilization: A Survey on Explainability for Large Language Models](https://mlsanity.com/paper/2401.12874/) (2024)
- [On the Explainability of Natural Language Processing Deep Models](https://mlsanity.com/paper/2210.06929/) (2022)
- [Prompt Packer: Deceiving LLMs through Compositional Instruction with Hidden Attacks](https://mlsanity.com/paper/2310.10077/) (2023)
- [Diff-eRank: A Novel Rank-Based Metric for Evaluating Large Language Models](https://mlsanity.com/paper/2401.17139/) (2024)
- [SPML: A DSL for Defending Language Models Against Prompt Attacks](https://mlsanity.com/paper/2402.11755/) (2024)
- [Does Representation Matter? Exploring Intermediate Layers in Large Language Models](https://mlsanity.com/paper/2412.09563/) (2024)
- [exBERT: A Visual Analysis Tool to Explore Learned Representations in Transformers Models](https://mlsanity.com/paper/1910.05276/) (2019)
- [Prompt-Augmented Linear Probing: Scaling beyond the Limit of Few-shot In-Context Learners](https://mlsanity.com/paper/2212.10873/) (2022)
- [Model Explainability in Deep Learning Based Natural Language Processing](https://mlsanity.com/paper/2106.07410/) (2021)
- [Towards Agile Text Classifiers for Everyone](https://mlsanity.com/paper/2302.06541/) (2023)
- [Large Language Models Are Overparameterized Text Encoders](https://mlsanity.com/paper/2410.14578/) (2024)
- [MINI-LLM: Memory-Efficient Structured Pruning for Large Language Models](https://mlsanity.com/paper/2407.11681/) (2024)
- [Logistic Regression makes small LLMs strong and explainable "tens-of-shot" classifiers](https://mlsanity.com/paper/2408.03414/) (2024)

## Built on by (most cited first, 5 of 5)

- [GuardReasoner: Towards Reasoning-based LLM Safeguards](https://mlsanity.com/paper/2501.18492/) (2025)
- [UniGuardian: A Unified Defense for Detecting Prompt Injection, Backdoor Attacks and Adversarial Attacks in Large Language Models](https://mlsanity.com/paper/2502.13141/) (2025)
- [LLM Safety From Within: Detecting Harmful Content with Internal Representations](https://mlsanity.com/paper/2604.18519/) (2026)
- [Gandalf the Red: Adaptive Security for LLMs](https://mlsanity.com/paper/2501.07927/) (2025)
- [Disentangled Safety Adapters Enable Efficient Guardrails and Flexible Inference-Time Alignment](https://mlsanity.com/paper/2506.00166/) (2025)

## Abstract (arXiv)

In this paper, we introduce a novel technique for content safety and prompt injection classification for Large Language Models. 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.