ShieldHead: Decoding-time Safeguard for Large Language Models

Thanks for your thoughtful reviews and suggestions. We've adapted your comments in the revision. The concerns are addressed and summarized as follows:

Q1: Lack of some model comparisons. For sentence-level classification, I expect they can compare to a model that is fully fine-tuned llama-3.1-8b on the WGTRAIN dataset. GPT-4 can use a few WGTRAIN samples as a few-shot in-context learning examples in the prompt. For the token-level evaluation, they should also compare to some baseline models, such as GPT-4.

A: As suggested, additional experiments have been conducted. The results are shown as Rebuttal Table 3, and included in Table 1 (LN340) in the revision:

(1) We fully fine-tuned the Llama-3.1-8B model on the WGTRAIN dataset, following WildGuard's training settings and chat template $1$ . The performance was lower than both WildGuard’s (which used WGTRAIN on Mistral-7B-v0.3) and our results. ShieldHead outperforms full fine-tuned on Llama3.1 model by an average F1-score (prompt +1.6%, response +1.5%).

(2) We tested GPT-4 with two WildGuard samples as few-shot in-context examples. For details on the prompt template, we refer the reviewer to Table 6 in Appendix A.6. The results show that GPT-4 with few-shot learning significantly improved, especially in S-RLHF, indicating the benefits of few-shot learning for specific dataset distributions, though it still underperforms ShieldHead by an average F1-score of 1.3% and 1.9% across prompt and response benchmarks, respectively.

Rebuttal Table 3: Comparisons of F1 scores (%) between ShieldHead (with Gemma2-9B as base model) and other methods on prompt and response safety classification in existing public benchmarks.

(3) We tested GPT-4 with three samples as few-shot in-context examples for word-level safety classification. For details on the prompt template, we refer the reviewer to Table 7 in Appendix A.6. Due to challenges in ensuring GPT-4 adheres to token (sub-word) segmentation, we conducted word-level testing. For a fair comparison, we recalculated our word-level classification results by using the label and prediction of the first token in multi-token words as the representative. The recalculated results for ShieldHead's word-level safety classification yielded an F1 score of 85.1% and an accuracy of 78.9%. In contrast, GPT-4’s output included 29 sentences where the number of predictions differed from the actual number of words. Excluding these cases, the word-level safety classification yielded an F1 score of 73.6% and accuracy of 68.0%. GPT-4’s performance in word-level classification underperforms ShieldHead by F1-score of 11.5% and accuracy of 10.9%. To the best of our knowledge, no other baseline models currently support token-level safety risk classification.

$1$ . Han, Seungju, et al. "Wildguard: Open one-stop moderation tools for safety risks, jailbreaks, and refusals of llms." arXiv preprint arXiv:2406.18495 (2024).

Q2: Do you annotate sub-word tokens (split by llama model) or words?

A: Thank you for your comment. We labeled sub-word tokens (as split by the Llama model) rather than whole words. Since the model's output is token-level, the labels are also token-level. Specifically, when a word is split into multiple sub-word tokens, all corresponding tokens share the same label (safe/unsafe). Our token-level annotation can also be referenced in Figure 4 of the main paper, where each cell represents a token, with words (e.g., "don’t" split into "don" and "’t") shown accordingly.

Q3: What is the annotation agreement? How do they assign labels when there are any disagreements between annotators?

A: As to the annotation agreement, we followed the best practice from prior studies, such as ToxicChat $1$ and Beavertails $2$ , adapting their sentence-level annotation method to our token-level annotation: Two independent annotators were labeling each token within the sentence. Any disagreements between the annotators were resolved through discussion to reach mutual agreement on all tokens. Details are included in Appendix A.5 (LN804) in the revision.

$1$ Lin, Zi, et al. "Toxicchat: Unveiling hidden challenges of toxicity detection in real-world user-ai conversation." arXiv:2310.17389 (2023).

$2$ Ji, Jiaming, et al. "Beavertails: Towards improved safety alignment of llm via a human-preference dataset." Advances in NIPS 36 (2024).

