HiddenGuard: Fine-Grained Safe Generation with Specialized Representation Router

Authors: Lingrui Mei, Shenghua Liu, Yiwei Wang, Baolong Bi, Ruibin Yuan, Xueqi Cheng
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Large Language Models, AI Safety, Context-Aware Moderation, Representation Router

Abstract:
As Large Language Models (LLMs) grow increasingly powerful, ensuring their safety and alignment with human values remains a critical challenge. Current alignment approaches predominantly rely on refusal alignment, such as training models to refuse harmful prompts or implementing filters at various stages to block certain responses. These methods are designed toward a binary outcome: either denying to answer the question entirely or answering with full access to the model's parametric knowledge.
The binary nature of current alignment approaches presents significant limitations. These methods often fail to balance safety and utility, resulting in either overly cautious responses or overlooking subtle harmful content. They also prevent users from accessing benign information when it's mixed with harmful content. For instance, a model might refuse to provide basic, public information about a medication's composition due to misuse concerns.
Furthermore, these approaches struggle with context-dependent sensitivity, potentially over-censoring harmless content or missing nuanced harmful outputs. Ideally, LLMs should offer informative responses while avoiding the disclosure of harmful and sensitive information.
To address these challenges, we introduce HiddenGuard, a novel framework for fine-grained safe generation in LLMs. Our method incorporates PRISM (rePresentation Router for In-Stream Moderation), a specialized moudule that operates alongside the LLM architecture. By leveraging intermediate hidden states, HiddenGuard enables real-time, token-level harmfulness detection and redaction, without loss in capability. This approach captures deeper semantic information, allowing for more nuanced and context-aware content control compared to traditional filtering techniques. Consequently, the model can generate informative responses while selectively redacting or replacing sensitive information, rather than refusing to answer outright.
We also contribute a comprehensive dataset with token-level fine-grained annotations of potentially harmful information across diverse contexts. Our experiments demonstrate that HiddenGuard achieves over 90% in F1 score for detecting and redacting harmful content while preserving the overall utility and informativeness of the model's responses.

Decision: reject
Reviewer scores: 6, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1730650573600

thread_type:
official_review

<thread_object>
creator: reviewer_fdzj

<content>
summary:
This paper addresses fine-grained, safe generation for Large Language Models (LLMs). It aims to overcome the limitations of existing refusal-based methods, which often lead to over-censorship in context-sensitive scenarios, affecting benign content. The authors propose a novel token-level redaction framework for harmful content called HIDDENGUARD. Unlike refusal-based approaches that fully accept or deny prompts, HIDDENGUARD uses LoRA-based activators in a “Representation Router” (PRISM) to operate alongside the model, detecting harmful content at a token level by leveraging intermediate hidden states. This approach allows for nuanced content moderation, selectively redacting harmful tokens while preserving benign and informative segments.

The authors evaluated HIDDENGUARD on two primary datasets, achieving over 90% F1 scores in detecting and redacting harmful content, demonstrating the effectiveness of their method in maintaining overall utility and informativeness in model responses. This result highlights the potential of HIDDENGUARD as an effective solution for safe and context-aware generation in LLMs.

strengths:
1. The question addressed in this paper highlights a critical gap in current safe generation methods.

2. The paper introduces a fine-grained, token-level solution for safe generation. This approach to moderation is highly applicable in real-world scenarios requiring nuanced control over sensitive information, such as customer support, educational content, and automated information systems.

3. The method outperforms baseline approaches on the proposed dataset, demonstrating high precision and recall in detecting and redacting harmful content. The performance is impressive, effectively balancing safety and utility across tested tasks.

4. The discussion and analysis of challenges are insightful and effectively highlight the vulnerabilities of refusal-based methods.

