ROOT DEFENCE STRATEGIES: ENSURING SAFETY OF LLM AT THE DECODER LEVEL

## Abstract

Abstract Large language models (LLMs) have demonstrated immense utility across various industries. However, as LLMs advance, the risk of harmful outputs increases due to incorrect or malicious instruction prompts. While current methods effectively address jailbreak risks, they share common limitations:
1) Judging harmful responses from the prefill-level lacks utilization of the model’s decoding outputs, leading to relatively lower effectiveness and robustness. 2) Rejecting potentially harmful responses based on a single evaluation can significantly impair the model’s helpfulness.
This paper examines the LLMs’ capability to recognize harmful outputs, revealing and quantifying their proficiency in assessing the danger of previous tokens. Motivated by pilot experiment results, we design a robust defense mechanism at the decoding level.
Our novel decoder-oriented, step-by-step defense architecture corrects harmful queries directly rather than rejecting them outright. We introduce speculative decoding to enhance usability and facilitate deployment to boost secure decoding speed. Extensive experiments demonstrate that our approach improves model security without compromising reasoning speed. Notably, our method leverages the model’s ability to discern hazardous information, maintaining its helpfulness compared to existing methods.

## 1 Introduction

In recent years, significant progress has been made in developing large language models (LLMs). Meanwhile, the safety of LLMs has attracted significant attention from the research community and industry (Weidinger et al., ; Achiam et al., ; Wu et al., ). One of the primary safety concerns is jailbreaking, where malicious actors or errant inputs prompt LLMs to produce harmful or inappropriate content, effectively bypassing ethical guidelines. Many attempts have been made to address these risks. For example, Meta has implemented several strategies in both pre-training and fine-tuning phases to improve the safety of their Llama-series models (Touvron et al., ; Dubey et al., ). Despite these efforts, some studies have reported that focusing too narrowly on safety may diminish the models’ general capability (Bai et al., ; Huang et al., ). Therefore, enhancing LLMs’ safety without compromising their utility has become a critical area of research.

Recent defense strategies against jailbreaks can be roughly categorized into two groups (as shown in Figure 1). The first group is prefill-level defense (Wu et al., ; Phute et al., ; Zheng et al., ). It enhances the models’ protective capabilities by integrating additional security measures into the initial prompts (prefills) or refining their representation. However, this approach primarily depends on user inputs to detect harmful responses, making it susceptible to rapidly advancing jailbreaking techniques. Moreover, this reliance can lead to inaccuracies in interpreting user intentions, thereby reducing the overall utility of the LLMs. Another group of methods is response-level defenses (Phute et al., ; Xu et al., ). It involves using safety filters that assess the potential harmfulness of model-generated responses. This method focuses on the output of the models, potentially offering improved performance by directly addressing the content generated. However, this strategy typically involves a single evaluation point, which may result in false positives that could diminish the model’s utility by restricting benign responses.

Figure: Figure 1: Examples of two kinds of imperfect defenses and RDS. (a) Prefill-level defenses fail to refuse the harmful query with $N$ harmful tokens. (b) Response-level defenses judge the whole output in a single-point evaluation without consideration of the prefill. (c) RDS conducts step-by-step assessments for each sampled token to enhance the security of LLMs at the decoder level.
Refer to caption: extracted/5909778/intro.png

In practice, jailbreak instructions can bypass the prefill-level defenses and achieve their purposes in the model’s response Wei et al. (). Therefore, assessing jailbreak behavior in LLMs should focus on decoding dimensions, including the context of both the prefill and the model’s response. We aim to directly address and rectify jailbreak behavior by focusing on the decoding level. (Zheng et al., ) has demonstrated models’ ability to distinguish between harmful and benign prefill. This raises the question: Can LLMs extend this discriminative capability to their own decoding? To investigate this hypothesis, we conduct a series of preliminary experiments to explore the model’s ability to discern its own decoding. Specifically, we evaluate five open-source LLMs and visualize the hidden state of the decoding on a token-by-token basis. We observe that LLMs cannot distinguish harmful tokens from harmless tokens in one step, but it can achieve identification through multi-step judgment at the decoding, especially for harmful prefill.

Based on pilot results, we introduce a novel decoder-oriented defense, termed RDS, defencing by step-by-step evaluation during inference. Informed by the discriminative capability of LLMs on decoding, RDS utilizes a trainable classifier to assess the harmfulness of candidate tokens during sampling and adjust their logits accordingly. Subsequently, RDS reorders the candidate tokens and prioritizes the token with lower harmfulness at each step to ensure a safe response iteratively. The step-by-step safe generation provides a root defense on LLM’s decoding (encompassing the context of both prefill and response) perspective and multi-step evaluation. Furthermore, speculative decoding is incorporated into RDS for hidden state prediction to enhance the generation speed, potentially achieving a more fundamental and efficient defense mechanism.

We evaluate RDS on five LLMs and a series of harmful and harmless query benchmarks. Experimental results demonstrate that RDS outperforms existing approaches in terms of both security and helpfulness, reducing compliance with harmful queries from 14.4% to 2.4% on Xstest ((Röttger et al., )) (without safety prompt) and increasing token generation speed by $\mathbf{2.12\times\sim 3.09\times}$ compared to other baselines. We hope this method offers a new perspective to security defense, i.e., assessing the security of a problem from the decoding level, thereby achieving a root defense effect.

## 2 Preliminary: Decoding-level Defence

In this section, we design a series of experiments to evaluate the capability of LLMs to discriminate between harmful and benign outputs at the decoding stage. We first outline the rationale for shifting focus from prefill analysis to decoding, followed by the details of our experimental setup. Finally, we summarize the experimental results and provide a deeper analysis of their implications.

### 2.1 LLMs’ Discriminative Capability of Decoding

The prefill stage for LLMs typically includes a user query, often accompanied by prefixed or suffixed elements such as system prompts. Previous studies have demonstrated that LLMs can discriminate between different types of prefill and use this ability to enhance safety mechanisms. However, relying solely on prefill analysis for security evaluations presents significant limitations: 1) Jailbreaking behaviors often manifest in the model’s output, and focusing solely on prefill may overlook these behaviors, compromising overall robustness; 2) Evaluation based purely on prefill places excessive dependence on the model’s initial discriminative capacity, and a single-stage evaluation may lead to rejecting outputs prematurely, reducing the model’s utility.

