ReGA: Model-Based Safeguard for LLMs via Representation-Guided Abstraction

1. Problem Statement

Large Language Models (LLMs) have achieved significant success but face critical safety and security challenges, including the generation of harmful content and vulnerability to jailbreaking attacks (adversarial prompts designed to bypass safety filters).

Existing solutions face two main limitations:

• Detection-based methods: Approaches like LLM-Judges (using another LLM to check safety) or Perplexity filters often suffer from high computational costs, over-refusal (blocking benign inputs), and a lack of interpretability.
• Model-based analysis: Traditional techniques (e.g., extracting Discrete Time Markov Chains or DFAs) used for smaller Deep Neural Networks (DNNs) fail to scale to LLMs. The vast feature spaces of LLMs require millions of abstract states to model universal properties, creating a scalability bottleneck that renders these methods ineffective for large-scale models.

The core challenge is to develop a safeguard that is scalable (works with large models), efficient (low inference overhead), and interpretable, without relying on expensive black-box classifiers.

2. Methodology: ReGA Framework

The authors propose ReGA (Representation-Guided Abstraction), a model-based analysis framework that leverages recent discoveries regarding low-dimensional safety-critical representations in LLMs. Instead of modeling all hidden states, ReGA focuses on specific directions in the hidden state space that correlate with safety concepts.

The framework operates in three phases:

Phase I: Safety Representations Extraction

• Contrastive Dataset Construction: A dataset $D$ is created containing four subsets: Safe Prompts ($R_S$), Safe Conversations ($C_S$), Harmful Prompts ($R_H$), and Harmful Conversations ($C_H$).
• Feature Selection: The authors extract hidden states from the middle layers of the LLM (specifically layer $L/2$), as these layers are found to contain the most discriminative safety concepts before the model makes a final refusal/generation decision.
• Dimensionality Reduction: Principal Component Analysis (PCA) is applied to the difference between safe and harmful features in the contrastive dataset. This extracts a small set of safety representations ($r_1, \dots, r_K$), which are the principal directions distinguishing safe from harmful inputs. This reduces the dimensionality from thousands (e.g., 4096) to a manageable size (e.g., 8–64).

Phase II: Abstract Model Construction

• Concrete State Mapping: Input tokens are projected onto the extracted safety representations to form concrete safety states.
• Abstraction (Clustering): Unsupervised clustering (K-Means) maps these concrete states into a small set of abstract states ($S$).
• Transition Modeling: A Discrete Time Markov Chain (DTMC) is constructed. The transition probabilities between abstract states are modeled using only safe inputs. This establishes a "safe trajectory."
• Safety Scoring:
  • State Score ($u$): Based on the proportion of safe inputs falling into a specific cluster.
  • Transition Score ($v$): Based on the probability of a transition occurring within the safe distribution.
  • Total Score: A combined score $p(x)$ is calculated for an input sequence based on the last $m$ states and their transitions.

Phase III: Runtime Safeguarding

• During inference, ReGA calculates the safety score for the user prompt (pre-generation) and the full conversation (post-generation).
• If the score falls below a defined threshold, the system triggers a refusal.
• Dual-Level Check: It checks both the prompt level and the conversation level to catch cases where a model might initially refuse but then generate harmful content in a multi-turn dialogue.

3. Key Contributions

1. Novel Framework: Introduction of ReGA, the first model-based safeguard for LLMs that utilizes representation-guided abstraction to overcome scalability issues.
2. Scalability via Representations: By focusing on low-dimensional safety-critical directions rather than universal feature spaces, ReGA reduces the abstract state space from millions to a few dozen (e.g., 32 states), making it feasible for LLMs.
3. Comprehensive Evaluation: Extensive experiments across six open-source LLMs (Vicuna, Llama-2, Qwen, Mistral, Koala, Baichuan) and various datasets.
4. Open Source: Release of the code and practical guidelines for deployment.

4. Experimental Results

The authors evaluated ReGA against baselines including MLP classifiers, Universal DTMCs, LLM-Judges, and Perplexity filters.

• Effectiveness (RQ1):

  • AUROC: ReGA achieved 0.975 at the prompt level and 0.985 at the conversation level for distinguishing safe vs. harmful inputs.
  • Cross-Dataset: It maintained high accuracy (Avg Acc $\approx$ 0.87 for harmful, 0.97 for safe) on unseen datasets, demonstrating strong generalization.
  • Ablation: Both state scores (semantic safety) and transition scores (distributional normality) were shown to be complementary and necessary for optimal performance.

• Generalizability (RQ2):

  • Jailbreaking: ReGA successfully detected jailbreaking attacks (e.g., role-playing, misspellings) with an average accuracy of ~0.35–0.50 (depending on the threshold), significantly outperforming baselines which often failed completely on adversarial inputs.
  • Safety Perspectives: It generalized well across different safety categories (hate speech, inappropriate topics, etc.).
  • Robustness: The framework remained stable across different hyperparameters (number of states, PCA dimensions) and model sizes (scaling from 7B to 14B parameters).

• Comparison (RQ3):

  • Performance: ReGA outperformed all baselines in accuracy. For instance, on the HarmBench dataset, ReGA (AccMCA 0.81) significantly beat MLP-Cls (0.56) and Uni.-DTMC (0.21).
  • Efficiency: ReGA has negligible inference overhead (0.01s per prompt), comparable to MLP classifiers and far superior to LLM-Judges (2.86s per prompt).
  • Interpretability: Unlike black-box classifiers, ReGA provides interpretable abstract states and transition probabilities, allowing developers to understand why an input was flagged.

5. Significance

• SE4AI Advancement: ReGA bridges the gap between Software Engineering for AI (SE4AI) and LLM safety. It demonstrates that formal model-based analysis techniques can be adapted for LLMs if guided by representation engineering.
• Practical Deployment: It offers a lightweight, scalable, and interpretable solution for real-world LLM deployment, addressing the "over-refusal" and "computational cost" issues of current safeguards.
• New Paradigm: The paper establishes that safety in LLMs can be modeled effectively through low-dimensional abstractions of hidden states, paving the way for future research into formal verification and monitoring of generative AI.

In conclusion, ReGA provides a robust, efficient, and interpretable mechanism to safeguard LLMs, effectively balancing safety detection with model utility and scalability.