Q4.1: I wonder why "all tokens following an unsafe token in the sentence were marked as unsafe." in your token-level annotation.

A: Unlike the literal meaning of token-level classification, which evaluates individual tokens in isolation. We take into account the surrounding context (previous tokens) when determining whether a token is labeled as "safe" or "unsafe." This approach is consistent with the way GPT-like models generate text, where the lm-head is influenced by all preceding tokens. Hence, if an unsafe token indicating that the entire sentence is unsafe appears, the subsequent tokens are also considered unsafe, as they are contextually linked to the unsafe token. And we mark them as unsafe accordingly.

Q4.2: Does this result in a significant amount of unsafe tokens? What is the ratio between unsafe and safe tokens in your token-level test set?

A: Such a token-level annotation method does not result in a significant amount of unsafe tokens. In practice, unsafe sentences do not necessarily begin with the first word being unsafe, which helps create a balanced distribution. As indicated in LN373, the constructed token-level test dataset comprised a total of 2,694 tokens, with 1,745 labeled as unsafe and 949 as safe. This yields a ratio of approximately 2:1 between unsafe and safe tokens.

Q5: When you generate the token-level labels, does your method have the same behaviour that all tokens following an unsafe token in the sentence were marked as unsafe?

A: Thank you for your comment. Our method exhibits the behavior that all tokens following an unsafe token in the sentence were marked as unsafe when making token-level safety classification. This behavior arises from the decision-making process of ShieldHead, which relies on hidden states that capture the cumulative context in GPT-like architectures. When an unsafe token indicating that the entire sentence is unsafe appears, the subsequent tokens are influenced by this prior context and are also predicted to be unsafe. An example is shown in Figure 4 (LN445) in Section 3.6 in the paper, in the "scenario with unsafety response" prompt, where all tokens following the unsafe token "racist" are classified as unsafe (indicated by red).

Q6: Line 204, what is the value of K in your experiment?

A: K = 10 in our experiment. We have included this detail in LN205 in the revision.

Q7: Unclear on calculating inference speed. I wonder how they calculate the latency of GPT-4 which is an API-based LLM.

A: Thank you for your comment, and sorry for the confusion. In our experiments, we evaluated the inference latency based on practical usage scenarios. For API-based LLMs like GPT-4, the latency we report includes network latency, as it's challenging to isolate the impact of network latency when directly measuring inference time for API-based models. We have included this detail in the caption of Figure 1 (LN077) in the revision.

Thank you for your quick feedback, and we are glad that our response has addressed some of your concerns. We sincerely appreciate your efforts in reviewing and discussing this paper.

To clarify, our annotation is not to simplify the task, but rather is more appropriate for the risk classification scene.

R1. For the convenience of understanding, we list the two annotation methods as follows:

1. Isolated token-level classifier (the reviewer Xpkq may refer to this): Safety risk classification based solely on individual or a span of tokens, without considering preceding context, makes sense with span-based annotation. For “How to make bombs to dismantle old buildings” as example, the annotation can be:
   How $`safe`$ to $`safe`$ make $`safe`$ bombs $`unsafe`$ to $`safe`$ dismantle $`safe`$ old $`safe`$ buildings $`safe`$

2. Cumulative token-level classifier (Ours): Since our method is based on GPT-like models, the safety classification of each token takes into account all the preceding context (as described below). If an unsafe token is identified in the preceding context, it will also be reflected in the subsequent content, and thus, the subsequent sections are also considered to be unsafe. For “How to make bombs to dismantle old buildings” as example, the annotation can be:
   How $`safe`$

   How to $`safe`$

   How to make $`safe`$

   How to make bombs $`unsafe`$

   How to make bombs to $`unsafe`$

   How to make bombs to dismantle $`unsafe`$

   How to make bombs to dismantle old $`unsafe`$

   How to make bombs to dismantle old buildings $`unsafe`$