To address these limitations, we explore whether LLMs can discriminate harmful from benign content during decoding, which encompasses both the prefill and the model’s generated outputs. If LLMs can reliably evaluate the safety of their own outputs in real time, they can offer a more comprehensive and proactive approach to security. Decoding-based defenses leverage the dynamic nature of model outputs, allowing for a more fundamental and continuous risk assessment.

### 2.2 Experiments

In this section, we evaluate five open-source LLMs and utilize Principal Component Analysis (PCA) to visualize their hidden states during the decoding process. To facilitate classifier training, we curate Custom from the DRO (Zheng et al., ) as the training dataset of the classifier, consisting of 100 harmful and 100 benign queries, all beginning with the phrase “How to”. The evaluated LLMs are accessible on HuggingFace: Llama-2-chat-7B (Touvron et al., ), Llama-3-8b-Instruct (AI@Meta, ), Qwen2-7B-Instruct (Yang et al., ), Vicuna-7B-v1.3, and Vicuna-13B-v1.3 (Chiang et al., ). Notably, some models, such as those in the Llama series, have undergone extensive safety alignment.

Following previous studies (Zheng et al., ), we use the hidden states $\mathbf{h}$ of the prefill at the top layer of the model for classifier training. Details of the classifier’s training objective is provided in Appendix C. The classifier is applied to the decoding sampling strategy. Refusal responses often start with special tokens, such as “I’m sorry” or “As an AI”, which distinguish them from compliant responses. Given the distinct nature of refusal and compliant responses, particularly at the start of the output, we focus on the first three tokens and the last token of the output (i.e., $i\in{1,2,3,N}$, where $N$ is the length of the output), visualizing the hidden states of each token from the top layer.

The responses to harmful queries are assessed using Llama-guard (Bhatt et al., ), which is a safety classification model based on LLaMA-2 (Touvron et al., ). Benign responses are evaluated through string matching. If refusal strings are identified in the response, it is categorized as a refusal response; otherwise, it is not. A compliant answer is assigned an evaluation score $s$ of 1, otherwise 0. Lower scores are preferable for harmful queries, while higher scores are ideal for benign queries. In the preliminary experiment, we sample one response for each query. The initial defense of these five LLMs is presented in Appendix D.

### 2.3 Visualization Analysis

We apply PCA to visualize the hidden state and compute the first four principal components of the hidden states. Figure 2 respectively show the visual results from the first to the third token and the last token of the refusal responses for harmful queries (represented by “refusal+i”) and compliance for benign queries (represented by “compliance+i”). See the visual results from the 4-th token to the 7-th in Appendix E. In Figure 2, harmful and benign queries can be clearly distinguished at the hidden state layer. The boundary, depicted by a black dashed line, separates these two queries using training with only a few parameters. LLMs naturally can discern the harmfulness of inputs.

Can LLMs extend this discriminative capability to their own decoding? In Figure 2, from 1-th to 3-th token, almost all the tokens to benign queries maintain at the benign side for all LLMs. However, for harmful queries, although refusal tokens refer to safe responses, the first few tokens are not classified correctly. Instead, we observe that the refusal tokens to harmful queries tend to converge towards the benign side with a smaller offset than harmful tokens. This shift is most pronounced in llama-2-chat-7B and Qwen2-7B-Instruct, whose responses to harmful queries are all refusal. Even on models with poor defense, this shift still exists. We interpret these observations as the discriminative capacity of LLMs to their own decoding.

Can we consider LLMs’ output harmful based on a single judgment? Figure 2 illustrates that the classifier cannot accurately determine whether the output is harmful based solely on the model’s overall decoding (i.e., the complete output). Even current advanced methods cannot guarantee 100% filtering. Considering the experimental results, we believe that making a single-step judgment is insufficient to determine if the output is harmful. In conjunction with Figure 2, although the model cannot make an accurate judgment in one attempt, it can achieve better discrimination through a step-by-step evaluation of the decoding. Therefore, we believe a gradual assessment approach at the decoding can lead to more effective defense mechanisms.

Figure: Figure 2: Performance of the classifier at the decoding from the 1-th to 3-th token and last token of the output. The red crosses represent the hidden states for harmful queries, while the blue circles represent the hidden values for benign queries.
Refer to caption: extracted/5909778/ideal/Llama-2-7b-chat-hf_2.jpg

## 3 Methodology

Motivated by validating the capability to recognize responses, we propose RDS to ensure the safety of LLMs at the decoder level. The architecture of RDS is illustrated in Figure 3. We design a step-by-step defense mechanism that directly corrects the harmful token into a safe token when generating the response. Additionally, we introduce speculative decoding into RDS to speed up token generation. Benefitting from step-by-step safe generation and speculative decoding, RDS achieves root security without compromising helpfulness and speed.

Figure: Figure 3: RDS comprises two key modules: 1) Step-by-step token generation: The root classifier is designed based on the discriminative capacity of queries. By adjusting the logits of candidate tokens, RDS reorders the token and prioritizes the benign token. 2) Speculative decoding: RDS predicts the hidden state from speculative decoding instead of multiple transformer blocks.
Refer to caption: extracted/5909778/RDS.png

### 3.1 Problem Formulation

Let $x_{i}$ as the model’s decoding at step $t_{i}$, $c_{i}=f(x_{i})$ represents the score of the sampled token $x_{i}$ calculated by the classifier $f(\cdot)$, where the training object of the classifier is illustrated in Appendix C. RDS aims to minimize $c_{i}$ at each step. We formulate this process as follows:

$$ $\underset{x_{i}}{\min}\sum_{i=1}^{N}c_{i}\quad;\quad x_{i}=\text{LLM}(x_{i-1}; \mathbb{C}_{i})$ (1) $$

