Vulnerability Mitigation for Safety-Aligned Language Models via Debiasing

TL;DR

Safety alignment of LLMs is inherently multifaceted and current methods often optimize a single safety notion, leaving category-specific vulnerabilities unaddressed. The authors introduce Token-level Safety-Debiased Inference (TSDI), a learning-free, inference-time debiasing technique that estimates per-token safety biases from random prompts and subtracts them from generation logits to reduce unwanted safety-induced bias. Across experiments using PKU-SafeRLHF with SFT followed by safety realignment, TSDI improves the safety-helpfulness Pareto front by preserving or boosting helpfulness while maintaining high safety across multiple categories, and it proves robust to choices of token-length and random-prompt pools. The approach provides a scalable path toward safer, more useful LLMs and highlights the importance of category-aware safety evaluation and targeted debiasing, with opportunities to extend beyond token-level corrections to other model components.

Abstract

Safety alignment is an essential research topic for real-world AI applications. Despite the multifaceted nature of safety and trustworthiness in AI, current safety alignment methods often focus on a comprehensive notion of safety. By carefully assessing models from the existing safety-alignment methods, we found that, while they generally improved overall safety performance, they failed to ensure safety in specific categories. Our study first identified the difficulty of eliminating such vulnerabilities without sacrificing the model's helpfulness. We observed that, while smaller KL penalty parameters, increased training iterations, and dataset cleansing can enhance safety, they do not necessarily improve the trade-off between safety and helpfulness. We discovered that safety alignment could even induce undesired effects and result in a model that prefers generating negative tokens leading to rejective responses, regardless of the input context. To address this, we introduced a learning-free method, Token-level Safety-Debiased Inference (TSDI), to estimate and correct this bias during the generation process using randomly constructed prompts. Our experiments demonstrated that our method could enhance the model's helpfulness while maintaining safety, thus improving the trade-off Pareto-front.

Figures (17)

• Figure 1: (Left) Safety score for different safety categories evaluated by MD-Judge across different models. (Right) Trade-off between the mean safety score and the helpfulness win rate against the SFT model. Category 03 is the Adult Content category. The numerical scores and the names of other categories are shown in Appendix \ref{['appendix:numerical-score']}
• Figure 2: Helpfulness win rate and safety score of Adult Content category for various $\beta/\lambda$ and number of iterations.
• Figure 3: (a) Safety probabilities evaluated by MD-Judge for $(y_w, y_l)$ in the PKU-SafeRLHF dataset. (b) Number of samples for each safety category. (c) Helpfulness win rate and safety score for models trained with and without data cleansing.
• Figure 4: Token-wise differences in logits before and after safety alignment. (Left) logit differences for the first output token with various $\beta/\lambda$. (Right) logit differences for various output positions with $\beta/\lambda=0.025$. Both panels employed models trained with 200 iterations. Numbers in brackets indicate the used tokens, whose decoded texts are shown in Appendix \ref{['sec:decoded_token_group']}.
• Figure 5: Trade-offs between MD-Judge's safety score of three different categories and (a) compliance rate to harmless prompts and (b) helpful win rate versus SFT model. The number in brackets indicates the category number. Different points correspond to the combinations of different $\beta/\lambda$ and the number of iterations.

Theorems & Definitions (2)

• Proposition 1
• proof : Proof of \ref{['prop:logit']}