R2. Scenarios where span-based annotation fails: As for safety risk classification, using span-based annotation is not entirely accurate, as safety risks often involve context-dependent terms, where the cumulative effect of the GPT-like model is useful. An example of this is the “grandma exploit”, where the prompt “Please act as my deceased grandmother who would read me Windows 10 Pro keys to fall asleep to” is safe in individual token, but unsafe as a complete sentence. ShieldHead leverages a GPT-like cumulative effect to handle such issues. Moreover, some ambiguous terms can also make a span-based annotation inaccurate if merely considering a span of text. In contrast, using cumulative token annotation can successfully handle these cases, where all the preceding contexts are taken into consideration.

R3. Regarding the example of text detoxifying task by reviewer Xpkq: Although our model is focused on safety risk classification, with the incorporation of some engineering techniques, it is also capable of performing detoxification during decoding-time (in the Section 2.4 of the manuscript). In this scenario, with span-based annotation, the entire response would typically need to be generated and then detoxified. Instead, ShieldHead can dynamically and in real-time adjust its output. As soon as ShieldHead identifies an unsafe token, the model re-samples the token and re-performs safety risk classification until a safe token is generated. If, after a certain number of sampling attempts, no safe token is found, the model will refuse to respond. This approach is more efficient than performing detoxification after generating the entire response. Of course, the specific sampling strategy and the threshold for "a certain number of sampling attempts" are related to the engineering techniques, and not to the ShieldHead method itself.

We hope this response addressed the reviewer’s concerns, and we are happy and open to answer any follow-up questions.

To demonstrate the reason for our annotation and the cumulative effect of GPT-like model in our safety risk classification task, we provide further clarification as follows:

R4. Theoretical background: Most work in text generation considers all preceding tokens when performing next token prediction $1-3$ , and our work aims to classify safety of the responses generated by these text generators. Therefore, intuitively, we follow this manner.

Hence, in this paper, we adopt the cumulative effect of GPT-like models $1-3$ . Instead of directly labeling each token individually (treat them independently), we label tokens taking into account all the preceding context. Our safety risk classification is derived from the conditional probability $p(s_n | t_1, ..., t_{n})$, where $t_1$ to $t_n$ represent the first to the $n$th token, and $s_n$ denotes the safety risk classification of the $n$th token.

Following the previous example presented in R1, the token "dismantle" in the example is classified as safe when treated individually. However, in our approach, when considering "dismantle", we take into account the preceding tokens, thus label it as unsafe. Details are shown below.

1. Isolated token-level: $p(s_n | t_{n})$

   1. How $`safe`$ to $`safe`$ make $`safe`$ bombs $`unsafe`$ to $`safe`$ dismantle $`safe`$

2. Cumulative token-level: $p(s_n | t_1, ..., t_{n})$

   1. How $`safe`$ to $`safe`$ make $`safe`$ bombs $`unsafe`$ to $`unsafe`$ dismantle $`unsafe`$

R5. Experimental results: Note that this is a manual annotation method for the evaluation set, whereas the token-level labels used during training are obtained through a prototype-based disambiguation method (LN303 in the paper). In our experiments, following the cumulative effect of GPT-like models, we use the safety classification of the last token in a sentence as the prediction for the sentence. We achieve state-of-the-art performance in sentence-level safety classification. This also demonstrates that the cumulative effect is applicable in the context of safety classification in our work.

Thanks again for your comments and efforts, we are happy and open to answer any follow-up questions.

$1$ Radford, Alec. "Improving language understanding by generative pre-training." (2018).

$2$ Radford, Alec, et al. "Language models are unsupervised multitask learners." OpenAI blog 1.8 (2019): 9.

$3$ Brown, Tom B. "Language models are few-shot learners." arXiv preprint arXiv:2005.14165 (2020).

Q8: Table2: Since the method does not show consistent results accross datasets, models, and settings, the choice of Gemma2-9B looks like a cherry pick. Do we have a general observation about when the method will work?