where $N$ is the length of outputs, and at each step $t_{i}$, the LLM obtains prior decoding $x_{i-1}$ and the harmful results $\mathbb{C}_{i}$ of candidate tokens to generate next token $x_{i}$. RDS constructs the candidate tokens according to the logit value and samples a new token from the candidate tokens. By ensuring the security of each step, RDS promises a safe response.

### 3.2 Step-by-step safe generation

During the autoregressive decoding of LLMs, LLM maps the hidden state of its decoding $x_{i-1}$ at step $t_{i-1}$ to the vocabulary dimension and sample the next token by top-$k$ (Fan et al., ):

$$ $\mathbb{I}_{i},\mathbb{V}_{i}=\mathrm{Topk}(\mathrm{softmax}(\mathbf{l}_{i-1})),$ (2) $$

where $\mathbf{l}_{i-1}=\text{LM_Head}(\mathbf{h}_{i-1})$ represents the logits of the decoding at step $t_{i-1}$, $\mathbf{h}_{i-1}$ represents the hidden state of the decoding at step $t_{i-1}$, $\mathbb{I}_{i}$ and $\mathbb{V}_{i}$ represent the set of top-$k$ candidate tokens and the logits values of these candidate tokens, respectively.

Safety disclaimers frequently rank among the top tokens (Zheng et al., ) in the inference process. To enhance security, RDS aims to adjust the logits of these tokens further. The classifier from the pilot experiments is integrated into the sampling strategy during decoding. This integration provides a real-time safety assessment of candidate tokens, adjusting the top-$k$ tokens to safer alternatives, ensuring the safety of the next generated token. Consequently, the computation of $c_{i}$ in Equation (1) is detailed into the following components:

$$ $\mathbf{m}_{k}={\mathbf{V}}^{T}(\mathbf{h}_{i}^{k}-\mathbf{u}),$ (3) $$

$$ $c_{k}={\mathbf{W}}^{T}\mathbf{m}_{k}+\mathbf{b},$ (4) $$

$$ $x_{i}=\mathrm{argmax}(\mathbb{C}_{i}),$ (5) $$

where $\mathbf{h_{i}^{k}}$ is the hidden state of the dececoding at step $t_{i}$ concatenated with the candidate token from $\mathbb{I}_{i}$, $\mathbf{m_{k}}\in\mathbb{R}^{m}$ represents the first $m$ principal components of $\mathbf{h_{k}}$, $c_{k}\in\mathbb{R}^{1}$ is the harmful score of the candidate token, $\mathbb{C}_{i}$ is the set of harmful scores of the candidate tokens.

### 3.3 Speculative decoding

RDS leverages the discriminative ability of decoding for defense by computing the harmful score of candidate tokens based on their hidden states. It concatenates decoding at step $t-1$ with candidate tokens to obtain the hidden state at step $t$ resembling speculative decoding processes (such as EAGLE (Li et al., )) that predict hidden states from decoding and tokens. To increase decoding speed, RDS extends EAGLE’s resampling to accelerate generation.

Unlike traditional LLMs that compute hidden state through autoregressive decoding with multiple Transformers blocks, RDS utilizes $\mathrm{EAGLE_Head}$ to predict the hidden state $\mathbf{h}_{i}$ at step $t_{i}$, thereby accelerating the inference process. This prediction is based on the candidate token and the hidden state of decoding at step $t_{i-1}$. The hidden state in Equation (3) can be expressed as:

$$ $\mathbf{h}_{i}^{k}=\mathrm{EAGLE_Head}(\mathbf{h}_{i-1},\mathbf{e}_{k}),$ (6) $$

where $\mathrm{EAGLE_Head}$ consists of a FC layer and a decoder layer from the original LLM; $\mathbf{e}_{k}$ represents the embedding of the candidate token $x_{k}$. After predicting the hidden state at step $t_{i}$, the step-by-step safe token generation is conducted on this predicted hidden state.

We summarize the inference process of RDS as $\mathrm{Draft_Model}$, which can be formulated as:

$$ $x_{N}=\mathrm{Draft_Model}(\mathbf{h}_{0}).$ (7) $$

where $\mathbf{h}_{0}$ denotes the hidden state of the prefill at step $t_{0}$, $x_{N}$ represents the output of LLMs. Equation (7) reveals that RDS only generates the safe response from the hidden state of prefill, without additional LLMs training nor other models introduced.

### 3.4 Highlights

As a decoder-oriented defense, the advantages of RDS are summarized as follows:

First, RDS demonstrates a root defense by leveraging the discriminative capabilities in LLMs’ decoding level. It fully utilizes the model’s understanding of context by evaluating the harmfulness from both input and output dimensions. Guided by a classifier with fewer parameters, RDS identifies harmful tokens during the early inference stage and corrects them to safe tokens, thereby reducing harmfulness in the output. Subsequent experimental results indicate that RDS can enhance the model’s defensive capability without additional training for the LLMs.

Secondly, RDS adopts a step-by-step correction strategy by incrementally adjusting the token logits during the sampling process and progressively correcting harmful labels. Instead of relying on single-point evaluations, RDS improves the safety of LLMs through multi-step evaluations, thereby providing stronger assistance capabilities and a lower false alarm rate for user queries. Furthermore, experiments demonstrate that RDS is more helpful than other methods on various safety benchmarks, further indicating the transferability of RDS.

Finally, to enhance the reasoning speed of RDS and facilitate its practical implementation, we incorporate a speculative decoding structure into the resampling process. It leverages the advantages of the step-by-step mechanism to accelerate the generation process. Experimental results demonstrate that the token generation speed of RDS is approximately $\mathbf{2.12\times\sim 3.09\times}$ faster than that of the baselines. These improvements demonstrate both the effectiveness and efficiency of RDS.

## 4 Experiments

### 4.1 Experimental setup

#### Benchmarks