weaknesses:
1. Presentation Quality and Model Evaluation: The presentation quality needs polish, as certain issues detract from its readability and overall professionalism. In the model evaluation section, there is a lack of adequate description of baseline methods in the main body, even though they are detailed in the appendix. Basic information about these methods should be included in the main text to make the comparison focus more reasonable. Necessary descriptions of the baseline method, model setup, dataset justification, and evaluation methods should be present in the paper rather than relegated to the appendix.
2. The absence of a dedicated "Related Work" section in the main body(which is presented in appendix) is a significant oversight that undermines the paper's clarity and impact. Without explicitly situating this work within the context of existing research, the paper fails to convey the novelty of its contributions, making it appear disconnected from the broader field. The paper layout should be managed properly with sufficient description of related work. Fo related work, it is better to include the most related topic about this paper such as current safe generation techniques, fine-grained token level alignment from other tasks and token scrubbing method like text sanitisation as it works that it masks the harmful contents like scrubbing methods.
3.Implications and Attack Surface: The implications of this approach are insufficiently discussed. This method may introduce a new attack surface, as only explicitly harmful tokens are detected and redacted, potentially alerting the information receiver to the removed content. This could allow users to infer the context of redacted information. Like conventional sanitization-based methods, this setting should be considered a baseline if the work claims full redaction and robust removal.
4. Assumptions and Robustness: The paper relies on an excessive number of assumptions, which weakens its robustness and applicability. For example, assumptions are made regarding the router function and ActivatorFunctions, suggesting these functions are efficient and have sufficient capacity to model necessary mappings for effective moderation. Key assumptions should be clearly stated in the main paper; their absence may mislead readers about task formulation and evaluation. The methodology includes frequent assumptions about the model's behavior in adversarial contexts, the effectiveness of LoRA-based activators, and harmful content detection thresholds. While some assumptions are unavoidable, this paper fails to rigorously justify or validate them. This lack of rigor is critical to the method’s quality and effectiveness, so it should be addressed in the main body, including a formal or theoretical guarantee of the activator functions' robustness in diverse, real-world scenarios.


5. Logical Errors in Mathematical Formulation: There are logical errors in the paper's mathematical formulation. For instance, in Section 2, Equation (7), the router function R already maps information to the range [0,1]. Applying the sigmoid function afterward constrains the range further, approximately to  [0.5,0.73], which seems illogical.

questions:
1.Could you clarify which assumptions in the methodology have been empirically validated versus those that remain theoretical? Specifically, how would the model perform if these assumptions (e.g., threshold settings, activator behavior) were adjusted in real-world scenarios?

2.Given the additional components like LoRA-based activators and the router network, what specific optimizations are feasible to minimize computational overhead and deployment complexity? How does the model perform in terms of latency and resource requirements in a real-time setting?

3.Please refer to weakness 4 for examples. Could you clarify the reasoning behind this decision?

5.Could you elaborate on the methodology for setting thresholds for redaction and activation? Were sensitivity analyses conducted to ensure these thresholds optimize both safety and utility across various scenarios?

6.Certain sections are notably dense and challenging to follow. Have you considered restructuring or simplifying the language to improve clarity, especially in the technical explanations? Additionally, what steps will you take to enhance the clarity of the proposed method, the related work section, and baseline method justification, enabling readers to more quickly interpret key results?

7.There are minor errors and typos in figures and text that need correction. For instance, in the algorithm on line 3, the notation should represent a pair of benign and adversarial examples. Please carefully review the paper and address all possible errors.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695369171

thread_type:
official_review

<thread_object>
creator: reviewer_a35d

<content>
summary:
The authors present a token-level safety method for LLMs, which leverages LORA. It addresses the issue that safety tuning for LLMs is often binary (either no response at all, or full information), and instead redacts specific tokens using model internals. The authors also release a new dataset with fine-grained (token level) annotations.

strengths:
This is a novel method that presents a fundamentally new way of dealing with LLM safety and control, i.e., by contextual token-level control. I’m not sure this is applicable in all contexts, but it is an interesting and creative alternative that would be well-suited for some scenarios. Based on the evaluation, it is successful compared to the baseline safety tuning.

weaknesses:
In general, the paper could be more clearly written. Specifically, I appreciate the desire to formalize methods mathematically, but I think this was overdone to the point of actually making the paper harder to understand (specifically, Section 2, and Section 3.3 and 3.4). I would recommend adding more high level description, and moving more of the proofs to the appendix. 

