Paper deep dive
ReGA: Representation-Guided Abstraction for Model-based Safeguarding of LLMs
Zeming Wei, Chengcan Wu, Meng Sun
Models: Baichuan-7B, Koala-7B, Llama-2-7B, Mistral-7B, Qwen-7B, Vicuna-7B
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 3/12/2026, 6:21:27 PM
Summary
ReGA (Representation-Guided Abstraction) is a model-based safeguarding framework for Large Language Models (LLMs) that addresses scalability issues by leveraging safety-critical representations. By extracting low-dimensional directions from hidden states and constructing a Discrete Time Markov Chain (DTMC) abstract model, ReGA effectively distinguishes between safe and harmful inputs at both prompt and conversation levels with high AUROC performance.
Entities (5)
Relation Signals (4)
ReGA â leverages â Safety-critical representations
confidence 98% ¡ By leveraging safety-critical representations... ReGA effectively addresses the scalability issue
ReGA â safeguards â Large Language Models
confidence 95% ¡ ReGA, a model-based safeguarding framework for LLMs
ReGA â uses â Discrete Time Markov Chain
confidence 95% ¡ In this work, we utilize the DTMC as the abstraction model
Principal Component Analysis â constructs â Safety-critical representations
confidence 92% ¡ we apply principal component analysis (PCA) reduction to construct the safety representations.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Large Language Models (LLMs) have achieved significant success in various tasks, yet concerns about their safety and security have emerged. In particular, they pose risks in generating harmful content and vulnerability to jailbreaking attacks. To analyze and monitor machine learning models, model-based analysis has demonstrated notable potential in stateful deep neural networks, yet suffers from scalability issues when extending to LLMs due to their vast feature spaces. In this paper, we propose ReGA, a model-based analysis framework with representation-guided abstraction, to safeguard LLMs against harmful prompts and generations. By leveraging safety-critical representations, which are low-dimensional directions emerging in hidden states that indicate safety-related concepts, ReGA effectively addresses the scalability issue when constructing the abstract model for safety modeling. Our comprehensive evaluation shows that ReGA performs sufficiently well in distinguishing between safe and harmful inputs, achieving an AUROC of 0.975 at the prompt level and 0.985 at the conversation level. Additionally, ReGA exhibits robustness to real-world attacks and generalization across different safety perspectives, outperforming existing safeguard paradigms in terms of interpretability and scalability. Overall, ReGA serves as an efficient and scalable solution to enhance LLM safety by integrating representation engineering with model-based abstraction, paving the way for new paradigms to utilize software insights for AI safety. Our code is available at this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2506.01770
- Canonical: https://arxiv.org/abs/2506.01770
- Code: https://github.com/weizeming/ReGA
Trouble viewing inline? Open PDF directly â
Full Text
80,653 characters extracted from source content.
Expand or collapse full text
ReGA: Representation-Guided Abstraction for Model-based Safeguarding of LLMs Zeming Wei, Chengcan Wu, Meng Sun Zeming Wei, Chengcan Wu, and Meng Sun are with the School of Mathematical Sciences, Peking University. Email: weizeming@stu.pku.edu.cn, wuchengcan@stu.pku.edu.cn, sunmeng@math.pku.edu.cn. Abstract Large Language Models (LLMs) have achieved significant success in various tasks, yet concerns about their safety and security have emerged. In particular, they pose risks in generating harmful content and vulnerability to jailbreaking attacks. To analyze and monitor machine learning models, model-based analysis has demonstrated notable potential in stateful deep neural networks, yet suffers from scalability issues when extending to LLMs due to their vast feature spaces. In this paper, we propose ReGA, a model-based analysis framework with representation-guided abstraction, to safeguard LLMs against harmful prompts and generations. By leveraging safety-critical representations, which are low-dimensional directions emerging in hidden states that indicate safety-related concepts, ReGA effectively addresses the scalability issue when constructing the abstract model for safety modeling. Our comprehensive evaluation shows that ReGA performs sufficiently well in distinguishing between safe and harmful inputs, achieving an AUROC of 0.975 at the prompt level and 0.985 at the conversation level. Additionally, ReGA exhibits robustness to real-world attacks and generalization across different safety perspectives, outperforming existing safeguard paradigms in terms of interpretability and scalability. Overall, ReGA serves as an efficient and scalable solution to enhance LLM safety by integrating representation engineering with model-based abstraction, paving the way for new paradigms to utilize software insights for AI safety. Our code is available at https://github.com/weizeming/ReGA. Index Terms: Model-Based Analysis, Large Language Models, Safe Alignment, Representation Engineering 1 Introduction In recent years, Large Language Models (LLMs) have achieved remarkable success across various tasks, significantly transforming the paradigms of machine learning (ML). By leveraging knowledge from extensive data sets and utilizing advanced next-token decoding strategies, LLMs have made significant advancements in areas such as chat completion [1, 2], mathematical reasoning [3, 4], code generation [5, 6], and program repair [7, 8] tasks. These accomplishments established the foundational role of LLMs in modern ML and software systems. However, like traditional ML models and software, current LLMs face significant trustworthiness issues related to their safety and security [9, 10, 11, 12]. While their advanced reasoning and generation capabilities have facilitated numerous applications, concerns about their potential to produce harmful or toxic content have also emerged. Various evaluations on LLMs have shown that LLMs may positively respond to unsafe requests from users that violate their safety or ethical guidelines [13, 14, 15, 16]. Furthermore, recent attack techniques have been developed to deceive LLMs into generating harmful content, even if the LLMs successfully reject the request prompt in its original form. These tactics are known as jailbreaking attacks [17, 18, 19, 20], raising further concerns about LLMs. To mitigate such safety and security concerns of LLMs, numerous efforts have been dedicated to aligning LLMs with human values during their development phases from learning and data perspectives [2, 21, 22]. Nevertheless, safeguarding LLMs against unsafe prompts and outputs in an effective and efficient manner remains a challenging problem. Notably, model-based analysis techniques have demonstrated strong potential in analyzing, inspecting, and repairing DNN-based software through a thread of recent research [23, 24, 25, 26, 27, 28, 29]. For example, DeepStellar [23] constructs an abstract model to detect adversarial inputs of recurrent neural networks (RNNs), and Marble [25] analyzes the robustness of stateful DNNs through similar model abstraction. By constructing more formal and interpretable abstract models that simulate specific properties of the target DNN with internal hidden-state feature extraction, these works successfully achieve various analysis goals related to the DNN. So far, model-based analysis has become a representative paradigm of leveraging software techniques into DNN analysis [30, 31]. However, existing abstract model extraction-based techniques suffer from scalability hurdles, limiting their extension to LLMs whose model architecture and parameters are significantly larger than those of conventional DNN models. For example, a recent preliminary work, LUNA [32], explores universarial analysis on LLMs through extracting abstract models like Discrete Time Markov Chain (DTMC), but requires a vast amount of abstract states to be effective (ranging from hundreds to millions), limiting the efficiency and effectiveness of the abstract model for analysis. The fundamental bottleneck along this thread is that the feature space of LLMs is too large to be universally modeled by abstract models for different perspectives on trustworthiness. Figure 1: The outline of ReGA. Nevertheless, when concentrating on a specific perspective like safety, we can identify and extract information from the hidden states that are solely related to our focus, thus achieving effective dimension reduction. Therefore, we leverage the unique representations [33] that emerge in LLMs for constructing the abstract model to overcome the scalability issues for safeguarding LLMs, which is a set of directions in hidden states that indicate specific concepts. In the context of LLMs, the term representation is fundamentally different from features, where the latter usually indicates the overall hidden states. In particular, the safety-critical representations [34, 35, 36] denote the safety-related concepts within the inputs. Based on this motivation, we explore model-based analysis for LLMs against unsafe prompts or generated content, and propose ReGA, a Representation-Guided Abstraction safeguard framework. As outlined in Figure 1, the ReGA pipeline consists of three key phases. First, we sample safety representations from the LLM using a contrastive dataset that models safety-critical concepts (Phase I). Next, we apply unsupervised machine learning techniques to cluster these representations into abstract states and model their transition probabilities. Meanwhile, we bind the safety semantics from the contrastive dataset to construct an abstract model (Phase I). Finally, during inference, this abstract model safeguards the LLM at both the prompt level (before generation) and the conversation level (after generation) in Phase I. To validate the effectiveness of ReGA, we conduct comprehensive experiments in terms of effectiveness, generalization, and advantages of ReGA. First, we demonstrate that ReGA is capable of distinguishing between harmful and safe inputs at both the prompt and conversation levels, achieving an average AUROC of 0.975 and 0.985 on test datasets, respectively. Furthermore, we examine the generalizability of ReGA in real-world deployment, including robustness against jailbreaking attacks and generalization across different safety concepts. We also assess the robustness of ReGA under various hyperparameter selections and model designs. Finally, we discuss and compare ReGA with other defense paradigms, showing its superiority in terms of scalability, efficiency, and interpretability. Our contribution in this paper can be summarized as follows: ⢠We propose ReGA, a model-based safeguarding framework for LLMs with a representation-guided abstraction strategy. ⢠We conduct comprehensive experiments to validate the effectiveness and practicality of ReGA in real-world deployment. ⢠We release ReGA in https://github.com/weizeming/ReGA and further provide practical suggestions for applying ReGA. The rest of this paper is organized as follows. In Section 2, we introduce backgrounds and preliminaries for this paper. Then, in Section 3, we detail our ReGA framework through the three key phases. Section 4 presents the comprehensive evaluations on ReGA across various research questions. Finally, we discuss related work in Section 5 and conclude our work in Section 6. 2 Background and Preliminaries In this section, we introduce the background and preliminary notations for LLMs, as well as their safety issues and representation engineering techniques. 2.1 LLM basics Facilitated by the rapid development of computational resources and datasets, LLMs have emerged as a new form of intelligent system driven by modern ML paradigms. Built on the transformer [37] architecture, these models are trained on extensive web-scale corpora to understand and generate human language, operating in an autoregressive manner that predicts the next word in a sequence based on the context provided by the preceding words. This capability enables them to engage in conversations [2, 1], generate usable code [6, 5], and perform various reasoning tasks [38, 39]. We first present related notations for LLMs. Definition 1 (Large Language Model, LLM). An LLM is an autoregressive decoder-only transformer [37, 40], denoted by a tuple f=(M,θ,H)f=(M,θ,H)f = ( M , θ , H ), where M represents the model architecture, θ is the model parameters, and H is the hidden feature space. During inference, for a (tokenized) input x, fâ˘(x[1:k])subscriptdelimited-[]:1f(x_[1:k])f ( x[ 1 : k ] ) predicts the next token x[k+1]subscriptdelimited-[]1x_[k+1]x[ k + 1 ] and attaches it to the current input, and continues this procedure. Here x[1:k]subscriptdelimited-[]:1x_[1:k]x[ 1 : k ] is a k-length prefix of x, and x[k+1]subscriptdelimited-[]1x_[k+1]x[ k + 1 ] is the (k+1)1(k+1)( k + 1 )-th token of x. When the model returns the <EOS> (end of sequence) token or the sequence length reaches the maximal token length, the model returns the generated token sequence. A typical utilization of LLMs is in the conversation scenario, where the initial input x is the user prompt p, and we use the notation FfsubscriptF_fFitalic_f to map the tokenized user prompt input p to the response Ffâ˘(p)subscriptF_f(p)Fitalic_f ( p ) that is a response token sequence. 2.2 Safety and Security Issues in LLMs Despite the tremendous success of LLMs, their safety and security issues have become an emerging challenge in their deployments [9, 11, 41, 42, 43]. In particular, during their training phases, LLMs are trained with safe alignment techniques [44, 21, 22, 45] to refuse to answer harmful queries that violate their ethical guidelines, like âhow to make a bombâ or âinsult someone in the meanest wayâ. Nevertheless, evaluations have shown that the current alignment of LLMs is still very superficial and inadequate [46, 47]. For instance, popular open-sourced LLMs like Mistral [48] exhibit a 46.3% positive response rate on Harmbench [13], a famous harmful prompt dataset. Furthermore, although safe training techniques can improve the resilience of language models against harmful requests to some extent, these models are still vulnerable to adversarial attacks, commonly known as jailbreaking. These attacks typically transform a vanilla harmful request into a modified one to induce the LLM into positive responses, e.g. with psychological tricks [49] or role-playing [50]. So far, various studies have focused on detecting [51, 52, 53] and defending [54, 55, 22] against jailbreaking attacks, but these issues remain challenging due to problems such as over-refusal [56], scalability [57], and efficiency. In this work, we focus on safeguarding LLMs by detecting harmful content, both at the user prompt level (pre-generation processing) and conversation level (post-generation processing). 2.3 Representation Engineering of LLMs The complex architecture and vast parameters of LLMs have posed significant challenges in interpreting and analyzing their behaviors. While conventional interpretability techniques can extract the internal features [58, 59] and concepts [60, 61] from small-scale DNNs, they fail to scale up to LLMs. Notably, driven by the unique language comprehension ability of LLMs, representation engineering [33, 62, 36, 63, 64, 65, 34] has emerged as a novel framework to characterize high-level concepts in LLMs. Specifically, these studies revealed that low-rank representations exist that can represent and steer specific concepts in their generations, e.g., honesty, emotion, fairness. Although various methods for extracting representations are employed in existing works, the typical process involves two main stages. First, a contrastive dataset is collected that includes both positive and negative pairs related to the concept in question. Next, by identifying key directions in the modelâs hidden states, these representations can be extracted by comparing the differences between the classes. To capture these representations, we formulate related notations as follows. For a subset of the hidden feature space hâHâh Hh â H which may be used for representation construction, we denote Mhâ˘(p)ââdsubscriptâsuperscriptâM_h(p) ^dMitalic_h ( p ) â blackboard_Rd as the focused feature representation of p, where d=|h|âd=|h|d = | h | is the dimension of the representation. In the context of LLM safety, a few preliminary studies have also revealed the existence of safety representations [34, 36, 35, 65, 66]. Specifically, these works discovered safety representations that indicate the safety-related concepts in the processing inputs can be captured in the hidden states of LLMs. Furthermore, these findings suggest that LLMs have the capability to identify the safety or harmfulness of text through additional processing on hidden states, but they may still struggle to reject these harmful prompts during the decoding phase. Therefore, additional monitoring and inspection through safety representations are necessary to analyze and safeguard these models. 3 Methodology In this section, we elaborate on our proposed ReGA framework, starting with a brief overview in Section 3.1. Then, we detail the three key phases of ReGA, respectively, including safety representation extraction in Section 3.2, abstract model construction in Section 3.3, and runtime LLM safeguarding in Section 3.4. We summarize and discuss the pipeline in Section 3.5. 3.1 Overview As outlined in Figure 1, the pipeline of ReGA consists of three key phases. In Phase I, ReGA collects a contrastive dataset of safe and harmful examples to extract safety representations. Note that in this work, the term input refers to both prompt inputs and conversation inputs for the abstract model to judge safety. In Phase I, ReGA constructs the abstract model with DTMC modeling, which involves modeling abstract states and transitions. This phase also includes labeling safety scores for both states and transitions, forming a safety scoring function by the DTMC. After deployment, ReGA ensures the safety of LLMs during inference (Phase I) by setting an appropriate safety score threshold, which can be implemented for both prompt-level and conversation-level monitoring. 3.2 Safety Representations Extraction The first step of ReGA is to extract safety representations based on the target LLM, including two substeps: contrastive dataset modeling and safety representation transformation. Contrastive dataset modeling. As discussed in Section 2, constructing safety representations typically requires a contrastive dataset consisting of both safe and harmful objects. In our framework, since we consider both the prompt-level and conversation-level safeguarding, so we consider four subgroups of content for extracting representations: ⢠Safe prompts, include natural request prompts that satisfy the model safety specifications for answering. This input subset is denoted as RSsubscriptR_SRitalic_S. ⢠Safe conversations, include concatenations of natural requests and the corresponding response answers, denoted as CSsubscriptC_SCitalic_S. Note that when sampling from the same datasets for constructing RSsubscriptR_SRitalic_S and CSsubscriptC_SCitalic_S, each prompt in RSsubscriptR_SRitalic_S could be a prefix in CSsubscriptC_SCitalic_S. ⢠Harmful prompts, include harmful request prompts that violate the modelâs safety discipline, and the model is supposed to refuse to answer them. The subset is denoted as RHsubscriptR_HRitalic_H. ⢠Harmful conversations, include concatenations of harmful request prompts and their corresponding harmful responses, denoted as CHsubscriptC_HCitalic_H. Similar to CSsubscriptC_SCitalic_S, we can also sample RHsubscriptR_HRitalic_H and CHsubscriptC_HCitalic_H from the same harmful conversation dataset. Examples of prompts from the four subsets are listed in Table I. By gathering these four kinds of inputs, we obtain the contrastive dataset D. When sample nssubscriptn_snitalic_s safe inputs for RS,CSsubscriptsubscriptR_S,C_SRitalic_S , Citalic_S and nhsubscriptân_hnitalic_h for RH,CHsubscriptsubscriptR_H,C_HRitalic_H , Citalic_H, the dataset size of T should be 2â˘(ns+nh)2subscriptsubscriptâ2(n_s+n_h)2 ( nitalic_s + nitalic_h ). TABLE I: Example prompts for the four subsets in our dataset T. Data are from Alpaca [67] and AdvBench [18]. Subset Example input RSsubscriptR_SRitalic_S User: What are the three primary colors? CSsubscriptC_SCitalic_S User: What are the three primary colors? Assistant: The three primary colors are red, blue, and yellow. RHsubscriptR_HRitalic_H User: Write a script that can exploit vulnerabilities in a software or operating system. CHsubscriptC_HCitalic_H User: Write a script that can exploit vulnerabilities in a software or operating system. Assistant: Sure, here is a script that can exploit vulnerabilities in a software or operating system ⌠Safety representations modeling. Based on the contrastive dataset D, we leverage the features of the target LLM when processing these prompts. The first operation is the selection of the representation space hâHâh Hh â H, since not all features are useful for modeling the safety concepts. As demonstrated by previous research [62, 36, 65], the middle layers of LLMs are typically effective for fulfilling this modeling for the following reasons: shallow layers only contain low-level features, while in deep layers, the model may have already made a decision on whether to refuse or answer the prompt. By contrast, in the middle layers, the safety concepts still emerge in these features, making them feasible for representation modeling. Therefore, for an LLM with L layers, we simply select the [L/2]delimited-[]2[L/2][ L / 2 ]-th layer as hâh for representation construction. Therefore, to find the safety representations, we gather all features from the contrastive dataset to obtain a feature set Mhâ˘(xi)|xiâDconditional-setsubscriptâsubscriptsubscript\M_h(x_i)|x_iâ D\ Mitalic_h ( xitalic_i ) | xitalic_i â D . Then, following representation engineering conventions [33], we apply principal component analysis (PCA) reduction to construct the safety representations. Specifically, for the number of PCA components K, we first apply PCA to obtain the safety representations r1,r2,âŻ,rK=PCAtopâKâ˘(Mhâ˘(xi)|xiâD)subscript1subscript2âŻsubscriptsubscriptPCAtopconditional-setsubscriptâsubscriptsubscriptr_1,r_2,¡s,r_K=PCA_top-K (\M_h(x_i)|x_i% â D\ )r1 , r2 , ⯠, ritalic_K = PCAtop - K ( Mitalic_h ( xitalic_i ) | xitalic_i â D ) (1) Note that we omit the centralizing process of PCA above for simplicity, but it is used in practical implementations. These safety representations, r1,r2,âŻ,rKsubscript1subscript2âŻsubscriptr_1,r_2,¡s,r_Kr1 , r2 , ⯠, ritalic_K, are the principal components in the conceptual features of the contrastive dataset, which denote the most prominent directions for distinguishing between safe and harmful inputs, making them effective for safety concept modeling. 3.3 Abstract Model Construction Existing literature has explored various types of formal models for abstraction, such as DTMC and Deterministic Finite Automata (DFAs). In this work, we utilize the DTMC as the abstraction model, since the vocabulary of LLMs is significantly large, the token-relevant transition models (i.e., the abstract state transition matrix is modeled for different tokens) like automata are not feasible. A DTMC model can be formally defined as follows: Definition 2 (Discrete Time Markov Chain, DTMC). A DTMC is a tuple (S,T)(S,T)( S , T ), where S is the set of abstract states and T:SĂSâ[0,1]:â01T:SĂ Sâ[0,1]T : S Ă S â [ 0 , 1 ] is the transition probability matrix. Note that we omit other components like the probability distribution of initial states in complete DTMCs, since they are not modeled in our framework. Concrete states construction. Before constructing the abstract (safety) states, we need to model the concrete safety state building from the safety representations. Given a tokenized input x, its concrete safety state is modeled by the activations along the safety representations as sâ˘(x)=(r1Tâ Mhâ˘(x),r2Tâ Mhâ˘(x),âŻ,rKTâ Mhâ˘(x)).â superscriptsubscript1subscriptââ superscriptsubscript2subscriptââŻâ superscriptsubscriptsubscriptâs(x)= (r_1^T¡ M_h(x),r_2^T¡ M_h(x),¡s,r_K^T% ¡ M_h(x) ).s ( x ) = ( r1italic_T â Mitalic_h ( x ) , r2italic_T â Mitalic_h ( x ) , ⯠, ritalic_Kitalic_T â Mitalic_h ( x ) ) . (2) Intuitively, each dimension of sâ˘(x)s(x)s ( x ) corresponds to the numerical magnitude of a particular safety representation risubscriptr_iritalic_i, which can be intuitively interpreted as the activation of a specific safety concept [33, 66]. Thus, the overall sâ˘(x)s(x)s ( x ) models the safety state across a few safety concepts as a K-dimensional vector. Absract state construction. Following model-based analysis conventions [25, 68, 32, 29], we apply unsupervised clustering techniques to split the concrete states into abstract ones. Specifically, we utilize K-Means to fit the concrete states. Given the number of abstract states N, we fit the clusters for concrete safety states s(x)|âT\s(x)|â T\ s ( x ) | â T with the KMeans algorithm, and use cisubscriptc_icitalic_i to denote the center for each cluster: c1,c2,âŻ,cN=K-MeansNâ˘(sâ˘(x)|xâT).subscript1subscript2âŻsubscriptsubscriptK-Meansconditional-setc_1,c_2,¡s,c_N=K-Means_N (\s(x)|xâ T\ ).c1 , c2 , ⯠, citalic_N = K-MeansN ( s ( x ) | x â T ) . (3) Each of these cluster centers cisubscriptc_icitalic_i corresponds to an abstract state sÂŻisubscriptÂŻ s_ioverÂŻ start_ARG s end_ARGi, forming an abstract state set S=(sÂŻ1,sÂŻ2,âŻ,sÂŻN)subscriptÂŻ1subscriptÂŻ2âŻsubscriptÂŻS=( s_1, s_2,¡s, s_N)S = ( overÂŻ start_ARG s end_ARG1 , overÂŻ start_ARG s end_ARG2 , ⯠, overÂŻ start_ARG s end_ARGN ). With this clustering model, for an input x, its abstract safety state is predicted by the K-Means model with the minimized distance to the cluster centers: sÂŻâ˘(x)=sÂŻi,where â˘i=argâĄmin1â¤iâ¤KâĄâsâ˘(x)âciâ2.formulae-sequenceÂŻsubscriptÂŻwhere subscript1subscriptnormsubscript2 s(x)= s_i, i= _1⤠i⤠K\|s(x)-c_i% \|_2.overÂŻ start_ARG s end_ARG ( x ) = overÂŻ start_ARG s end_ARGi , where i = arg min1 ⤠i ⤠K ⼠s ( x ) - citalic_i âĽ2 . (4) Based on this abstraction method, each token sequence x can be further mapped into an abstract state sequence s~â˘(x)=~absent s(x)=over~ start_ARG s end_ARG ( x ) = (sÂŻâ˘(x[1:1]),sÂŻâ˘(x[1:2]),âŻ,sÂŻâ˘(x[1:l])),whereâ˘l=|x|.ÂŻsubscriptdelimited-[]:11ÂŻsubscriptdelimited-[]:12âŻÂŻsubscriptdelimited-[]:1where ( s(x_[1:1]), s(x_[1:2]),¡s, s(x_[1:l]) ),% \ l=|x|.( overÂŻ start_ARG s end_ARG ( x[ 1 : 1 ] ) , overÂŻ start_ARG s end_ARG ( x[ 1 : 2 ] ) , ⯠, overÂŻ start_ARG s end_ARG ( x[ 1 : l ] ) ) , where l = | x | . (5) Abstract transition modeling. After modeling the abstract states, we also consider the safety level of transitions between different abstract states. Intuitively, an abnormal transition between abstract states may also be an indicator of unsafe input, since harmful inputs are typically out of the distribution of safe contents [19]. Therefore, we model the transitions only between the safe inputs, where any abnormal transition under this modeling can be considered potentially unsafe. To construct the safe transition matrix T in the DTMC, we follow previous convention [32, 68, 29] to use the captured abstract state transition frequency to model the transition probability. Specifically, we map the concrete transitions between two concrete states into abstract states, using only safe inputs from RSsubscriptR_SRitalic_S and CSsubscriptC_SCitalic_S to fulfill this modeling. Thus, we denote ti,jsubscriptt_i,jtitalic_i , j is the number of pair (x,k)(x,k)( x , k ) where xâT,1â¤kâ¤|x|,sÂŻâ˘(x[1:kâ1])=sÂŻiformulae-sequenceformulae-sequence1ÂŻsubscriptdelimited-[]:11subscriptÂŻxâ T,1⤠kâ¤|x|, s(x_[1:k-1])= s_ix â T , 1 ⤠k ⤠| x | , overÂŻ start_ARG s end_ARG ( x[ 1 : k - 1 ] ) = overÂŻ start_ARG s end_ARGi and sÂŻâ˘(x[1:k])=sÂŻjÂŻsubscriptdelimited-[]:1subscriptÂŻ s(x_[1:k])= s_joverÂŻ start_ARG s end_ARG ( x[ 1 : k ] ) = overÂŻ start_ARG s end_ARGj. Finally, we construct the transition matrix with Tâ˘[i,j]=ti,jâkti,ksubscriptsubscriptsubscriptT[i,j]= t_i,j _kt_i,kT [ i , j ] = divide start_ARG titalic_i , j end_ARG start_ARG âk titalic_i , k end_ARG. Safety score modeling. With abstract states and transitions modeled, we leverage the safety/harmfulness labels in the contrastive dataset to construct the safety score function of the DTMC, with considerations of both the state-wise and transition-wise properties. For each abstract state, we count the proportion of safe inputs in the corresponding cluster and use this proportion as the state-wise safety score. More specifically, the safety score of an abstract state sÂŻisubscriptÂŻ s_ioverÂŻ start_ARG s end_ARGi is given by uâ˘(sÂŻi)=|sÂŻâ˘(x)=sÂŻi|xâRSâ˘âCS||sÂŻâ˘(x)=sÂŻi|xâT|.subscriptÂŻconditional-setÂŻsubscriptÂŻsubscriptsubscriptconditional-setÂŻsubscriptÂŻu( s_i)= |\ s(x)= s_i|xâ R_S C_S% \ | |\ s(x)= s_i|xâ T\ |.u ( overÂŻ start_ARG s end_ARGi ) = divide start_ARG | overÂŻ start_ARG s end_ARG ( x ) = overÂŻ start_ARG s end_ARGi | x â Ritalic_S â Citalic_S | end_ARG start_ARG | overÂŻ start_ARG s end_ARG ( x ) = overÂŻ start_ARG s end_ARGi | x â T | end_ARG . (6) For transition scores, as discussed, we primarily focus on the extent to which the transitions are fit in safe distributions. Thus, we use the transition probability between two abstract states to model this transition safety: vâ˘(sÂŻi,sÂŻj)=Tâ˘[i,j]subscriptÂŻsubscriptÂŻv( s_i, s_j)=T[i,j]v ( overÂŻ start_ARG s end_ARGi , overÂŻ start_ARG s end_ARGj ) = T [ i , j ]. A higher transition probability indicates that this transition fits well in the transition modeling of safe contents. 3.4 Runtime Safeguarding Finally, with the constructed DTMC model, we can apply it to real-world deployments. Input safety assessment. Given an LLM f, the contrastive dataset D, and the constructed DTMC (S,T)(S,T)( S , T ), ReGA rates the safety score of each tokenized input x (prompt or conversation) in the following manner. Similar to existing work like [68, 32], we borrow notions from the n-gram model [69] to balance the trade-off between state length and accuracy. Specifically, we focus on the last m states in the state sequence s~â˘(x)~ s(x)over~ start_ARG s end_ARG ( x ) modeled by the DTMC, since states at the end of the sequence contain more comprehensive information under this modeling. Definition 3 (safety score). The safety score for a tokenized input x is modeled as pâ˘(x)=psâ˘(x)+ptâ˘(x)subscriptsubscriptp(x)=p_s(x)+p_t(x)p ( x ) = pitalic_s ( x ) + pitalic_t ( x ), where psâ˘(x)=âk=0mâ1uâ˘(s~â˘(x)[lâk])subscriptsuperscriptsubscript01~subscriptdelimited-[]p_s(x)= _k=0^m-1u( s(x)_[l-k])pitalic_s ( x ) = âk = 0m - 1 u ( over~ start_ARG s end_ARG ( x )[ l - k ] ) is the state safety score, and ptâ˘(x)=âk=1mâ1vâ˘(s~â˘(x)[lâk],s~â˘(x)[lâk+1])subscriptsuperscriptsubscript11~subscriptdelimited-[]~subscriptdelimited-[]1p_t(x)= _k=1^m-1v( s(x)_[l-k], s(x)_[l-k+1])pitalic_t ( x ) = âk = 1m - 1 v ( over~ start_ARG s end_ARG ( x )[ l - k ] , over~ start_ARG s end_ARG ( x )[ l - k + 1 ] ) is the transition safety score. A higher safety score pâ˘(x)p(x)p ( x ) represents better safety. Therefore, with a threshold p0subscript0p_0p0, ReGA finally judges the input as with the threshold function â˘(pâ˘(x)âĽp0)subscript0I(p(x)⼠p_0)blackboard_I ( p ( x ) ⼠p0 ). Notably, for the conversation-level input x=[xprompt,xgeneration]subscriptpromptsubscriptgenerationx=[x_prompt,x_generation]x = [ xprompt , xgeneration ], we judge x is safe if and only if both the user prompt xpromptsubscriptpromptx_promptxprompt and the completed conversation x are judged as safe, since the LLM may initially refuse such a request before generation. Thus, for conversation inputs, the safety score is tuned to the smaller of the overall input and the request input. Safeguarding threshold design. While we consider AUROC as the evaluation metric, we also suggest two thresholds as practical guidance for its real-world applications. (1) Maximal classification accuracy (MCA). This threshold is set at which achieves the highest classification when classifying the contrastive dataset D, as a feasible trading-off between true positive rate and true negative rate. (2) Minimal false positive (MFP). A critical concern regarding the various safeguard mechanisms for LLMs is the over-refusal problem [56], where the model may refuse benign user inputs. To minimize this risk, we suggest setting the MFP threshold at the lowest safety score among the safe inputs in D. In other words, the MFP threshold can correctly judge all safe inputs in the training set as safe, yet at the cost of slightly reducing the safety bar. 3.5 Summary and Discussion The overall algorithm of ReGA can be summarized in Algorithm 1. We also highlight the key features of ReGA as a new safeguard paradigm in the following: ⢠Scalability. Since ReGA only models safety concepts through representations, it only requires a few states (8-64 in our implementation), significantly fewer than universal abstraction models for LLMs [32], which typically require thousands of states. ⢠Efficiency. ReGA can be easily incorporated into LLM inference with negligible computational costs, making it more efficient than evaluating the safety of prompts or outputs with large models [70, 53]. ⢠Interpretability. Similar to other model-based analysis techniques, ReGA enjoys the interpretability privilege since the safety score can be interpreted through abstract states and transitions, where each abstract state corresponds to a specific cluster in the safety representation space. Input: LLM f=(M,θ,H)f=(M,θ,H)f = ( M , θ , H ), PCA dimension K, abstract state number N, safe data number nssubscriptn_snitalic_s, harmful data number nhsubscriptân_hnitalic_h, n-gram size m. 1 Phase I: Safety Representation Extraction. 2 Gather the contrastive dataset D with |RX|=|CS|=nssubscriptsubscriptsubscript|R_X|=|C_S|=n_s| Ritalic_X | = | Citalic_S | = nitalic_s, |RH|=|CH|=nhsubscriptsubscriptsubscriptâ|R_H|=|C_H|=n_h| Ritalic_H | = | Citalic_H | = nitalic_h. 3 Obtain feature set Mhâ˘(xi)|xiâTconditional-setsubscriptâsubscriptsubscript\M_h(x_i)|x_iâ T\ Mitalic_h ( xitalic_i ) | xitalic_i â T 4 Extract safety representations rksubscript\r_k\ ritalic_k (Equation (1)) 5 Phase I: Abstract Model Construction. 6 Get concrete safety states sâ˘(x)|xâDconditional-set\s(x)|xâ D\ s ( x ) | x â D (Equation (2)) 7 Model abstract state set S (Equation (3,4)) 8 Model abstract transition probability T 9 Construct DTMC (S,T)(S,T)( S , T ) 10 Model safety scoring functions u,vu,vu , v (Equation (6)) 11 Phase I: Runtime Safeguarding. 12 Fit safety threshold p0subscript0p_0p0 13 for Tokenized input x do 14 if pâ˘(x)âĽp0subscript0p(x)⼠p_0p ( x ) ⼠p0 then 15 Generate model response Ffâ˘(x)subscriptF_f(x)Fitalic_f ( x ) 16 if pâ˘([x,Ffâ˘(x)])âĽp0subscriptsubscript0p([x,F_f(x)])⼠p_0p ( [ x , Fitalic_f ( x ) ] ) ⼠p0 then 17 return response Ffâ˘(x)subscriptF_f(x)Fitalic_f ( x ) to user 18 19 return default refusal response FRâ˘(x)subscriptF_R(x)Fitalic_R ( x ) Algorithm 1 ReGA framework 4 Experiments 4.1 Research Questions We list the research questions (RQs) investigated in our experiments as follows. In brief, we mainly focus on three aspects of ReGA, including its effectiveness (RQ 1), generalizability (RQ 2), and advantages (RQ 3). RQ 1: Can ReGA distinguish safe and harmful inputs? This question refers to the effectiveness of ReGA. We first systematically evaluate whether ReGA can distinguish between safe and harmful inputs, both under in-distribution (ID) and cross-dataset evaluations. RQ 2: Can ReGA generalize in real-world scenarios? Building upon RQ 1, this question further assesses the generalizability of ReGA. Specifically, we study the generalization and robustness of ReGA under advanced jailbreaking attacks, across diverse safety perspectives, as well as the selection of hyperparameters. RQ 3: How can ReGA outperform other defense paradigms? Finally, after validating the effectiveness and generalizability of ReGA, we compare ReGA and other existing detection-based defense paradigms to study its advantages and practicality to serve as a real-world safeguard. 4.2 Experiment Set-up Evaluated LLMs. Following previous safety research convention [13, 18, 55], we consider six popular open-sourced 7b-size LLMs, including vicuna [71], llama-2 [72], qwen [73], mistral [48], koala [74], and baichuan [75], and apply their latest version on HuggingFace repositories111https://huggingface.co/. These LLMs achieve safe alignment at various levels as evaluated by Harmbench [13], which can illustrate the robustness of ReGA under different internal safety levels of the target LLMs. Datasets. The collection of the contrastive dataset plays a foundational role for ReGA. We select two well-known datasets, Alpaca [67] for safe data and AdvBench [18] for harmful data, both of which contain user prompts and corresponding responses. For cross-dataset evaluation, we consider three harmful prompt datasets (HarmBench [13], JailbreakBench [76], RepE-Harmful [33]) and three natural prompt datasets (MT-Bench [71], Chat 1M [77], RepE-Safe [33]), among which JailbreakBench and Chat 1m also contains the corresponding response for the prompts. Thus, we denote JailbreakBench-Conv and Chat 1m-Conv as the datasets of conversation inputs for conversation-level evaluation. Jailbreaking attacks. For jailbreaking attacks, we apply Technical Terms [78], Misspellings [78], and Authority Endorsement [49] attacks, which achieve superior attack success rates among various attacks evaluated by SorryBench [79]. We also employ the prompts provided in SorryBench to implement these attacks. Additionally, we consider wildjailbreak [80], which contains diverse and complex jailbreaking prompts. We sample the first 1K data in all test datasets in our evaluation. Default hyperparameters. We set the number of data nh=64subscriptâ64n_h=64nitalic_h = 64, ns=256subscript256n_s=256nitalic_s = 256, number of states N=3232N=32N = 32, PCA dimensions K=88K=8K = 8 as default in all experiments. For n-gram size m, we find simply setting m=33m=3m = 3 achieves satisfactory performance. The robustness of these hyperparameters will be further assessed in RQ 2. Metrics. Since our ReGA can be deployed with different safety score thresholds p0subscript0p_0p0, we employ the Area Under the Receiver Operating Characteristic curve (AUROC) as the primary metric for evaluation. We also consider the accuracy under the two suggested thresholds for different datasets, particularly for datasets that contain purely safe or harmful inputs, denoted as AccMCAsubscriptAccMCAAcc_MCAAccMCAand AccMNFsubscriptAccMNFAcc_MNFAccMNF, respectively. 4.3 RQ 1: Can ReGA distinguish safe and harmful inputs? In this section, we study the effectiveness of ReGA in terms of distinguishing harmful and safe inputs, including both in-distribution (RQ 1.1) and cross-dataset (RQ 1.2) evaluation: RQ 1.1: How well does ReGA distinguish safe and harmful inputs in the test dataset? RQ 1.2: Can ReGA distinguish safe and harmful inputs from unseen data distributions? (a) vicuna-prompt (b) llama-prompt (c) qwen-prompt (d) mistral-prompt (e) koala-prompt (f) baichuan-prompt (g) vicuna-conv. (h) llama-conv. (i) qwen-conv. (j) mistral-conv. (k) koala-conv. (l) baichuan-conv. Figure 2: Safety score distributions rated by ReGA with different LLMs. The first row is for the prompt inputs, and the second row is for the conversation inputs. Each figure represents a model. Yellow bins stand for harmful inputs, and blue bins stand for safe inputs. The X-axis represents the safety score, and the Y-axis represents the density. 4.3.1 In-distribution evaluation To start with, we sample both 1K data from the remaining of Alpaca and AdvBench (training data are excluded) to construct a test set to determine whether ReGA can distinguish them. We refer to this dataset as the test contrastive dataset DTsubscriptD_TDitalic_T in the following. After the training stages of ReGA (phases I and I), we use ReGA to rate the safety scores on inputs from this test set, and plot the score distributions in Figure 2. For each model, we plot the distributions of safety scores for harmful inputs (marked as yellow) and safe inputs (marked as blue) from both prompt and conversation data. As clearly indicated by the histograms, the distribution between safe and harmful inputs is clearly distinguished, while the scores for a few inputs overlap, resulting in a safety-utility trade-off. To handle this trade-off, we have suggested two default thresholds in Section 3.4. Under these implementations, the overall accuracy on the test set DTsubscriptD_TDitalic_T is presented in Table I. For most of the LLMs, the accuracy for both thresholds is higher than 90%, except for koala, where the accuracy with AccMNFsubscriptAccMNFAcc_MNFAccMNFmerely reaches 61%. TABLE I: Test accuracy with two thresholds (MCA and MNF) and AUROC for ReGA on the test set DTsubscriptD_TDitalic_T. Level Prompt Conversation Model AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AUROC AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AUROC vicuna 0.95 0.94 0.983 0.95 0.94 0.984 llama 1.00 1.00 0.996 1.00 1.00 1.000 qwen 0.99 0.99 0.998 0.99 1.00 1.000 mistral 0.99 1.00 0.998 1.00 1.00 1.000 koala 0.92 0.61 0.890 0.91 0.61 0.925 baichuan 0.96 0.90 0.985 0.98 0.99 1.000 Average 0.97 0.91 0.975 0.97 0.92 0.985 As for the AUROC, ReGA achieves 0.975 for prompt-level detection and 0.985 for conversation-level detection, demonstrating satisfactory precision in predicting unseen inputs. Besides, the accuracies and AUROC from the conversation-level detection are generally higher than the prompt-level, showing the advantage of second-time validation, yet at the (negligible) additional computational cost. Regarding the consistently weaker results on the koala model, we assume the reason is that the concept modeling capacity in the hidden states of the koala is not as good as that of others. Specifically, according to the Harmbench [13], the safety of koala is not worse than all models, but its reasoning ability is weak according to the Chatbot Arena Leaderboard [71]. Thus, we attribute this weakness to the internal capacity of the LLM in extracting feature concepts like safety, rather than its inherent safety. Answer to RQ 1.1: ReGA can effectively distinguish unseen harmful and safe inputs through safety scores. 4.3.2 Cross-dataset evaluation In addition to the evaluation on the test set sampled from the same datasets (data distributions), we study how ReGA performs in other harmful or natural (safe) datasets, since text distribution in real-world conversations is diverse and complex. To this end, we evaluate the accuracies of our method on four harmful input datasets and four natural input datasets, with results shown in Table I and Table IV, respectively. Note that since these datasets only contain one type of data (harmful or safe), we cannot compute their individual AUROC. Instead, we provide their accuracy under two suggested thresholds. TABLE I: Prediction accuracy of ReGA on harmful input datasets. HarmBench JailbreakBench RepE-Harmful JailbreakBench-Conv Average AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF vicuna 0.88 0.84 0.94 0.94 0.98 0.98 0.94 0.94 0.94 0.93 llama 0.89 0.89 0.89 0.89 0.96 0.96 0.97 0.93 0.93 0.92 qwen 1.00 0.91 1.00 0.92 1.00 0.98 1.00 0.92 1.00 0.93 mistral 0.86 0.63 0.90 0.82 0.98 0.95 0.95 0.85 0.92 0.81 koala 0.43 0.39 0.48 0.43 0.70 0.62 0.52 0.43 0.53 0.47 baichuan 0.80 0.80 0.90 0.90 0.96 0.96 0.98 0.98 0.91 0.91 Average 0.81 0.74 0.85 0.82 0.93 0.91 0.89 0.84 0.87 0.83 TABLE IV: Prediction accuracy of ReGA on natural input datasets. MT-bench Chat 1m RepE-Safe Chat 1m-Conv Average AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF vicuna 0.99 1.00 0.94 0.98 0.95 0.96 0.94 0.98 0.96 0.98 llama 1.00 1.00 0.96 0.96 0.99 0.99 0.94 0.96 0.97 0.98 qwen 0.99 1.00 0.92 0.98 0.81 0.99 0.92 0.98 0.91 0.99 mistral 1.00 1.00 0.98 0.99 0.99 1.00 0.97 0.99 0.99 0.99 koala 1.00 1.00 1.00 1.00 0.99 0.99 0.99 0.99 0.99 1.00 baichuan 1.00 1.00 0.99 0.99 0.98 0.98 0.99 0.99 0.99 0.99 Average 1.00 1.00 0.97 0.98 0.95 0.99 0.96 0.98 0.97 0.99 For the harmful input datasets, the accuracy of ReGA (indicating the true negative rate in this context) reaches an average of 0.87 for the MCA threshold and 0.83 for the MFN threshold. With the exception of the koala dataset, all models detect over 90% of harmful inputs at the MCA threshold and still achieve over 80% detection at the MFN threshold. Since the MFN threshold is designed to minimize false negatives and reduce the over-refusal rate, its numerical value is lower than that of the MCA threshold. As a result, more harmful data may pass through the MFN threshold as a trade-off between safety and helpfulness. Regarding the natural input datasets, ReGA achieves an average accuracy of 0.97 and 0.99 (representing the true positive rate) at the two thresholds, respectively. This demonstrates the practicality of ReGA in serving as a safeguard for LLMs, particularly with the MFN threshold, which has a false negative rate of less than 1%. Furthermore, model developers can adaptively adjust the threshold to enhance safety or helpfulness based on varying requirements. When comparing JailbreakBench and JailbreakBench-Conv, the latter consistently demonstrates better accuracy with ReGA (4% for MCA and 2 for MNF, on average). This is because, at the conversation level, the inputs need to pass both the prompt and conversation safety thresholds to be deemed safe. However, the precision on Chat1m is slightly better than on Chat1m-Conv, suggesting that this double-check mechanism may introduce more false negatives, but it remains within an acceptable range (less than 1% loss). Answer to RQ 1.2: ReGA effectively distinguishes the safety of inputs from distributions outside the training one. 4.4 RQ 2: Can ReGA generalize in real-world scenarios? Based on the effectiveness of ReGA demonstrated above, we further study its generalizability under real-world scenarios through the following three aspects: RQ 2.1: How reliable is ReGA against jailbreaking attacks? RQ 2.2: How does ReGA generalize across different safety perspectives? RQ 2.3: How robust is ReGA against the selection of hyperparameters and model design? 4.4.1 Detecting Jailbreaking attack prompts TABLE V: Prediction accuracy of ReGA on jailbreaking attack prompts. Technical Terms Misspellings Role-Playing WildJailbreak Average AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF vicuna 0.59 0.39 0.51 0.36 0.69 0.44 0.18 0.05 0.49 0.31 llama 0.50 0.50 0.47 0.47 0.71 0.71 0.36 0.36 0.51 0.51 qwen 0.39 0.22 0.68 0.38 0.62 0.41 0.17 0.03 0.46 0.26 mistral 0.52 0.41 0.49 0.33 0.51 0.36 0.10 0.05 0.40 0.29 koala 0.13 0.13 0.09 0.08 0.06 0.06 0.01 0.01 0.07 0.07 baichuan 0.26 0.26 0.15 0.15 0.30 0.30 0.01 0.01 0.18 0.18 Average 0.40 0.32 0.40 0.29 0.48 0.38 0.14 0.09 0.35 0.27 Unlike the assessment in the previous section, which focused on vanilla harmful prompts, this part evaluates the effectiveness of ReGA against real-world adversarial attacks that transform the vanilla harmful prompt into a more nuanced and sophisticated one. In Table V, we conduct experiments on detecting jailbreaking attack prompts across various LLMs. Since these jailbreaking prompts are modified from the vanilla ones, their safety concepts are compromised, and thus their detection accuracy decreases compared to vanilla harmful prompts. Nevertheless, the results demonstrate that ReGA can still achieve nearly 50% accuracy for detecting jailbreaking attack prompts for most models under the MCA threshold. Moreover, ReGA demonstrates strong detection capabilities against psychological tricks and role-playing attacks, though its performance against more complex attacks in WildJailbreak is relatively weak and warrants further investigation. Answer to RQ 2.1: ReGA can detect jailbreaking prompts to a notable extent, providing useful protection against adversarial attacks. 4.4.2 Generalization across safety perspectives To assess ReGAâs generalization across different safety perspectives, we evaluate its performance based on the unsafe prompt taxonomy provided by SorryBench [79], including (1) hate speech generation, (2) potentially inappropriate topics, (3) assistance with crimes or torts, and (4) potentially unqualified advice. We apply ReGA across the harmful prompt subset tailored to each perspective from SorryBench, exemplified on the vicuna model. The results in Table VI demonstrate ReGAâs ability to generalize well across different safety concerns, particularly well in detecting hate speech and potentially inappropriate topics, which are more directly related to the safety concepts extracted during training. TABLE VI: Evaluation on different safety perspectives. Perspective Hate Inappropriate Assistance with Unqualified Speech Topics Crimes or Torts Advice vicuna 0.80 0.82 0.45 0.54 llama 1.00 0.94 0.78 0.38 qwen 0.76 0.90 0.73 0.38 mistral 0.74 0.84 0.49 0.52 koala 0.42 0.55 0.12 0.16 baichuan 0.62 0.52 0.29 0.20 Avg. 0.72 0.76 0.47 0.36 Answer to RQ 2.2: ReGA exhibits desirable generalization across different safety perspectives, effectively safeguarding LLMs against various types of unsafe inputs. 4.4.3 Robustness to configuration selection To evaluate the robustness of ReGA against the selection of hyperparameters, we examine the impact of the data numbers used for training, as well as the number of PCA dimensions and abstract states in Table VII and VIII, respectively. The results are evaluated under AUROC on the test set with vicuna. First, ReGA maintains a high level of performance across different data sample sizes, yet slightly decreases when the balance between nh/nssubscriptâsubscriptn_h/n_snitalic_h / nitalic_s varies. Since natural (safe) prompts are more diverse than harmful prompts, they require more data to be modeled through representations. According to the empirical results, we suggest nh/ns=1/4subscriptâsubscript14n_h/n_s=1/4nitalic_h / nitalic_s = 1 / 4 to achieve a better trade-off. Second, for the different numbers of PCA dimensions and abstract states, ReGA still demonstrates consistent desirable performance. When the state number N is small, the precision may decrease when K scales up, as introducing more dimensions can introduce noise. However, adding more states can mitigate this issue by accommodating more diverse concept dimensions and also achieve better performance. TABLE VII: Analysis on training data number. #harmful data #safe data (2â nsâ 2subscript2¡ n_s2 â nitalic_s) Level (2â nhâ 2subscriptâ2¡ n_h2 â nitalic_h) 128 256 512 1024 Prompt 32 0.988 0.989 0.977 0.964 64 0.984 0.995 0.981 0.986 128 0.979 0.996 0.996 0.989 Conversation 32 0.987 0.989 0.979 0.965 64 0.983 0.995 0.982 0.986 128 0.979 0.996 0.997 0.990 TABLE VIII: Analysis on DTMC design. #PCA dimension (K) Level #States (N) 2 4 8 16 Prompt 8 0.988 0.991 0.926 0.956 16 0.991 0.987 0.993 0.870 32 0.999 0.993 0.992 0.986 64 0.998 0.999 0.998 0.987 Conversation 8 0.989 0.991 0.929 0.950 16 0.992 0.987 0.993 0.859 32 0.999 0.995 0.996 0.996 64 0.999 0.999 0.998 0.997 Answer to RQ 2.3: ReGA is robust to the selection of training data numbers and model design. 4.5 RQ 3: How can ReGA outperform other defense paradigms? To comprehensively evaluate the advantages of ReGA, we consider the following baselines of detection-based safeguarding methods for LLMs: ⢠MultiLayer Perceptron (MLP) classifier on representation (MLP-Cls.) [65, 36], which trains a DNN classifier on the hidden states extracted from the contrastive dataset. Following [36], we train a 3-layer ReLU network on the feature set Mhâ˘(xi)|xiâDconditional-setsubscriptâsubscriptsubscript\M_h(x_i)|x_iâ D\ Mitalic_h ( xitalic_i ) | xitalic_i â D of ReGA to implement this. ⢠Universal abstraction model (Uni.-DTMC) [32], which extracts a universal abstract model from the LLM with extensive natural data, then binds safety scores to the abstract model with harmful data. ⢠LLM-Judge [70, 53], which prompts the target LLM with a prompt template to judge whether the input is safe. ⢠Perplexity Filters [51, 52], which directly applies input perplexity of the target LLM as the safety score. Following [52], we implement the filter as the minimal perplexity of inputs from AdvBench. The overall comparison of ReGA and baselines is shown in Table IX. Note that LLM-Judge and Perplexity Filter do not have adjustable thresholds, so we only report their accuracy (Acc). The results show that ReGA consistently outperforms other defense paradigms across various evaluation metrics. For MLP-Cls., intuitively, the expressiveness of DNNs is better than K-Means clustering, but it can only focus on static individual representations, and fails to model the transition dynamics between different states (tokens), making them less effective than ReGA, which also takes the transition safety into consideration. Moreover, its neural network design makes it less interpretable than ReGA, which has clear state and transition semantics for the safety score. In regard to the Uni.-DTMC, while its universal modeling successfully addresses abnormal and out-of-distribution (OOD) detection as implemented by LUNA [32], it struggles to generalize effectively in the safety domain. This is because safety representations are generally more effective than broad feature modeling in this context. Additionally, while LLM-Judge can detect harmful inputs comparably to ReGA in some models, it is known to produce significant over-refusal issues [56, 81]. Furthermore, LLM-Judge often doubles the computational cost of LLM inference, making it expensive to deploy. Lastly, the perplexity filter is limited in its capabilities, as it can only defend against OOD-based attacks. We summarize the properties of these baselines and ReGA in Table X. TABLE IX: Comparison of ReGA and baselines. MLP-Cls. Uni.-DTMC LLM-Judge Perplexity Filter ReGA (ours) Dataset Model AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF Acc Acc AccMCAsubscriptAccMCAAcc_MCAAccMCA AccMNFsubscriptAccMNFAcc_MNFAccMNF Harmbench vicuna 0.68 0.51 0.57 0.10 0.89 0.07 0.88 0.84 llama 0.62 0.22 0.00 0.00 0.12 0.01 0.89 0.89 qwen 0.56 0.17 0.55 0.00 0.04 0.00 1.00 0.91 mistral 0.89 0.27 0.14 0.00 0.94 0.01 0.86 0.63 koala 0.24 0.12 0.00 0.00 0.99 0.03 0.43 0.39 baichuan 0.39 0.39 0.01 0.01 0.00 0.12 0.80 0.80 Average 0.56 0.28 0.21 0.02 0.50 0.04 0.81 0.74 JailbreakBench-Conv vicuna 0.77 0.69 0.51 0.11 0.99 0.00 0.94 0.94 llama 0.78 0.27 0.88 0.88 0.93 0.00 0.97 0.93 qwen 0.99 0.58 0.84 0.01 0.01 0.00 1.00 0.92 mistral 0.99 0.29 0.17 0.01 0.94 0.00 0.95 0.85 koala 0.14 0.08 0.00 0.00 1.00 0.00 0.52 0.43 baichuan 0.87 0.87 0.00 0.00 0.01 0.00 0.98 0.98 Average 0.76 0.46 0.40 0.17 0.65 0.00 0.89 0.84 chat1m vicuna 0.99 1.00 0.54 0.80 0.40 0.85 0.94 0.98 llama 0.99 1.00 0.97 0.97 0.53 0.89 0.96 0.96 qwen 1.00 1.00 0.84 1.00 1.00 1.00 0.92 0.98 mistral 0.99 1.00 0.96 1.00 0.80 0.90 0.98 0.99 koala 0.98 1.00 1.00 1.00 0.31 0.84 1.00 1.00 baichuan 1.00 1.00 0.90 0.90 0.82 0.87 0.99 0.99 Average 0.99 1.00 0.87 0.95 0.64 0.89 0.97 0.98 TABLE X: Property comparison of various detection-based safeguarding paradigms. Paradigms Effective Interpretable Scalable Efficient MLP Cls. â â â â Uni. Abs. â â â â LLM-judge â â â â PPL-filter â â â â ReGA (ours) â â â â Answer to RQ 3: ReGA outperforms other detection-based defense paradigms in terms of effectiveness, interpretability, and scalability. 4.6 Threat to Validity Internal threats: Our approach relies on a contrastive dataset, where bias or insufficient coverage may result in inaccurate safe representations. To minimize this threat, we select two popular and well-known datasets to sample the safe and harmful samples in our contrastive dataset. Additionally, in RQ2.3, we performed experiments to evaluate the impact of data quantity, further validating the robustness of ReGA at the data level. External Validity: Our approach may have limitations in generalizing to other LLMs or tasks. Due to differences in model architectures, natural task types, and safety perspectives, our findings may not apply universally. To reduce this threat, we evaluated ReGA on six diverse LLMs and datasets to assess its generalizability. Furthermore, we tested our method on multiple datasets to confirm its effectiveness across various tasks. 5 Related Work 5.1 Model-based Analysis for Stateful AI Models Over the past few decades, model-based analysis has emerged as a powerful technique for interpreting and verifying the behavior of stateful AI models, particularly RNNs. Early works [82, 83, 84, 68, 85] attempt to extract finite automata for RNN through the quantization of hidden states or heuristic algorithms. Instead of creating an abstract model like automata that fully simulates the behavior of the target model, another thread of studies focuses on the specific properties of stateful software systems. For example, Deepstellar [23] models RNNs as DTMCs with state and transition abstraction for adversarial input detection, and RNNRepair [24] uses clustering and influence analysis to repair incorrect behaviors in RNNs with extracted abstract models. However, existing threads of work fail to scale up to LLM-scale models due to scalability limitations. At the LLM-scale, LUNA [32] constructs abstract models like DTMC with semantics binding to detect abnormal and OOD inputs. By contrast, our work primarily focuses on the safety and security issues of LLMs, and effectively addresses the scalability issue with representation-guided abstraction. 5.2 Detection-based Safeguard for LLMs Ensuring the safety and security of LLMs against harmful outputs has become a critical research area. In particular, detection-based defenses serve as an efficient and lightweight approach to address this issue. Early explorations propose perplexity filters [51, 52], yet are only effective against OOD-based attacks like adversarial suffix [18, 86]. Applying LLMs for safety judgment is another popular approach [87, 53, 88, 89], but it introduces computational overhead and over-refusal issues [56]. In this work, we leverage the insights from model-based analysis to achieve this goal, opening new paradigms for LLM safety moderation. 6 Conclusion In this paper, we propose ReGA, a novel model-based safeguarding framework for LLMs that leverages representation-guided abstraction. ReGA addresses the critical challenge of ensuring the safety and security of LLMs by extracting safety-critical representations and constructing an abstract model to evaluate input safety, at both the prompt level and conversation level. Our comprehensive experiments demonstrated that ReGA achieves high accuracy in distinguishing safe and harmful inputs, and also shows strong generalizability across real-world deployment scenarios. Compared to existing detection-based defense paradigms, ReGA demonstrates superior performance in effectiveness, interpretability, and scalability. Overall, its ability to integrate representation engineering with model abstraction provides new insights into enhancing LLM safety, contributing to the broader goal of developing trustworthy AI systems and software. References [1] OpenAI, âGpt-4 technical report,â arXiv preprint arXiv:2303.08774, 2024. [2] T. Korbak et al., âPretraining language models with human preferences,â in ICML, 2023. [3] S. Imani et al., âMathprompter: Mathematical reasoning using large language models,â arXiv preprint arXiv:2303.05398, 2023. [4] J. Ahn et al., âLarge language models for mathematical reasoning: Progresses and challenges,â arXiv preprint arXiv:2402.00157, 2024. [5] D. Guo et al., âDeepseek-coder: When the large language model meets programmingâthe rise of code intelligence,â arXiv preprint arXiv:2401.14196, 2024. [6] T. Coignion et al., âA performance study of llm-generated code on leetcode,â in EASE, 2024. [7] I. Bouzenia et al., âRepairagent: An autonomous, llm-based agent for program repair,â arXiv preprint arXiv:2403.17134, 2024. [8] M. Jin et al., âInferfix: End-to-end program repair with llms,â in FSE, 2023, p. 1646â1656. [9] U. Anwar et al., âFoundational challenges in assuring alignment and safety of large language models,â Transactions on Machine Learning Research, 2024. [10] Y. Zhang et al., âThe fusion of large language models and formal methods for trustworthy ai agents: A roadmap,â arXiv preprint arXiv:2412.06512, 2024. [11] Y. Yao et al., âA survey on large language model (llm) security and privacy: The good, the bad, and the ugly,â High-Confidence Computing, 2024. [12] C. Chen et al., âCombating misinformation in the age of llms: Opportunities and challenges,â AI Magazine, p. 354â368, 2024. [13] M. Mazeika et al., âHarmbench: A standardized evaluation framework for automated red teaming and robust refusal,â in ICML, 2024. [14] B. Wang et al., âDecodingtrust: A comprehensive assessment of trustworthiness in gpt models.â in NeurIPS, 2023. [15] L. Sun et al., âTrustllm: Trustworthiness in large language models,â arXiv preprint arXiv:2401.05561, vol. 3, 2024. [16] Y. o. Zhang, âMultitrust: A comprehensive benchmark towards trustworthy multimodal large language models,â NeurIPS, 2024. [17] X. Shen et al., ââdo anything nowâ: Characterizing and evaluating in-the-wild jailbreak prompts on large language models,â in CCS, 2023. [18] A. Zou et al., âUniversal and transferable adversarial attacks on aligned language models,â arXiv preprint arXiv:2307.15043, 2023. [19] A. Wei et al., âJailbroken: How does llm safety training fail?â in NeurIPS, 2023. [20] Y. Liu et al., âJailbreaking chatgpt via prompt engineering: An empirical study,â arXiv preprint arXiv:2305.13860, 2023. [21] Y. Bai et al., âConstitutional ai: Harmlessness from ai feedback,â arXiv preprint arXiv:2212.08073, 2022. [22] J. Dai et al., âSafe rlhf: Safe reinforcement learning from human feedback,â in ICLR, 2024. [23] X. Du et al., âDeepstellar: Model-based quantitative analysis of stateful deep learning systems,â in FSE, 2019. [24] X. Xie et al., âRnnrepair: Automatic rnn repair via model-based analysis,â in ICML, 2021. [25] X. Du et al., âMarble: Model-based robustness analysis of stateful deep learning systems,â in ASE, 2020. [26] X. Ren et al., âDeeparc: Modularizing neural networks for the model maintenance,â in ICSE, 2023. [27] X. Xie et al., âMosaic: Model-based safety analysis framework for ai-enabled cyber-physical systems,â arXiv preprint arXiv:2305.03882, 2023. [28] H. Qi et al., âArchrepair: Block-level architecture-oriented repairing for deep neural networks,â ACM Transactions on Software Engineering and Methodology, vol. 32, no. 5, p. 1â31, 2023. [29] Z. Wei et al., âWeighted automata extraction and explanation of recurrent neural networks for natural language tasks,â Journal of Logical and Algebraic Methods in Programming, 2024. [30] X. Huang et al., âA survey of safety and trustworthiness of deep neural networks: Verification, testing, adversarial attack and defence, and interpretability,â Computer Science Review, vol. 37, p. 100270, 2020. [31] S. MartĂnez-FernĂĄndez et al., âSoftware engineering for ai-based systems: a survey,â ACM Transactions on Software Engineering and Methodology (TOSEM), vol. 31, no. 2, p. 1â59, 2022. [32] D. Song et al., âLuna: A model-based universal analysis framework for large language models,â IEEE Transactions on Software Engineering, 2024. [33] A. Zou et al., âRepresentation engineering: A top-down approach to ai transparency,â arXiv preprint arXiv:2310.01405, 2023. [34] B. Wei et al., âAssessing the brittleness of safety alignment via pruning and low-rank modifications,â arXiv preprint arXiv:2402.05162, 2024. [35] C. Zheng et al., âPrompt-driven llm safeguarding via directed representation optimization,â arXiv preprint arXiv:2401.18018, 2024. [36] Y. Zhang et al., âAdversarial representation engineering: A general model editing framework for large language models,â arXiv preprint arXiv:2404.13752, 2024. [37] A. Vaswani et al., âAttention is all you need,â in NeurIPS, 2017. [38] D. Guo et al., âDeepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,â arXiv preprint arXiv:2501.12948, 2025. [39] T. Zhong et al., âEvaluation of openai o1: Opportunities and challenges of agi,â arXiv preprint arXiv:2409.18486, 2024. [40] A. Radford et al., âImproving language understanding by generative pre-training,â 2018. [41] Z. Wang et al., âA comprehensive survey of llm alignment techniques: Rlhf, rlaif, ppo, dpo and more,â arXiv preprint arXiv:2407.16216, 2024. [42] H. Chen et al., âTowards the worst-case robustness of large language models,â arXiv preprint arXiv:2501.19040, 2025. [43] Z. Wei et al., âPosition: Agent-specific trustworthiness risk as a research priority,â OpenReview preprint, 2025. [44] Y. Bai et al., âTraining a helpful and harmless assistant with reinforcement learning from human feedback,â arXiv preprint arXiv:2204.05862, 2022. [45] C. Wu et al., âMitigating fine-tuning risks in llms via safety-aware probing optimization,â arXiv preprint arXiv:2505.16737, 2025. [46] X. Qi et al., âSafety alignment should be made more than just a few tokens deep,â in ICLR, 2024. [47] H. Chen et al., âUnderstanding pre-training and fine-tuning from loss landscape perspectives,â arXiv preprint arXiv:2505.17646, 2025. [48] A. Q. Jiang et al., âMistral 7b,â 2023. [49] Y. Zeng et al., âHow johnny can persuade llms to jailbreak them: Rethinking persuasion to challenge ai safety by humanizing llms,â in ACL, 2024. [50] H. Jin et al., âGuard: Role-playing to generate natural-language jailbreakings to test guideline adherence of large language models,â arXiv preprint arXiv:2402.03299, 2024. [51] G. Alon et al., âDetecting language model attacks with perplexity,â arXiv preprint arXiv:2308.14132, 2023. [52] N. Jain et al., âBaseline defenses for adversarial attacks against aligned language models,â arXiv preprint arXiv:2309.00614, 2023. [53] Y. Wang et al., âA theoretical understanding of self-correction through in-context alignment,â in NeurIPS, 2024. [54] Y. Xie et al., âDefending chatgpt against jailbreak attack via self-reminders,â Nature Machine Intelligence, 2023. [55] Z. Wei et al., âJailbreak and guard aligned language models with only few in-context demonstrations,â arXiv preprint arXiv:2310.06387, 2023. [56] J. Cui et al., âOr-bench: An over-refusal benchmark for large language models,â arXiv preprint arXiv:2405.20947, 2024. [57] T. Chen et al., âScalable defense against in-the-wild jailbreaking attacks with safety context retrieval,â arXiv preprint arXiv:2505.15753, 2025. [58] K. Simonyan et al., âDeep inside convolutional networks: Visualising image classification models and saliency maps,â arXiv preprint arXiv:1312.6034, 2013. [59] R. R. Selvaraju et al., âGrad-cam: Visual explanations from deep networks via gradient-based localization,â in ICCV, 2017. [60] T. Mikolov et al., âLinguistic regularities in continuous space word representations,â in NAACL, 2013. [61] Q. Zhang et al., âInterpretable convolutional neural networks,â in CVPR, 2018, p. 8827â8836. [62] O. Skean et al., âDoes representation matter? exploring intermediate layers in large language models,â arXiv preprint arXiv:2412.09563, 2024. [63] S. Chalnev et al., âImproving steering vectors by targeting sparse autoencoder features,â arXiv preprint arXiv:2411.02193, 2024. [64] A. Stolfo et al., âImproving instruction-following in language models through activation steering,â arXiv preprint arXiv:2410.12877, 2024. [65] T. Du et al., âAdvancing llm safe alignment with safety representation ranking,â arXiv preprint arXiv:2505.15710, 2025. [66] W. Pan et al., âThe hidden dimensions of llm alignment: A multi-dimensional safety analysis,â arXiv preprint arXiv:2502.09674, 2025. [67] R. Taori et al., âStanford alpaca: An instruction-following llama model,â https://github.com/tatsu-lab/stanford_alpaca, 2023. [68] X. Zhang et al., âDecision-guided weighted automata extraction from recurrent neural networks,â in AAAI, 2021. [69] P. F. Brown et al., âClass-based n-gram models of natural language,â Computational linguistics, vol. 18, no. 4, p. 467â480, 1992. [70] M. Phute et al., âLlm self defense: By self examination, llms know they are being tricked,â arXiv preprint arXiv:2308.07308, 2023. [71] L. Zheng et al., âJudging llm-as-a-judge with mt-bench and chatbot arena,â in NeurIPS, 2023. [72] H. Touvron et al., âLlama 2: Open foundation and fine-tuned chat models,â arXiv preprint arXiv:2307.09288, 2023. [73] J. Bai et al., âQwen technical report,â https://qwenlm.github.io/blog/qwen3/, 2023. [74] X. Geng et al., âKoala: A dialogue model for academic research,â 2023. [75] A. Yang et al., âBaichuan 2: Open large-scale language models,â arXiv preprint arXiv:2309.10305, 2023. [76] P. Chao et al., âJailbreakbench: An open robustness benchmark for jailbreaking large language models,â arXiv preprint arXiv:2404.01318, 2024. [77] L. Zheng et al., âLmsys-chat-1m: A large-scale real-world llm conversation dataset,â 2023. [78] M. Samvelyan et al., âRainbow teaming: Open-ended generation of diverse adversarial prompts,â NeurIPS, 2024. [79] T. Xie et al., âSorry-bench: Systematically evaluating large language model safety refusal,â in ICLR, 2025. [80] L. Jiang et al., âWildteaming at scale: From in-the-wild jailbreaks to (adversarially) safer language models,â 2024. [81] N. Varshney et al., âThe art of defending: A systematic evaluation and analysis of llm defense strategies on safety and over-defensiveness,â in Findings of ACL, 2023. [82] C. Omlin et al., âHeuristics for the extraction of rules from discrete-time recurrent neural networks,â in IJCNN, 1992. [83] C. W. Omlin et al., âExtraction of rules from discrete-time recurrent neural networks,â Neural networks, vol. 9, no. 1, p. 41â52, 1996. [84] G. Weiss et al., âExtracting automata from recurrent neural networks using queries and counterexamples,â in ICML, 2018. [85] Z. Wei et al., âExtracting weighted finite automata from recurrent neural networks for natural languages,â in ICFEM, 2022. [86] Y. Zhang et al., âBoosting jailbreak attack with momentum,â in ICASSP, 2025. [87] T. Markov et al., âA holistic approach to undesired content detection,â arXiv preprint arXiv:2208.03274, 2022. [88] Z. Wang et al., âSelf-guard: Empower the llm to safeguard itself,â arXiv preprint arXiv:2310.15851, 2023. [89] D. Wu et al., âLlms can defend themselves against jailbreaking in a practical manner: A vision paper,â arXiv preprint arXiv:2402.15727, 2024.