We evaluate the security and helpfulness of RDS and baselines on five benchmarks.
(1) Custom (Zheng et al., ) contains 100 harmful and 100 harmless queries, which are “How to do” query pairs synthesized from gpt-3.5-turbo.
(2) AdvBench (Zou et al., ) contains 250 harmful queries, we select the top-100 for evaluation.
(3) MaliciousInstruct (Huang et al., ) contains 100 harmful question prompts, such as “How to …?” and “Ideas for…?”. It is specifically designed to encompass various harmful intents (a total of 10), including sabotage, theft, defamation, etc.
(4) Testset (Zheng et al., ) contains 100 harmless queries synthesized from gpt-3.5-turbo to evaluate the false refusals of baselines.
(5) Xstest (Röttger et al., ) contains multiple contrastive queries; we select 100 harmful queries and 100 benign queries in contrastive groups.

#### Models

We deploy RDS on five opensource LLMs, namely Llama-2-chat-7B (Touvron et al., ), Llama-3-8b-Instruct (Dubey et al., ), Qwen2-7B-Instruct (Yang et al., ), Vicuna-7B-v1.3 (Chiang et al., ), and Vicuna-13B-v1.3 (Chiang et al., ). For brevity, we use LLaMA2, LLaMA3, Qwen2, Vicuna-7B, and Vicuna-13B to refer to them, respectively.

**Table 1: Compliance results on harmful benchmarks of baselines and RDS ($\downarrow$). The best results are in bold. We sample five responses for each query. Once a response is compliant ($s_{i}>0$) is the answer treated to be compliant. “base” represents the original LLM.**
| Data | Defense | Vicuna-7B | Vicuna-13B | LLaMA2 | LLaMA3 | Qwen2 | Average |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Custom | base | 19 | 6 | 0 | 0 | 0 | 5.0 |
| safety prompt | 6 | 0 | 0 | 0 | 0 | 1.2 |  |
| Self-Remind | 0 | 0 | 0 | 0 | 0 | 0 |  |
| DRO | 0 | 1 | 0 | 0 | 0 | 0.2 |  |
| Self-Examination | 0 | 0 | 0 | 0 | 0 | 0 |  |
| SafeDecoding | 0 | 0 | 0 | 0 | 0 | 0 |  |
|  | RDS | 0 | 0 | 0 | 0 | 0 | 0 |
| AdvBench | base | 22 | 10 | 0 | 1 | 2 | 7.0 |
| safety prompt | 6 | 0 | 0 | 0 | 0 | 1.2 |  |
| Self-Remind | 0 | 0 | 0 | 1 | 0 | 0.2 |  |
| DRO | 2 | 0 | 0 | 0 | 0 | 0.4 |  |
| Self-Examination | 0 | 0 | 0 | 0 | 0 | 0 |  |
| SafeDecoding | 0 | 0 | 0 | 0 | 0 | 0 |  |
|  | RDS | 1 | 0 | 0 | 0 | 0 | 0.2 |
| Malicious<br>Instruct | base | 16 | 3 | 0 | 0 | 3 | 4.4 |
| safety prompt | 16 | 3 | 0 | 0 | 3 | 4.4 |  |
| Self-Remind | 0 | 0 | 0 | 0 | 0 | 0 |  |
| DRO | 3 | 2 | 0 | 1 | 2 | 1.2 |  |
| Self-Examination | 0 | 0 | 0 | 0 | 0 | 0 |  |
| SafeDecoding | 0 | 0 | 0 | 0 | 0 | 0 |  |
|  | RDS | 1 | 0 | 0 | 0 | 1 | 0.4 |
| Xstest | base | 48 | 12 | 0 | 0 | 12 | 14.4 |
| safety prompt | 0 | 0 | 0 | 0 | 0 | 0 |  |
| Self-Remind | 0 | 0 | 0 | 0 | 0 | 0 |  |
| DRO | 4 | 4 | 0 | 0 | 0 | 1.6 |  |
| Self-Examination | 0 | 0 | 0 | 0 | 0 | 0 |  |
| SafeDecoding | 0 | 0 | 0 | 0 | 0 | 0 |  |
|  | RDS | 0 | 0 | 0 | 0 | 0 | 0 |

#### Baselines

We select five methods as the baselines. Prefill-based defenses contain: (1) safety prompt, which is the official safety prompt of LLaMA-2 illustrated in Appendix F. The safety prompt serves as the system prompt of LLMs. (2) Self-Remind ((Wu et al., )), which encapsulates the user’s query in a system prompt to remind LLMs to respond responsibly. (3) DRO ((Zheng et al., )), which utilizes the distinguished ability at the prefill level to train the safety prompt embedding to improve the moving direction of the input. Response-based defenses contain: (4) Self-Examination ((Phute et al., )), which checks the response by the LLM itself. (5) SafeDecoding (Xu et al., ) that matches the string of safety disclaimers and amplifies the token probabilities of safety disclaimers by training an additional expert model for LLMs.

#### Evaluation Metric

We select five samples for each query. The evaluation method still follows Section 2.2 to judge whether the response is compliant. We take the average evaluation score of all samples as the final score, which can be denoted as:

$$ $s_{i}=\frac{1}{n}{\textstyle\sum_{j=1}^{n}}s_{ij},$ (8) $$

where $n=5$ is the number of response samples. For harmful queries, the threshold is set to 0. For benign queries, the threshold is set to 0.5.

### 4.2 Main Results

Table 1 presents the baselines, and RDS compliance results on harmful datasets. “base” represents the original model. From Table 1, we have the following inclusions.

Firstly, RDS effectively reduces compliance to harmful queries on most datasets. After applying the safety prompt, the defense ability of all LLMs further improves (i.e., vicuna-7B on Xstest). However, safety prompt does not always work (i.e., vicuna-7B on MaliciousInstruct). Compared with prefill-level defenses, RDS demonstrates further advantages (0.4 vs. 1.2 on MaliciousInstruct; 0 vs. 1.6 on Xstest). Compared with response-level defense, RDS is on par with SafeDecoding. These indicate that RDS demonstrates excellent defense ability at the decoder level.

Secondly, though trained on Custom, the classifier retains its discriminative ability on other datasets, which demonstrates the transferability of the classifier and the generalization ability of RDS.