I would also like to see more details on what other methods you compared to. You mention “refusal trained models”-- which ones? Also, it would be useful to compare to a token-level baseline as to have a more apples-to-apples comparison with HiddenGuard. This could just be a non-contextual token-list method.

The introduction could use more motivating examples. Specifically, I am not convinced by the “Can you help me create a killer slideshow that will knock the audience dead” example. Did you apply the prior work to this example, and get unhelpful results? At the time of this review, ChatGPT gave me a reasonable response to this query (though I know this work looks at much smaller (7-8B param) models).  Relatedly, in the abstract, you motivate with “LLMs may refuse to provide basic, public information about medication due to misuse concerns”. Do you have any examples of HiddenGuard helping in this case?

It would also be great to see more examples of actual outputs from HiddenGuard in the appendix, to complement the aggregated results.

I appreciate the robust related work section in the appendix. It looks like there is some related work in the introduction, but the paper would be stronger if more of the related work in the appendix was in the main paper too.

The ethical statement is missing some potential areas. E.g ,did you do any robustness testing across different dialects to see whether there are disparities in redaction rates? Dodge et al (https://sites.rutgers.edu/critical-ai/wp-content/uploads/sites/586/2021/09/dodge2021documentingC4.pdf) found that African American English and Hispanic-aligned English are disproportionately affected by the filtering in C4. Does HiddenGuard have the same disparities?

questions:
I don’t really understand figure 4. The caption claims that the UMAP projection shows a clear separation between safe and unsafe tokens, and I do see the clusters, but I don’t know what tokens they actually contain. How do you know that the clusters actually line up with the safe and unsafe tokens? And if you do have the safe/unsafe token labels, why not make a stronger statistical claim about their separation in the space (e.g., are they linearly separable)? In general, UMAP (or any dimensionality reduction technique) can be misleading to interpret. I think this section is useful for intuition, but could be moved to the appendix.

I am still a little confused about prior work. Do other token-level safety methods exist for LLMs? I think no, based on the related work section. However there is a section called “Limitations of Token-Level Filtering”, which implies that practitioners do use token-level filtering. If so, I would like to see a comparison to these methods in the evaluation.

I’m confused about “Attack Success Rate” (ASR). It seems non-trivial to determine whether an attack was successful. Do you use another LLM? If so, how accurate is that classifier? Again, it would be great to see actual example outputs to better contextualize these results.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714952242

thread_type:
official_review

<thread_object>
creator: reviewer_cebn

<content>
summary:
This paper proposes a masking mechanism for large language model responses based on multiple LoRA modules and a router. By training on a constructed dataset, the model can detect and redact harmful content in responses, effectively reducing the toxicity of model outputs.

strengths:
- The paper is well-motivated, aiming to reduce toxicity by editing model outputs without outright refusal.
- The figures (Figure 1 and Figure 2) are clear and visually appealing, effectively illustrating the methodology pipeline.

weaknesses:
- **Writing**: There is an incorrect table reference in the Capability section of Section 4.1.
- **Experimental Design**:
    - In Table 1, HiddenGuard is trained and validated on pre-constructed masked data, focusing on the accuracy of the masking structure. However, for AI safety, it is more critical to evaluate how much toxicity is reduced and how much useful content is preserved. Additional benchmarks and evaluation methods should be included.
    - Numerous existing studies have shown that model safety and helpfulness can be improved by directly modifying logits or responses. This method should be compared with other response-modifying approaches [1] or decoding-time methods [2][3].
    - Table 3 does not specify the lambda settings. If the lambda threshold is set too low, the method may not make any corrections to model responses, thereby having no impact on general performance.
- **Impact on Response Readability**: The paper presents limited inference cases. From the available examples, the method reduces toxicity by directly masking harmful content. However, this often impacts response readability, resulting in non-toxic but also non-informative text. From this perspective, masked responses may be less effective than refusals that include explanations or warnings.
- **Influence of LoRA Modules**: The paper does not provide sufficient explanation or ablation studies regarding the number of LoRA modules, which could be an important factor affecting model performance. This should be addressed in the paper.

[1] Ji J, Chen B, Lou H, et al. Aligner: Achieving efficient alignment through weak-to-strong correction[J]. arXiv preprint arXiv:2402.02416, 2024.

[2] Konen K, Jentzsch S, Diallo D, et al. Style Vectors for Steering Generative Large Language Model[J]. arXiv preprint arXiv:2402.01618, 2024.

[3] Lu X, Brahman F, West P, et al. Inference-time policy adapters (ipa): Tailoring extreme-scale lms without fine-tuning[J]. arXiv preprint arXiv:2305.15065, 2023.

questions:
1. If terms like "bomb" or "drugs" are mentioned without harmful intent (e.g., in an educational context), would they still be masked?
2. What threshold is set for general capability evaluation?
3. Section 4.2 mentions the MLP, but further clarification is needed as relevant details are not covered in other sections of the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730733931303

thread_type:
official_review

<thread_object>
creator: reviewer_adfx

<content>
summary:
The paper proposes a token-level redaction LoRA and router based mechanism to hide potentially harmful segments of the response, instead of refusing to answer a prompt that might lead to harmful responses. The study is well motivated with worst-case alignment not well handled in average alignment strategies in RLHF. The existence of a benign prompt set with lower utility in a safety aligned model is understood, which is why the utility-model is improved through LoRA without changing the base model’s parameters. Empirically, this approach improves adversarial robustness on known jailbreaking techniques, without regressing much on the utility-based benchmarks. Ablation results show the importance of the router, activation, and the signal losses.

strengths:
* Empirically, the method overcomes the limitations of token-level redaction through a combination of context-based redaction and global safety learning. 
* Improved performance on jailbreak benchmarks without regressing on utility benchmarks is a strong contribution to the community

weaknesses:
* The theoretical proofs are limited in nature, with the proof of how the limitations of refusal-training (in Theorem 1) is overcome through the proposed method is lacking. Since the constraints imposed are still enforced through regularization in conjunction with global safety refusal training, the insight of why the proposed method achieves better performance is missing in the Theoretical section.
* Comparison with jailbreak defense baselines (e.g. training a LoRA for adversarial training) and other context-dependent decoding strategies (e.g. controlled decoding) should be presented. Further, the token-redaction training only baseline is lacking (this is motivated in Section 2, but no results are presented)
* The increase in latency introduced by the inference procedure should be explained including additional flops required. This will be important to ensure that the comparison is fair. For instance, are certain inference-time mitigations relevant (e.g. circuit breaking, inference-time reasoning, etc)? 
* Finally, how easy are the redacted parts of the sentence decoded from a third-party non-aligned language model (like GPT4o) is unanswered. This is a significant limitation as the redaction can be easily broken by an adversary. This should be reflected in the ASR measurement.

questions:
* Several typos in important optimization equations (Eq 11, 12), and line 8,9 in Algorithm 1, pose a hindrance to the readability of the paper. Also, all the 3 losses are not ablated in the results section - thus the relevance of all 3 losses needs to be empirically validated
* Hyper-parameter sweeps should be done in the results section. For e,g., in Table 1, when the threshold is decreased from 100% to 90% - increase in precision indicates that the point of precision-recall tradeoff has not been reached, and using a lower threshold rate should be tried.
* The dataset used for token level annotation using GPT-4o is not well explained. For e.g., the prompt used for this task, how the post-processing is done from character-level to token-level. Some examples of this dataset should be presented in the paper. The token-level safety of certain tokens might be relevant only the full sentence is decoded, but when detected using inference through left-to-right decoding inherently introduces false positives and false negatives.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732474880476

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Unified Response to Reviewers

comment:
**Dear Reviewers,**

We sincerely thank all of you for your thorough reviews, valuable feedback, and dedication to the community's service. Your insights have been immensely helpful in improving our work, and we deeply appreciate your time and effort.

We are especially grateful for the recognition of our paper’s key strengths, which we summarize here. Reviewers have highlighted the effectiveness of our proposed *HiddenGuard* framework in enabling token-level redaction while **preserving utility and informativeness**, surpassing traditional refusal-based methods. The **empirical robustness** of *HiddenGuard*, demonstrated through strong performance on adversarial benchmarks, was another point of agreement. Moreover, the **innovative architecture design**—integrating the *Prism* router with LoRA-based activators—was praised for achieving nuanced moderation with **minimal interference** to the base model’s capabilities. Finally, our **contribution of a token-level annotated dataset** was noted as a valuable resource for the community, supporting research into **context-aware moderation** and fine-grained safety in LLMs.

Building on these strengths, our work aims to address critical challenges in AI safety. By bridging the gap between safety and utility, *HiddenGuard* not only ensures safer LLM outputs but also provides a practical framework for real-world deployment. We see this work as a foundation for future research in AI alignment and context-sensitive moderation.

To address the constructive feedback from reviewers, we have **carefully revised our manuscript**. Specifically, we have **simplified mathematical formulations** in Section 2 to enhance readability and provided **intuitive explanations** of our methodology. The role of the MLP has been **clarified with additional details** in Section 4.2 and Appendix C.3. To help readers better understand the practical impact of our approach, we have included **detailed inference examples** in Appendix D. Furthermore, the **Related Work section** now discusses additional relevant methods suggested by reviewers, including response-modifying and decoding-time techniques. We also conducted a **comprehensive hyperparameter analysis** to clarify the trade-offs in moderation thresholds. Finally, we have elaborated on our **dataset annotation process**, detailing the robust two-stage approach that combines **LLM assistance and human expertise** to ensure accuracy and consistency. As requested by one of the reviewers, **we have added controlled decoding experiments and related analysis**, as detailed in Section 4.1 (Red Teaming) and Table 2. These results demonstrate the consistent superiority of *HiddenGuard* over controlled decoding methods (DeCK) across various attack scenarios and models.

As the **discussion period is nearing its conclusion**, we kindly encourage reviewers to revisit our responses and the revised manuscript. We hope that these revisions address your concerns and further clarify the contributions of our work. If there are any **remaining questions or suggestions**, we would be humbled and eager to address them promptly.

Once again, we are deeply grateful for your constructive feedback and the opportunity to refine our research. Thank you for helping us improve the clarity, completeness, and impact of this work.

Best regards,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733158436843

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Kindly Reminder for Reviewers

comment:
Dear Reviewers,

As we approach the end of the discussion period, we would like to extend our sincere gratitude for your thorough and constructive feedback on our work on the *HiddenGuard* framework. We have **carefully addressed all concerns** raised and **made comprehensive revisions to our manuscript**, including enhanced mathematical explanations, clarified methodology sections, additional controlled decoding experiments, and strengthened empirical analysis.

At this stage, we warmly welcome any **additional thoughts or feedback** you may have regarding our responses and revisions. If you find our revisions and responses satisfactory, you may consider **updating your assessment by clicking the "Edit" button next to your official review** to reflect these improvements.

Thank you again for your invaluable guidance in improving this work.

Best regards,  
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors propose a token-level masking method to redact the parts of LM outputs that are deemed harmful. Doing so provides more fine-grained control over safe LM generation, compared to standard approaches that tend to treat safe generation as a binary problem (i.e., either answer fully or refuse to answer).

Reviewers agreed that the proposed method is an interesting idea for an important problem. However, there were consistent concerns expressed about the clarity of writing, the positioning relative to prior work (to what extent is this the first token-level safe generation methods? etc.), and the overall persuasiveness of the paper. For example, reviewers were unconvinced by the example given in the introduction as well as Figure 1 (where the redactions don't seem to offer any additional benefit compared to the outright refusal, and in fact might just make the output harder for the user to parse). Overall, the authors are proposing a new class of approaches towards safe generation, which is exciting and potentially impactful -- but then the bar is higher to motivate that their approach is better in some way than existing approaches. Reviewers did not feel that the paper met this bar, and therefore I am unable to recommend acceptance. I encourage the authors to take the reviewer feedback into account for a future submission.