about
Lightweight Safety Classification Using Pruned Language Models (arxiv.org)
19 points by sandijean90 on Dec 19, 2024 | hide | past | pdf | 3 comments on HN

In plain words: A simple classifier is trained on the hidden signals from a middle layer of a small language model to spot unsafe content and prompt injections. It beat GPT-4o and models fine-tuned for the task using fewer than 100 examples, with middle layers outperforming the final one.

Abstract

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.

Mason Sawtell, Tula Masterman, Sandi Besen, Jim Brown
arXiv:2412.13435 · cs.CL, cs.AI, cs.LG · submitted Dec 18, 2024
abstract · pdf · html

add comment on HN

This is really easy to set up - and is much more accurate than asking the LLM to predict True/False.

Just feed the outputs of an embedding API into logistic regression, e.g. from sklearn.

  import voyageai
  
  vo = voyageai.Client()
  # This will automatically use the environment variable VOYAGE_API_KEY.
  # Alternatively, you can use vo = voyageai.Client(api_key="<your secret key>")

  texts = ["Sample text 1", "Sample text 2"]

  result = vo.embed(texts, model="voyage-2", input_type="document")
  print(result.embeddings[0])

https://scikit-learn.org/1.5/modules/generated/sklearn.linea...
Are these models available for us to try out?
I'd pay good money for a local LLM with no "content safety" at all.