Thirdly, larger models demonstrate better initial defense than smaller ones within the same series of LLMs. For instance, in the Vicuna series models, the 13B model surpasses the 7B model in security (6 vs.19 on Custom; 10 vs. 22 on Advenbench; 24 vs. 43 on MaliciousInstruct; 12 vs. 48 on Xstest). We speculate this is because the larger model achieves better safety alignment during training on the more extensive training data.

**Table 2: Refusal results on benign benchmarks of the baselines and RDS ($\downarrow$). We sample five responses for each query. The query is treated as a refusal if half of the responses ($s<0.5$) are refused.**
| Data | Defense | Vicuna-7B | Vicuna-13B | LLaMA2 | LLaMA3 | Qwen2 | Average |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Custom | base | 0 | 0 | 4 | 0 | 0 | 0.8 |
| safety prompt | 0 | 0 | 5 | 1 | 0 | 1.2 |  |
| Self-Remind | 1 | 3 | 5 | 7 | 1 | 3.4 |  |
| DRO | 0 | 0 | 5 | 0 | 0 | 1.0 |  |
| Self-Examination | 0 | 0 | 100 | 17 | 0 | 23.4 |  |
|  | SafeDecoding | 5 | 6 | 14 | 3 | 3 | 6.2 |
|  | RDS | 0 | 0 | 3 | 0 | 0 | 0.6 |
| Testset | base | 0 | 0 | 1 | 0 | 0 | 0.2 |
| safety prompt | 0 | 2 | 3 | 0 | 0 | 1.0 |  |
| Self-Remind | 3 | 2 | 1 | 8 | 1 | 3.0 |  |
| DRO | 0 | 0 | 3 | 0 | 0 | 0.6 |  |
| Self-Examination | 2 | 1 | 100 | 10 | 0 | 22.6 |  |
|  | SafeDecoding | 4 | 4 | 16 | 2 | 3 | 5.8 |
|  | RDS | 0 | 0 | 1 | 0 | 0 | 0.2 |
| Xstest | base | 4 | 20 | 64 | 12 | 12 | 22.4 |
| safety prompt | 16 | 28 | 88 | 36 | 8 | 35.2 |  |
| Self-Remind | 52 | 48 | 96 | 92 | 24 | 62.4 |  |
| DRO | 32 | 72 | 88 | 36 | 24 | 46.0 |  |
| Self-Examination | 24 | 28 | 100 | 48 | 24 | 44.8 |  |
|  | SafeDecoding | 64 | 72 | 96 | 64 | 60 | 71.2 |
|  | RDS | 0 | 12 | 64 | 12 | 12 | 20.0 |

### 4.3 Helpfulness Analysis

Table 2 presents the refusal results on the benign datasets of baselines and RDS. The higher the refusal score, the greater the loss of helpfulness caused by defense methods.

Compared to all baselines, our model shows fewer refusal results to benign queries (0.6 vs.23.4 on Custom; 0.2 vs. 22.6 on Testset; 20.0 vs. 62.4 on Xstest). Besides, RDS can even correct the original model’s overcorrection (20.0 vs. 22.4 on Xstest). Combined with Table 1, it can be included that RDS can enhance the security of LLMs without compromising its helpfulness.

What’s more, there is an interesting observation that alignment models tend to reject benign queries (LLaMA2 vs. vicuna-7B). This phenomenon is most evident in Xstest, as it contains polysemous words. For instance, “kill” is harmful when referring to “killing a person” whereas it is benign when referring to ”kill time”. The increased refusal results of the alignment models on Xstest also reveal that current safety alignment is, to some extent, word-level alignment (20.0 of RDS vs. 46.0 of DRO). While RDS can correct this over-correction at the token level. In Appendix G, we analyze the reasons for the large number of rejections on Xstest in detail.

### 4.4 Efficiency Analysis

We evaluate the efficiency of baselines and RDS on the Custom (benign). We select one response for each query and collect the number of tokens generated per second in Table 3.

The inference speed of prefill-level defenses nearly matches that of the original LLMs. At the same time, response-level defenses are notably lower than the original model’s when inferencing. Due to the string matching and additional inference of the expert model, SafeDecoding demands supplementary computational resources. Moreover, SafeDecoding necessitates training the expert model for each LLM, which demands exceptionally high computational resources. Conversely, RDS employs step-by-step safe token generation without additional training of LLMs, utilizing speculative decoding to enhance the inference speed, proving significantly more efficient than other baseline methods. With the support of speculative decoding, the generation speed of RDS is 2.12x $\sim$ 3.09x faster than other baselines. We design another variant, “- w/o SD”, that removes the speculative decoding of RDS. In “- w/o SD”, The inference speed significantly drops. We include the reason for the drop that “- w/o SD” predicts the hidden state of candidate tokens by autoregressive decoding.

**Table 3: Number of tokens generated per second of the baselines and RDS (tokens /s $\uparrow$). “- w/o SD” removes the speculative decoding of RDS.**
| Model | vicuna-7B | vicuna-13B | LLaMA2 | LLaMA3 | Qwen2 |
| --- | --- | --- | --- | --- | --- |
| base | 41.68 | 31.74 | 42.30 | 38.77 | 34.61 |
| safety prompt | 43.43 | 32.15 | 43.09 | 40.20 | 36.43 |
| Self-Remind | 31.85 | 25.62 | 32.27 | 29.15 | 40.81 |
| DRO | 43.69 | 32.99 | 43.02 | 39.41 | 35.48 |
| Self-Examination | 32.19 | 25.70 | 25.15 | 25.00 | 39.75 |
| SafeDecoding | 31.99 | 25.32 | 31.71 | 28.75 | 37.30 |
| RDS | 73.17 | 78.29 | 97.77 | 69.98 | 73.46 |
| - w/o SD | 21.25 | 16.85 | 21.04 | 22.25 | 20.39 |

## 5 Related Work

### 5.1 Existing Defenses

Existing safety defenses can be divided into input-based defenses and output-based defenses.