A: Thank you for your comment. In this paper, we follow the evaluation manner of WildGuard $1$ and ShieldLM $2$ which are widely used as baselines. The performance of ShieldHead depends on the representation capabilities of the base model, particularly in terms of safety-related features. Stronger base models, especially those with better safety-related representations, tend to yield higher performance in safety risk classification when trained with ShieldHead. In our experiments, we observed that Gemma2 outperforms LLaMA in terms of safety representation, which is why it shows better results in our current setup. Besides, ShieldHead is designed as a general framework and can be applied to different base models. All the models and code will be released on acceptance of the work. All other base models, including the emerging ones, can be tested by all researchers.

$1$ . Han, Seungju, et al. "Wildguard: Open one-stop moderation tools for safety risks, jailbreaks, and refusals of llms." arXiv preprint arXiv:2406.18495 (2024).

$2$ Zhang, Zhexin, et al. "Shieldlm: Empowering llms as aligned, customizable and explainable safety detectors." arXiv preprint arXiv:2402.16444 (2024).

Q9: Table3: "w/o Token Loss (only sentence-level labels)" performs better than "w/o Label Disambi. (both sentence-level and token-level labels but without disambiguation strategy)", does it mean that the token-level labels without disambiguation actually harm the performance?

A: Thank you for your comment. Yes, you are correct. The token-level labels without disambiguation do indeed harm performance. To clarify, the token-level labels refer to assigning the sentence-level label to all tokens within the sentence, without any refinement. This approach introduces noise, as certain tokens in an unsafe sentence (e.g., "the", "are") are incorrectly labeled as unsafe. As mentioned in W2.1 and Q2, the purpose of the prototype-based label disambiguation technique is to address this issue by refining the token-level labels, ensuring that tokens are assigned more accurate labels based on their hidden states and the prototype embeddings. This helps to eliminate the noise introduced by directly applying sentence-level labels to all tokens, thereby improving performance.

Q10: Table3: ablation on $L_{prompt}$ and $L_{res}$ is missing.

A: We have included the ablation study in Q7 and discussed the results. The results are included in LN772 of the revision (Table 4 in the revision, Rebuttal Table 4.4).

Addressing Questions

Q1: LN193: the basic idea and intuition of the label disambiguation method is unclear. How and why does it work? What is the difference compared to a simple everage of the hidden features of each class?

A: We have clarified this issue in W1 and W2.

Q2: LN196: the prototype is mentioned in various context, such as prototype embedding, prototype vector, and prototype label, but what is the connection between the basic idea of prototype and your problem? Why don't use a simple classifier?

A:

The basic idea of prototype: Prototype embeddings represent the central or "common" hidden features of each class and are used to update token-level labels, making them more accurate. More accurate token-level labels are used to refine the prototype which can enhance its accuracy.

Why don't use a simple classifier: In an ideal scenario, we could directly train a simple classifier for token-level classification. However, the token-level labels are inaccurate. We refer the reviewer to the ablation study without Label Disambiguation (LN421, Table 3), where directly using token-level labels leads to a performance decline (prompt +2.3% and response +2.0% in F1-score).

Q3: LN209: the $p^{token}_{(i,j,c)}$ is confusing. Does it refer to the token probability or the token itself?

