$\text{R}^{2}$-Guard: Robust Reasoning Enabled LLM Guardrail via Knowledge-Enhanced Logical Reasoning
\faWarning WARNING: This paper includes content that may be considered offensive.

To address these limitations, we propose $\text{R}^{2}$-Guard, a robust and reasoning enabled LLM guardrail. $\text{R}^{2}$-Guard consists of two main components: (1) a data-driven category-specific learning component, and (2) a knowledge-enhanced reasoning component. The pipeline of $\text{R}^{2}$-Guard is illustrated in Figure 1. The category-specific learning component takes the LLM prompt as input and computes the probability that the prompt falls into different unsafe categories (e.g., the self-harm predictor assesses the likelihood that the prompt contains self-harm-related content). These unsafety probabilities are then forwarded to the reasoning component, which makes the final prediction of the overall probability that the prompt is unsafe based on logical inference. We employ PGMs to implement the reasoning component. By incorporating safety knowledge into the PGMs, we perform probabilistic inference for the final prediction reasoning.

We map the safety knowledge rules such as the relationships among safety categories as first-order logical rules, which are built upon two types of logical variables, the target logical variable which presents the final prediction (i.e., “unsafe") and the category logical variable which is realted to different safety categories (e.g., “self-harm", “sexual"). $\text{R}^{2}$-Guard encodes two types of safety knowledge: (1) direct rules with the form that category logical variables implicate the target logical variable (e.g., $``\textit{self-harm}"\implies``\textit{unsafe}"$), and (2) indirect rules that build implication logics among different category logical variables (e.g., $``\textit{self-harm/instructions}"\implies``\textit{self-harm}"$). Each logical rule is associated with a knowledge rule weight to specify the importance of the knowledge rule to the moderation task. These rules are integrated into probabilistic graphical models (PGMs), employing either Markov logic networks with complete knowledge compilation (Section 3.2) or probabilistic circuits with our improved graph structure for a better precision-efficiency balance (Section 3.3). Through probabilistic inference on these PGMs, the system mimics human logical deduction, initially understanding the semantics and relationships among safety categories (via indirect rules) and subsequently deducing prompt unsafety based on all considered categories (via direct rules). $\text{R}^{2}$-Guard facilitates effective and robust detection of unsafe content through explicit logical inference based on given safety knowledge while allowing for easy adaptation to new safety categories by merely editing the PGM reasoning component.

In $\text{R}^{2}$-Guard, we consider $n$ logical variables taking binary values (i.e., $0$ or $1$), including $n-1$ category logical variables $\{v_{c}^{(i)}\}_{i=1}^{n-1}$ (e.g., $``\textit{self-harm}"$, $``\textit{sexual}"$) and $1$ target logical variable $v_{t}$ (i.e., $``\textit{unsafe}"$). Given any input or output LLM prompt $x$, we denote ${\bm{p}}(x)=[p_{1}(x),...,p_{n}(x)]$ as a conditional unsafety likelihood vector for $n$ logical variables such that $p_{i}(x)={\mathbb{P}}[v_{c}^{(i)}=1|x]$ for $i\in\{1,...,n-1\}$ and $p_{n}(x)={\mathbb{P}}[v_{t}=1|x]$. The unsafety likelihood vector ${\bm{p}}$ can be computed by the data-driven category-specific learning component and serves as the input to the reasoning component, as shown in Figure 1. Suppose that we consider $L$ direct and indirect logical rules $\{R_{i}\}_{i=1}^{L}$, each associated with a knowledge weight $w_{i}\in{\mathbb{R}}~{}(i\in\{1,2,...,L\})$.

We define a possible world $\mu\in\bm{M}=\{0,1\}^{n}$ as a possible assignment to $n$ logical variables such that $\mu_{i}=v_{c}^{(i)}$ for $i\in\{1,..,n-1\}$ and $\mu_{n}=v_{t}$. Based on it, we define the factor function of a possible world $F:\{0,1\}^{n}\mapsto{\mathbb{R}}^{+}$ which takes as input a possible world $\mu$ and outputs the factor value of the world as the following:

where $\mathbb{I}[\mu\sim R_{i}]=1$ indicates that the world $\mu$ follows the logical rule $R_{i}$, and otherwise $\mathbb{I}[\mu\sim R_{i}]=0$. The factor function of a possible world $\mu$ given prompt $x$ consists of two parts: (1) data-driven likelihood, which computes the joint likelihood of the assignments to $n$ logical variables based on unsafety likelihood vector ${\bm{p}}(x)$ provided by category-specific learning models, and (2) logical likelihood measuring how likely the world conform to the defined logical rules, which computes the exponential-summation of the knowledge weights of satisfied logical rules in the possible world $\mu$. In summary, the factor function $F(\mu|x)$ computes the likelihood of the world $\mu$ given prompt $x$. The factor function involves the data-driven likelihood by data-driven category-specific guardrail models and the logical likelihood, which serves as a correction scalar according to the conformity of the world $\mu$ to the safety knowledge space.

$\text{R}^{2}$-Guard eventually outputs the probability that the given prompt $x$ is unsafe (i.e., ${\mathbb{P}}[\text{``unsafe"}=1|x]$ or ${\mathbb{P}}[\mu_{n}=1|x]$). This requires a marginal probability computation which marginalizes over all other logical variables as the following:

where the numerator sums the likelihoods of possible worlds in which the target logical variable is assigned as unsafe (i.e., $\mu_{n}=1$), and the denominator computes the partition function or normalization constant, which is the sum of the likelihoods of all possible worlds.

Probabilistic circuits (PCs) [8, 9, 22, 16, 7, 41] are a more expressive type of PGM compared to MLNs. PCs can represent a wide range of probabilistic distributions over a set of random variables through summation and multiplication operations. Structurally, PCs are organized as tree graphs, where leaf nodes represent individual probabilistic distributions of variables and multi-layered internal nodes capture their interconnections. In $\text{R}^{2}$-Guard, we exploit the observation that certain safety categories exhibit low logical correlation to each other (e.g., $``\textit{self-harm}"$ and $``\textit{sexual}"$ related categories). Therefore, we apply clustering algorithms to partition category logical variables and position different clusters of safety types in different layers of the PC graph, as illustrated in Figure 1. Each PC layer is thus able to concentrate on a specific type of safety knowledge (e.g., “self-harm" or “sexual") and perform logical inference within that layer, emulating MLN inference locally as shown Equation 2. This layered design facilitates a sequential reasoning process that conducts logical inference across different types of safety knowledge step by step, ultimately generating a final prediction. By segregating logically less correlated categories into separate layers, we reduce low-yield interactions among these logical variables, thereby enhancing inference efficiency while maintaining high reasoning precision.

In line 1, we first represent the category logical variables $\{v_{c}^{(i)}\}_{i=1}^{n-1}$ and the set of implication rules in a directed graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, where $\mathcal{V}~{}(|\mathcal{V}|=n-1)$ corresponds to $n-1$ category logical variables and the edges denote the logical implications: $\mathcal{E}_{ij}\in\mathcal{E}\iff(\mathcal{V}_{i}\implies\mathcal{V}_{j})\in\{R_{i}\}_{i=1}^{L}$. In line 2, we apply the spectral clustering algorithm [50] to the knowledge graph $\mathcal{G}$ to obtain $N_{c}$ clusters, each focusing on a specific type of safety knowledge. From lines 3 to 7, we perform layerwise sequential reasoning on the PC graph, where each layer corresponds to a specific cluster. Specifically, we use the unsafety likelihood vector for the categories in the cluster from category-specific learning models and the predefined safety knowledge to perform local MLN reasoning as Equations 1 and 2.

Given the layerwise reasoning pattern on tree graphs, the computational complexity of PC reasoning is $\mathcal{O}(\sum_{i=1}^{N_{k}}2^{|\bm{C}_{i}|})$, where $|\bm{C}_{i}|$ is the size of the $i$-th cluster $\bm{C}_{i}$. Given that $\sum_{i=1}^{N_{k}}|\bm{C}_{i}|=n-1$, the complexity of PC reasoning improves from the exponential-sum order $\mathcal{O}(2^{\sum_{i=1}^{N_{k}}|\bm{C}_{i}|})$ (MLN reasoning complexity) to a sum-exponential order $\mathcal{O}(\sum_{i=1}^{N_{k}}2^{|\bm{C}_{i}|})$. In practice, the safety categories in regulations are well-defined, leading to generally uniform partitions across different clusters [31, 36, 17, 34]. Consequently, PC inference empirically introduces significant efficiency improvements, as shown in Section A.3. The results in Table 5 in Section A.3 show that PC reasoning achieves comparable performance in content moderation while requiring only 6% of the inference time needed for MLN reasoning.

We evaluate the guardrail models on six datasets including five standard safety datasets OpenAI Mod [31],ToxicChat [26], XSTest [42], Overkill [46], BeaverTails [19] and our new safety dataset TwinSafety, introduced in Section 4. We consider four types of strong guardrail models as baselines: (1) industry guardrail APIs from Detoxify [2], Perspective [24], Azure [1], and OpenAI Mod [31], (2) fine-tuned guardrail model LlamaGuard [17] and ToxicChat-T5 [26], (3) LLM-based guardrail model via chain-of-thought prompting (CoT) [54], and (4) ensemble-learning based guardrail models [14, 60]. We directly evaluate the likelihood of unsafety by different APIs. We keep the default prompt template and parameters in Llamaguard and ToxicChat-T5. We use Llama2-7b-chat as the inference model for CoT and carefully select $3$ representative examples from corresponding datasets and manually write the reasoning process as demonstrations. Ensemble learning takes the maximal unsafety scores of category-specific learning models for different categories as the prediction. We use the category-specific learning models from OpenAI Mod, LlamaGuard, and ToxicChat-T5 since they demonstrate high guardrail performance empirically. $\text{R}^{2}$-Guard leverages the same category-specific learning models as the ensemble learning to construct the category-specific learning component and performs logical inference on PGMs, as illustrated in Section 3. We consider both the MLN inference in Section 3.2 and PC inference in Section 3.3 and refer to them as $\text{R}^{2}$-Guard (MLN) and $\text{R}^{2}$-Guard (PC). Following literature[17, 31, 26], We leverage AUPRC as the metric to evaluate the ability of guardrail models to discriminate between safe and unsafe prompts.

The results in Table 2 demonstrate that $\text{R}^{2}$-Guard outperforms various strong baselines while maintaining comparable inference time to local guardrail models such as LlamaGuard. The effectiveness of $\text{R}^{2}$-Guard surpasses CoT reasoning, which implicitly facilitates reasoning through in-context learning. This highlights the power of explicit reasoning by encoding safety knowledge and performing probabilistic inference on MLN and PC graphs. Compared to ensemble learning, the effectiveness of $\text{R}^{2}$-Guard underscores the importance of modeling interactions among unsafety categories and systematically performing inference for guardrails. In addition, $\text{R}^{2}$-Guard leads to marginal runtime overhead with logical inference.

Our TwinSafety dataset leads to overall lower AUPRC on different guardrail models, demonstrating the challenge of our datasets and motivating the development of future guardrail models.

Jailbreak attacks aim to bypass the detection mechanisms of guardrail models by injecting optimized strings. Therefore, it is crucial to evaluate the robustness of guardrail models against these attacks to ensure the security of LLM systems. We consider three types of SOTA jailbreak attack algorithms: (1) the white-box adaptive attack GCG [66], which optimizes an adversarial suffix via token gradients; (2) the black-box attack AutoDAN [28], which leverages genetic algorithms to optimize jailbreak prompts from a pool of seed prompts; and (3) the black-box LLM-based jailbreak algorithms PAIR [5] and TAP [33], which prompt LLMs to generate and refine jailbreak prompts through feedback from target models. Since GCG is a white-box attack, we cannot access the model weights for API-based guardrail models such as OpenAI Mod. Therefore, we consider three types of strong GCG-optimized adversarial suffixes on surrogate models: (1) universal strings optimized to jailbreak multiple LLMs (GCG-U1, GCG-U2); (2) jailbreak strings against the safety-aligned LLM Vicuna-7B (GCG-V) and the SOTA guardrail model LlamaGuard (GCG-L); and (3) jailbreak strings optimized against the distilled Gemma-2B model of $\text{R}^{2}$-Guard (GCG-R). Following the literature [28, 5, 33], we evaluate the robustness of the guardrail models using AdvBench [66], which consists solely of unsafe prompts, and measure the unsafety detection rate (UDR), the portion of flagged unsafe prompts with threshold $0.5$ (i.e., the prompt is recognized as unsafe if the unsafety probability exceeds $0.5$). Additional details are provided in Section A.1.

The results in Table 3 demonstrate that $\text{R}^{2}$-Guard is more robust against multiple SOTA jailbreaks compared to other guardrail models. Both universal jailbreak strings (GCG-U1, GCG-U2) and optimized jailbreak strings using safety-aligned LLMs (GCG-V) and the guardrail model LlamaGuard (GCG-L) do not perturb the UDR of $\text{R}^{2}$-Guard. Even more adaptive GCG attacks against the distilled model of $\text{R}^{2}$-Guard (GCG-R) and SOTA black-box attacks (AutoDAN) only slightly decrease the UDR of $\text{R}^{2}$-Guard, and $\text{R}^{2}$-Guard still outperforms other guardrail models by a significant margin. We evaluate UDRs against PAIR and TAP in Table 4 in Section A.2, which shows that the UDR of $\text{R}^{2}$-Guard is decreased but remains much higher than UDRs of other models. This reduction is because PAIR and TAP may reformulate the original prompt so that the modified prompt is semantically less harmful (e.g., reformulating "grab the gun" to "grab the water gun"), which highlights the need for future work to develop a fairer benchmark in this scenario.