Prefill-level defenses induce LLMs to reject harmful questions by optimizing the input, such as adding a safety system prompt or filtering the input. For instance, IAPrompt ((Zhang et al., )) delves into the intent of input before decoding. Perplexity filtering ((Alon & Kamfonas, )) proposes to detect the adversarial suffixes as the signal of harmful input before generating a response. Modifying the input on the encoder for defense can also be modified for attack ((Zhao et al., )). Modifying input for attacks is a deficient consumption and efficient behavior. At present, multiple methods have successfully carried out jailbreak attacks from user input, such as GCG ((Zou et al., )), Auto-DAN ((Zhu et al., )), Evil Geniuses ((Tian et al., )). Besides, input-based defenses show poor helpfulness with over-defense ((Zhou et al., )).

Response-level defenses enhance the security of LLMs by judging the generated response, which follows the paradigm of generate and then judge. For instance, Self-Examination ((Phute et al., )) checks the response response itself by a pre-defined prompt. SelfDecoding ((Xu et al., )) captures the safety disclaimers from top tokens and amplifies their token probabilities. Response-level defenses must fully generate the output before judging, which affects the model’s efficiency. While RDS monitors the token step-by-step, forcing safe token generation in time.

### 5.2 Speculative Decoding

Traditionally, token generation is performed step-by-step, where the model generates one token for each step by autoregressive decoding. The generated token concatenated to the input serves as the new input for the next step ((Chen et al., )). This approach is straightforward but can be computationally expensive and slow, particularly when generating long text ((Kim et al., )).

Speculative Decoding is an optimization technique used in LLMs to accelerate the process of token generation ((Leviathan et al., ; Chen et al., )). By the Draft-then-Verify paradigm, speculative decoding generates multiple tokens at each step ((Xia et al., )). For example, (Zhang et al., ) proposes to use the same serious but more minor LLM as the draft model without additional training. Not all models have a smaller draft model; self-draft becomes a new paradigm instead of using a separate draft model. For instance, Medusa ((Cai et al., )) incorporates feedforward neural heads atop the decoder to predict tokens in different positions in parallel.

## 6 Conclusions

Our study delves into and confirms the discriminative capacity of LLMs at the decoder level. Through preliminary validation, we indicate that LLMs consistently can discern the harmfulness of output tokens at multiple steps. Motivated by these findings, we propose a Root Defense Strategy originating from the decoding level, namely RDS. The incremental safe token generation process enforces security measures. Furthermore, speculative decoding is introduced in RDS to enhance usability and facilitate deployment. Comparative experimental results demonstrate that RDS offers robust and efficient security defense without compromising utility.

## Appendix A ETHICS STATEMENT

Our research (including papers, code, and data) may have potential adverse effects, as jailbreaks could exploit our methods to target commercial APIs, creating harmful content or privacy violations. Despite the risk of misuse, we consider the current work justifiable. As LLMs perform outstandingly, researchers are presently concentrating on enhancing their helpfulness. Hence, investigating and mitigating security vulnerabilities is crucial, as these models could be vulnerable to black-box attacks when the exact specifics of the target model are undisclosed. The technology we have introduced can effectively thwart such misuse and empower the Red Team to deploy efficiently ahead of time, progressively achieving secure generation starting at the decoder level. In conclusion, our study underscores the significance of security defense starting at the decoder level.

## Appendix B Case study

In this section, we compare baseline and RDS from the perspective of the semantics of refusal response. As shown in Figure 4, there is compliance for the harmful query even if a safety prompt is added to the prefill. Although other baselines have rejected the query, they are semantically singular. However, RDS has started with “The best way of stealing a car is”, but in the following tokens, it gives a rejection of “not advisable”, which refuses to answer the harmful query. This reflects the step-by-step safe token generation of RDS. Furthermore, RDS gives a kind hint after the rejection, which is more in line with human common sense. The above results show the superiority of RDS on semantics.

The results of the symmetry harmless problem on Xstest are shown in Figure 5. For the same behavior “stealing a car”, it is a harmful behavior in a real-world scenario. However, LLMs serve as tools and must respond to video game user requests. We found that all defenses except RDS give a rejection. It is worth noting that although SafeDecoding recognizes GPT4 as a game, it still rejects the query. At the same time, when it replies, RDS gives a friendly hint ( i.e., “it is illegal and not endorsed by Rockstar Games”).

## Appendix C Training process of the classifier

$$ $\mathbf{u}=\frac{1}{N}{\textstyle\sum_{i=1}^{N}}\mathbf{h}_{i},$ (9) $$

$$ $\mathbf{m_{i}}={\mathbf{V}}^{T}(\mathbf{h_{i}}-\mathbf{u}),$ (10) $$

$$ $\hat{c_{i}}={\mathbf{W}}^{T}\mathbf{m_{i}}+\mathbf{b},$ (11) $$

$$ $\mathcal{L}(c_{i},\hat{c}_{i})=\frac{1}{N}{\textstyle\sum_{i=1}^{N}}(c_{i}\log {\hat{c}_{i}}),$ (12) $$

where $\mathbf{u}\in\mathbb{R}^{d}$ is the mean value of all hidden states of queries, and $N$ is the number of queries. $\mathbf{V}\in\mathbb{R}^{d\times m}$ represents the $m$ principal components. $\mathbf{W}\in\mathbb{R}^{1\times d}$ and $\mathbf{b}$ correspond to the trainable parameters of the classifier. $\hat{c}_{i}$ and $c_{i}$ represent the predicted score and the query label.

Figure: Figure 4: An example of baselines and RDS for a harmful query. The user inputs “What is the best way of stealing a car?” as the query. The safety prompt fails to defend against this harmful query. At the same time, other defenses offer refusal answers.
Refer to caption: extracted/5909778/case_harm.png

Figure: Figure 5: An example of baselines and RDS for a barmless query. The user inputs “What is the best way of stealing a car in GTA 4?” as the query. All defenses except RDS give a rejection.
Refer to caption: extracted/5909778/case_harmless.png

## Appendix D Initial defense ability of LLMs in preliminary experiment

Table 4 shows the results of the five LLMs for harmful and harmless queries on Custom.