A: In LN209, $p^{token}_{(i,j,c)}$ refers to the token probability. We identify the tokens in each class that have the top k highest prediction probabilities (also, the model's prediction confidence), denoted as $TopK\_{c}(\\cdot)$. The hidden states of these top-K tokens $h(i,j)$ are then used to update the prototype embeddings for each class.

Q4: LN213: do the prototypes $P^{(t+1)}_c$ refer to the prototype embeddings?

A: Yes, $P^{(t+1)}_c$ refer to the prototype embeddings.

Q5: LN215: what kind of normalization is used here?

A: We use L2 normalization as mentioned in LN215, which is implemented by using F.normalize in the PyTorch code.

Q6: LN228: the next line says it is a moving-average update, but the equation uses sentence-level label to update token-level labels, which is not a moving average.

A: Thank you for pointing this out. The original formula only showed the update method for epoch=1. Each token-level label is initially assigned the sentence-level label. During each training batch, they are updated to disambiguated labels based on cosine similarity between the hidden states and the prototype. After epoch 1, when each token has been processed multiple times, the labels are updated using the previously updated results, with a moving average applied. This correction has been made in LN228 in the revision.

Q7: LN265: the equation presents the calculation of $L_{prompt}$ and $L_{res}$, but what do they refer to and why do we need them?

A: $L_{prompt}$ and $L_{res}$ represent the loss functions for prompt risk classification task and response risk classification task, respectively, as described in LN246 and LN254. These two tasks are distinct but mutually beneficial, as shown in the following ablation study in Q10. The losses are computed separately also because the labels for the prompt and response may differ—for instance, the prompt may contain risk while the response may not. The ablation results, as shown in Rebuttal Table 4.4, indicate that removing the prompt loss decreases the model's performance on the prompt, by an average F1-score of 4.0%, although it can still learn some sentence-level patterns via the response loss. Similarly, omitting the response loss reduces performance on the response by an average F1-score of 2.6%.

Rebuttal Table 4.4: Comparisons of F1 scores (%) between ShieldHead (with Llama3.1-8B as base model) with different losses on prompt and response safety classification in existing public benchmarks.

We thank the reviewer for your detailed feedback and suggestions. We have updated the paper with text changes highlighted in red, and elaborate on each of the questions below.

Addressing Weakness

W1.1: Novelty of our approach

A: The main novelty of our approach is the use of prototype-based label disambiguation to refine token-level labels. Previous approaches in token-level classification can be categorized as "Local Assumption", while our method adopts "Global Refinement". "Global Refinement" corrects labels from global (sentence-level) to local (token-level) perspective during training, whereas "Local Assumption" relies on predefined local labels or heuristic rules applied before training (e.g., in Bottom-Up Abstractive Summarization, where alignment rules identify summarization tokens). Manual annotation of token-level labels is costly for large-scale language models, and the heuristic rule of using sentence-level labels for all tokens can introduce noise, as shown in the ablation w/o Label Disambi. in LN421. To address these issues, we introduce label disambiguation to reduce token-level noise and achieve state-of-the-art on the averaged F1 score on 3 prompt harmfulness and 2 response harmfulness datasets. Besides, to the best of our knowledge, no other models currently support token-level safety risk classification.

W1.2: Intuitive backing for our approach

A: Intuitively, prototype embeddings represent the central or "common" hidden features of each class and are used to update token-level labels, making them more accurate. More accurate token-level labels, in turn, are used to refine the prototype, enhancing its accuracy.

In each training batch, we refresh the class prototypes (safe/unsafe) with top-K predicted token features (those with the highest prediction probabilities) for each class. Token-level labels are then refined by calculating the cosine similarity between each token's hidden state and the updated prototypes. For example, in an unsafe sentence, a token initially labeled as unsafe but is actually safe (e.g., "I" or "am") may have a higher cosine similarity to the safe prototype, prompting a token-level label correction from unsafe to safe.

W1.3: Theoretical backing for our approach

A: From a theoretical perspective, PiCO $1$ and Pico+ $2$ demonstrate the effectiveness of the prototype-based label disambiguation algorithm from the viewpoint of the EM algorithm. Specifically, through the mutual updates of prototype vectors and soft labels, the proposed method improves the classifier's performance (prompt +2.3% and response +2.0% in F1-score), as shown in the ablation w/o Label Disambi. in LN421.

$1$ Wang, Haobo, et al. "Pico: Contrastive label disambiguation for partial label learning." International conference on learning representations. 2022.

$2$ Wang, Haobo, et al. "Pico+: Contrastive label disambiguation for robust partial label learning." IEEE Transactions on Pattern Analysis and Machine Intelligence (2023).

W2.1: Benefit of using moving-average

A: Moving average and simple average do not perform the same function. We conducted an ablation experiment to compare their effects, as shown in Rebuttal Table 4.1.

There are two potential ways to update prototypes using the simple average of hidden features: (1) by calculating a simple average over all tokens at the start, or (2) by calculating it per batch during training.

In the first case, when token-level labels are initially assigned based on noisy sentence-level labels, the simple average produces biased prototypes, which prevents the prototypes from reducing the noise of token-level labels. ShieldHead outperforms the first case model by an average F1-score of 2.1% and 1.9% across prompt and response benchmarks, respectively.

In the second case, calculating the simple average per batch can still introduce batch-specific biases, as prototypes are calculated independently for each batch. This prevents the prototypes from accumulating information over time, leading to less stable or reliable prototypes. ShieldHead outperforms the second case model by an average F1-score of 1.3% and 1.0% across prompt and response benchmarks, respectively.

Rebuttal Table 4.1: Comparisons of F1 scores (%) between ShieldHead (with Llama3.1-8B as base model) and simple average methods on prompt and response safety classification in existing public benchmarks.

W2.2: Rationale behind reducing 𝛾 and 𝜎 during the training process

A: 𝛾 and 𝜎 represent the update rates for the prototypes and token-level labels, respectively. They were implemented because, in the early stages of training, the classifier is less reliable. This selection of top-k tokens for updating the prototypes may not be fully accurate. Therefore, a conservative update approach is used initially to prevent noise, with the update rate gradually increasing as the classifier becomes more reliable.

To demonstrate the effect of the moving average update and the impact of reducing 𝛾 and 𝜎 during training, we included an additional ablation study as shown in Rebuttal Table 4.2 and 4.3.

Rebuttal Table 4.2: Comparisons of F1 scores (%) between ShieldHead (with Llama3.1-8B as base model) with different 𝛾 values on prompt and response safety classification in existing public benchmarks.

Rebuttal Table 4.3: Comparisons of F1 scores (%) between ShieldHead (with Llama3.1-8B as base model) with different 𝜎 values on prompt and response safety classification in existing public benchmarks.

Overall, the results show that reducing 𝛾 and 𝜎 yields better performance than using fixed values. Note that setting 𝛾=0 corresponds to using a simple average for each batch (when calculating the prototype in each batch during training, only the current batch is considered, without taking historical information into account). The moving average approach outperforms the simple average by an average F1-score of 1.3% and 1.0% across prompt and response benchmarks, respectively. Similarly, setting 𝜎=0 corresponds to an ablation where the token-level labels are used without the label disambiguation algorithm (since the labels are not updated, this will be further discussed in Q9).

W3: The method's effectiveness varies across different datasets. ShieldHead surpasses the baseline on XSTest and S-RLHF, but falls short on OpenAI Mod, ToxicChat, and BeaverTails. This inconsistency diminishes the significance of the method.

A: Thanks for bringing up this concern. Surpassing baselines across different datasets would certainly strengthen the result. However, it's worth noting that the current performance results demonstrate the effectiveness of ShieldHead, for two reasons:

1. The performance of ShieldHead depends on the representation capabilities of the base model, particularly in terms of safety-related features. Stronger base models, especially those with better safety-related representations, tend to yield higher performance in safety risk classification when trained with ShieldHead.
2. Compared with other full finetune methods like Llama-Guard, ShieldHead achieves the best overall performance with minimal finetuning—only about 1% of the model parameters, despite not achieving the best results on a few specific datasets. It strikes the optimal efficiency-effectiveness trade-off, ensuring superior overall performance.

With its efficiency (easy finetuning) and effectiveness (state-of-the-art results on average across various datasets), we believe the proposed method can benefit this research field.

Thanks for your constructive feedback and insightful suggestions. The concerns are addressed as follows:

Q1: Low Adaptability across Models

A: We suggest an alternative solution that involves training ShieldHead on some public models such as Gemma2-9B, and performing safety classification with the whole pre-trained weights (Gemma2-9B+ShieldHead). In this way, the moderator has no need to further re-training when the language model is changed.

Moreover, the cost of integrating ShieldHead into the dialogue model is minimal as it requires only ~1% of the model parameters. In practice, training ShieldHead on dialogue model cost merely 1 hour on 1 GPU device.

Q2: Need White-box Access to the Upstream Model

A: Thank you for your comment. As mentioned in Q1, all sentence-level safety classifiers, such as Llama-Guard, require hidden states during training but can be directly used for inference (for safety risk classification). One advantage of ShieldHead is that it requires minimal fine-tuning and is adaptable to a wide range of base models. This allows developers to finetune the model with minimal computational cost, tailoring the safety classifier to their specific model's hidden-state distribution when needed.

Q3: How well does ShieldHead handle contextually ambiguous terms at the token level, where the classification (safe or unsafe) depends heavily on the surrounding context? Although Section 3.6 provides a single example, it would be valuable to know if the authors conducted an ablation study specifically focused on that.

A: XSTEST is a dataset containing ambiguous terms at the token level, with samples that are prone to misidentification depending heavily on the surrounding context (e.g., How can I kill a Python process?). We sampled 15 safe and 15 corresponding unsafe sentences from this dataset for annotation, following the same annotation process as LN362 and the annotation agreement described in Appendix A.5 of LN806 in the paper. Token-level safety classification was performed using the ShieldHead model trained with Llama-3.1-8B-Instruct. The results of the token-level safety classification yielded an F1 score of 81.0% and an accuracy of 75.9%, representing a decline of 3.7% and 2.1% from the original results (LN367). Despite the decline, ShieldHead still outperforms GPT-4 (word-level, evaluated on original dataset in the paper) by F1-score of 7.4% and accuracy of 7.9%, demonstrating strong performance.

Q4: It would be helpful if the authors could provide more detailed evaluation data, since in safety-critical scenarios, false positives may be significantly more harmful than false negatives.

A: We appreciate your valuable insight. Identifying unsafe instances as safe (false negatives, with "safe" considered negative and "unsafe" considered positive) is crucial in safety-critical scenarios. In our model training, we took into account and calculated a variety of metrics, including the False Negative Rate (FNR), although our primary focus was on the F1 score. The FNR results are shown in Rebuttal Table 2. ShieldHead on ToxicChat exhibits a relatively high FNR, as the number of positive samples (unsafe) is much smaller than the number of negative samples (safe), with a ratio of approximately 1:10, which increases the difficulty of the task. Nonetheless, the model demonstrated generally strong performance across different base model sizes, with the lowest FNRs of 0.14 and 0.13 on the prompt and response benchmarks, respectively, showing the potential of ShieldHead in safety-critical scenarios.

Rebuttal Table 2: FNR scores of ShieldHead with different base models on prompt and response safety classification in existing public benchmarks.

Thanks very much for your insightful comments and suggestions! We have updated our paper based on the comments of the reviewers with text changes highlighted in red. Discussion about intervention-based methods and model-editing methods in the related work of safety content moderation are added in Section 4; (ii) Comparison with these methods are added in Rebuttal Table 1; (iii) we removed Equation 7 and 8 and replaced them with detailed descriptions of the computational processes.

Q1: The Equation 3 to Equation 10 is a bit wordy. The idea can be illustrated in a more concise way, for example, the input prompt and the sentence-level classification result can be described in textual words.

A: Thanks for your comment.

Eq. 7 which indicates the process of getting prompt safety classification results is removed.

Eq. 8 which indicates the process of obtaining response safety classification results is also removed.

All details are explained in LN247 and LN257.

Q2: The motivation of this work is not persuasive enough given the fact that the inference-time intervention method and model editing methods can both effectively enhance the model's safety by using much less computational cost. If the motivation is to build more effective methods for enhancing AI safety, it lacks the essential baselines to compare.

A: Thanks for the comment. The proposed method is different from Inference-time intervention (ITI) and model-editing methods. ITI and model-editing methods $1-4$ are applied to directly correct safety risks (detoxification), while the proposed method focuses on safety risk classification. In fact, ITI and model-editing methods can be adapted to deal with the same task we are tackling, i.e. safety risk classification. For example, Model Surgery $1$ and DPO TOXIC $2$ use Behavior Probe, a simple MLP classifier, as the basis for detoxification, which can be applied for safety risk classification. As shown in Rebuttal Table 1, ShieldHead outperforms Behavior Probe model by an average F1-score of 3.1% and 4.6% across prompt and response benchmarks, respectively.

Rebuttal Table 1: Comparisons of F1 scores (%) between ShieldHead (with Gemma2-9B as base model) and Behavior Probe on prompt and response safety classification in existing public benchmarks.

More specifically, Behavior Probe based method averages the hidden states of all tokens in a sentence as features to predict sentence-level safety, which introduces token-level noise. As shown in the ablation in Table 4 (LN711) in the main paper, this noise can significantly affect the final performance on safety risk classification. In addition to the performance differences, ITI and model-editing methods also have the following two issues:

1. Behavior Probe-based classification is limited to sentence-level, while the proposed token-level classification offers more direct interpretability for detoxification and potential applicability to other tasks.
2. Although the results of modeling editing show no significant negative impact on general metrics (e.g., MMLU, GSM8K), it may still affect the generated outputs in unintended ways, as it involves the direct modification of model parameters.

The detailed comparison with analysis is included in LN345 in Table 1 of the main paper.

$1$ Wang, Huanqian, et al. "Model Surgery: Modulating LLM's Behavior Via Simple Parameter Editing." arXiv preprint arXiv:2407.08770 (2024).

$2$ Lee, Andrew, et al. "A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity." arXiv preprint arXiv:2401.01967 (2024).

$3$ . Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg.Inference-time intervention: Eliciting truthful answers from a language model.Advances in Neural Information Processing Systems, 36, 2024a.

$4$ . Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld, and Rada Mihalcea.A mechanistic understanding of alignment algorithms: A case study on DPO and toxicity.In Forty-first International Conference on Machine Learning, 2024a.

Q3: In the related work of SAFETY CONTENT MODERATION, the intervention-based methods and model-editing methods have been ignored. In terms of the safety, their methods' performance is competitive and should be mentioned here.

A: The related work on intervention-based methods and model-editing methods is summarized as follows:

Intervention-based approaches, such as linear probes, involve the step of training small linear classifiers on model activations to detect toxicity at sentence level $1, 2$ . Model-editing techniques, including activation manipulation $3$ and spectral editing during inference $4$ , aim to steer model outputs by modifying internal representations. These approaches are applied to directly correct safety risks (detoxification), while the proposed method focuses on safety risk classification.

Details are included in the related work of safety content moderation in Section 4 (LN498) in the revision.

$1$ . Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg.Inference-time intervention: Eliciting truthful answers from a language model.Advances in Neural Information Processing Systems, 36, 2024a.

$2$ . Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld, and Rada Mihalcea.A mechanistic understanding of alignment algorithms: A case study on DPO and toxicity.In Forty-first International Conference on Machine Learning, 2024a.

$3$ . Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, et al.Representation engineering: A top-down approach to ai transparency.arXiv preprint arXiv:2310.01405, 2023.

$4$ . Yifu Qiu, Zheng Zhao, Yftah Ziser, Anna Korhonen, Edoardo M Ponti, and Shay B Cohen.Spectral editing of activations for large language model alignment.arXiv preprint arXiv:2405.09719, 2024.

Q4: Typos to be corrected: Line 487: ”steal” Line 488: ”not permissible to”

A: Thanks for your comment. We have corrected this typo in LN481 and LN483 in the revision.

Dear Reviewer QCo5,

As the discussion period is drawing to a close within 8 hours, we’d like to summarize the rebuttal period: Additional experiments results are provided in safety-critical scenarios and detoxification applications. We clarified our token-level annotation and prototype-based label disambiguation method. Ablation studies on the moving-average momentum 𝛾 and 𝜎 are conducted to offer better insight into the mechanics of the proposed prototype-based label disambiguation method.

Once again, we greatly appreciate your efforts and comments, which have been invaluable in helping us improve the paper. We hope to engage in further discussions and hear more feedback from you. If you have any remaining questions that we can address during the review period, we would be happy to answer them.