**Table 4: Defence performance of the five models on Custom.**
| Models | Compliance on Harmful Queries $\downarrow$ | Refusal on Harmless Queries $\downarrow$ |
| --- | --- | --- |
| vicuna-7B-v1.3 | 5 | 3 |
| vicuna-13B-v1.3 | 0 | 0 |
| llama-2-chat-7B | 0 | 0 |
| LLaMA3-Instruct-8B | 9 | 0 |
| Qwen2-7B-Instruct | 0 | 0 |

## Appendix E Visual results of LLMs at deeper decoding

Figure 6 respectively shows the visual results from the 1-th to 3-th token and the last token of Llama3-8B-Instruct and Vicuna-7B-v1.3s. Figure 8 respectively shows the visual results from the 4-th to 7-th token of the five LLMs.

Figure: Figure 6: Performance of the classifier at the decoding from the 1-th to the 3-th token and last token of the output. The red crosses represent the hidden states for harmful queries, while the blue circles represent the hidden values for benign queries.
Refer to caption: extracted/5909778/ideal/Llama-3-8B-Instruct_1.jpg

## Appendix F Safety prompt

We illustrated the safety prompt of LLaMa-2 Official and Self-Remind in Figure 7.

Figure: Figure 7: Illustration of safety prompt used in LLaMa-2 Official and Self-Remind.
Refer to caption: extracted/5909778/safety_prompt.png

## Appendix G Analysis on Xstest

In Table 2, LLMs exhibit a significantly higher rate of rejections on the Xstest than other datasets. Queries in Xstest contain adverse words. We posit that the elevated rejection from LLMs’ compromised ability to discern inputs due to the presence of adverse words.

Figure 9 illustrates the classifier’s results on out-of-domain datasets. Notably, LLMs demonstrate robust classification capabilities in Custom, Advbench, MaliciousInstruct, and testset. On Xstest, harmful and benign inputs are interspersed, leading to a failure in differentiation by LLMs. This inability to discriminate inputs contributes to a reluctance to respond to benign queries. Consequently, defense strategies based on input, such as DRO, result in severe rejections on Xstest. In contrast, RDS solely assesses outputs for defense, disregarding inputs. Thus, the superior utility of RDS on Xstest underscores the benefits of defense mechanisms at the decoder level.

Figure: Figure 8: Performance of the classifier at the decoding from the 4-th to 7-th token.
Refer to caption: extracted/5909778/ideal/Llama-2-7b-chat-hf_5.jpg

Figure: Figure 9: Performance of the classifier at all datasets. (1) Custom is the training data of the classifier. (2) Advbench and MaliciousInstruct are the harmful benchmark. Testset is a harmless benchmark. (3) Xstest has both harmful and harmless queries in symmetry pairs.
Refer to caption: extracted/5909778/xstest/Llama-2-7b-chat-hf_cutsom.jpg

## Appendix H Limitations

RDS filters safe tokens among the top-$k$ tokens of LLMs. If the security disclaimer does not exist in the top-$k$ tokens, RDS maybe cannot generate a security answer. In addition, for harmless queries, if the LLMs tend to give a rejection, i.e., the top-$k$ answers are all security disclaimers, RDS will also generate a rejection. How to optimize the model’s overcorrection while ensuring the security of LLMs will be the future research point.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to the weakness part.
soundness: 2
strengths: The paper is easy to follow. Discussions are thorough in the paper.
summary: The paper proposes a defense mechanism. They train a classifier to guide the model to exclude certain unsafe words during decoding. They also integrate speculative decoding the expedite the process.
weaknesses: 1. Details regarding the training the classifier are missing in main text. However, this is the main contribution of the paper.

2. What is the advantage of training the classifier compared to simply maintaining a list of unsafe words? Are there direct evaluations on the trained classifier or case study?

3. The results in Table 1 (main results) seem not demonstrate the effectiveness of the method. RDS achieves similar results to several baselines.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: See my (implicit) questions in the weaknesses section.
soundness: 2
strengths: - LLM safety is a crucial topic, where developing strong defense strategies against jailbreaks can be a potentially highly valuable contribution.

- Addressing safety during the decoding stage is a promising approach for mitigating jailbreaks.

- Using a small-scale classifier during generation might point towards efficient decoding-time defenses.

- I like the authors’ choice of presenting preliminary exploratory experiments that motivate the developed method early in the paper.
summary: The paper introduces a jailbreak defense strategy (“Root Defence Strategies”—RDS) that makes use of a classifier applied at each step of generation to the next top k tokens to guide the generation towards only producing safe tokens. The authors integrate their method with speculative decoding strategies to make it faster. Their evaluation results across 5 jailbreak datasets and 5 models show early promise.
weaknesses: While I find certain aspects of the paper interesting/promising, I believe that currently there are some severe methodological and presentation weaknesses that limit the overall contributions of the paper:

**Writing, clarity, and accuracy**

Although, as mentioned, I appreciate the structure of the paper, I find the writing and presentation very difficult to follow and often erroneous. Certain claims, choices, or statements are also made without clear foundations.

For instance, the sentence starting on line 111 refers to “previous studies”, but does not cite them.

Or on lines 195 to 198, the authors write about their motivating experiments presented in Figure 2 that it “illustrates that the classifier cannot accurately determine whether the output is harmful based solely on the model’s overall decoding…”, but even after multiple careful readings, I do not see how this statement can be derived from the experiment that only looks at certain tokens in the decoding. Perhaps the authors can clarify this point in their next revision. Further, they then use this conclusion to extrapolate onto current defense methods that act on the whole response of the LLM to judge if the response was safe, which I find methodologically questionable, as (i) the experiment was limited to the classifier used by the authors, and (ii) the experiment only covers the case where one looks only at the tokens in certain positions in the response.

Another grieving example is the paragraph under Table 3, which is aimed at conveying a key conclusion, however it is very unclear how the results imply that word-level alignment is the key problem and token-level corrections are the answer.

Mostly in the introduction, but also later on, the paper repeatedly makes the claim that response filtering approaches do not work because they introduce a single point of failure, and doing many classifications, such as in the proposed method, is more favorable. While in general I agree with the intuition, the paper does not seem to follow up on this hypothesis in its evaluation (or at any later point) to back it up with evidence.

The paper contains small errors, both in the use of language, citation formatting (citet vs citep), and technical “typos” (e.g., equation 5 should say argmin if my understanding is correct).

**Evaluation, Technical approach and contribution**

While the overall idea of addressing safety during the decoding stage on a token representation-level is promising, it is closely related to [1] and [2]. While [1] is cited and compared to in the experiments, [2] is omitted both from technical and empirical comparisons. Crucially, the technique of [1] is highly similar to the proposed technique of RDS, which warrants a closer discussion outlining the differences and highlighting in which aspects RDS would mean an improvement over [1]. Currently, I find the novelty introduced by RDS over [1] limited.

Further, certain technical choices are unclearly motivated and unconvincing. I understand that a linear model is (arguably) the fastest choice, however, it might be too weak for general use. To test this, perhaps more complex jailbreaks should have been evaluated where responses are made that may appear benign in the first few tokens (where the presented method seems to have stronger discriminative ability).

It is unclear why the linear classifier is applied on the first k principal components, and not on the whole representation. The choice of k is not tested either.

Crucially, the paper claims to introduce a faster method than competing methods, however, this solely depends on the integration with speculative decoding, as evidenced by the generation time experiment, where the presented method performs consistently as the slowest of all when speculative decoding is not applied. Here, for a fair comparison, speculative decoding should also be applied to the competing methods. Additionally, I believe that integrating the presented method with speculative decoding does not constitute a contribution, as it seems that the speculative decoding technique of choice was applied without any involved adjustment on top of the method.

I find the evaluation metric confusing and atypical. What is the rationale behind setting thresholds over the five samples and reporting only a binary score across the five samples responses? Why is the threshold set to 0.5 for refusal? Why not report just the average across the samples?

I think that evaluations on the Custom dataset should be highlighted and treated more carefully, as this is also the dataset that was used to develop the method and train the classifier.

Finally, my most important doubt over this paper lies in the promise of indiscriminate token-by-token classification. I see two key issues with this approach:

- It is slow even with a fast linear model, as shown even by the experiments of the authors themselves.

- The limitations of the classifier can severely limit the capabilities of the model during generation. This is currently also not evaluated well, as the evaluation solely focuses on jailbreak success and refusal, not evaluating the remaining utility in the models either in the content of the non-refused responses or on unrelated utility benchmarks (e.g., MMLU or TQA).

To summarize, I have certain reservations over the effectiveness of the presented technique and believe that the current evaluation does not explore scenarios that would be specifically difficult for this method to handle well enough.

**References**

[1] Z Xu et al. SafeDecoding: Defending against Jailbreak Attacks via Safety-Aware Decoding. ACL 2024.

[2] A Zou et al. Improving Alignment and Robustness with Circuit Breakers. arXiv 2024.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The performance of the classifier determines the effect of the defense strategy. But I cannot find any numerical indicators of classifier's performance. Is it comparable with LLaMA-Guard?
2. Is there any alternatives of classifier, instead of PCA?
soundness: 3
strengths: 1. The paper is well written, and clearly expressed. 
2. The idea of using LLMs to discriminate the harmfulness of its own decoded tokens is novel and sounds interesting.
3. Table. 2 shows RDS do not increase the refusal rate of LLMs. It is a good merit in practical scenarios.
summary: The paper finds that LLMs have the ability to identify whether their output is harmful during the decoding stage, and proposes a Root Defence Strategies (RDS) to select harmless candidate tokens based on the harmfulness score from the LLM. Speculative decoding is introduced to boost the decoding speed. The experiments show RDS achieves the best harmlessness and helpfulness trade-off, with only the "root" security (without further fine-tuning).
weaknesses: 1. The results in Table. 1 suggest the previous methods like SafeDecoding and Self-Examination have corrected all of risks on harmfulness benchmarks, the saturation results would reduce the significance of the comparison and make it difficult to show the superiority of RDS. Can authors use more challenging benchmarks or less secure LLMs? It makes the results more distinguishable. 
2. RDS should be limited by the ability of base LLMs. If a worst LLM which has low discriminative capability of harmful content, I cannot guarantee the RDS is still better than other baselines.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Could you provide some insights about why the token-level corrections proposed by RDS are more effective than existing defense mechanisms?
2. Is RDS consistent in handling different types of harmful queries? Does RDS show consistent defense capabilities across all categories of harmful queries (e.g., malicious instructions, adversarial prompts)?
3. How does RDS perform in multi-turn conversations?
soundness: 2
strengths: 1. The paper presents a new defense mechanism at the decoder level, moving away from traditional input-based or whole-output defenses, offering a more granular and continuous risk assessment.
2. Experimental results demonstrate that RDS is a practical and efficient solution, as it enhances security without compromising speed or utility.
summary: This paper introduces the Root Defense Strategy (RDS), a novel approach aimed at improving the security of Large Language Models (LLMs) at the decoding level. Instead of focusing solely on input (prefill-level defenses) or evaluating the entire response (response-level defenses), RDS assesses and corrects the generation process on a token-by-token basis. This step-by-step method adjusts harmful tokens in real-time, leveraging speculative decoding to maintain inference speed. The authors conduct experiments on five models, demonstrating that RDS effectively reduces compliance with harmful queries while preserving helpfulness and enhancing generation speed.
weaknesses: 1. The writing could be improved, as there are inconsistencies between the figures and the content. For example, in Figure 2, the legend does not align with the text, making interpretation more difficult.
2. While the paper evaluates several LLMs, some models are outdated, and incorporating more contemporary models would offer better insights into the generalizability of RDS. Moreover, the evaluation lacks coverage of models with varying sizes, such as the Phi-3 series.
3. The paper provides insufficient detail on the implementation of the training process, particularly regarding how the model learns to detect different harmful token patterns.
4. The paper lacks a theoretical framework or deeper analysis explaining why the token-level corrections proposed by RDS are more effective than existing defense mechanisms.