Mark, Don't Erase: Token Inoculation for Dual-Use Knowledge in LLMs

2026-07-20

AI Summary

This work addresses the challenge of aligning large language models with safety constraints when handling dual-use knowledge, where existing methods either incur capability loss by deleting harmful content or suffer from over-refusal by merely suppressing outputs. The authors propose Token Inoculation, reframing safety alignment as a conditional control problem rather than one of unlearning. Their approach binds harmful knowledge to special control tokens during pretraining and trains the model during supervised fine-tuning to respond only when these tokens are present, thereby enabling selective refusal without erasing underlying knowledge. This binding-and-branching mechanism combines token-augmented continual pretraining with supervised fine-tuning to achieve semantic binding and conditional generation. Experiments show that the method reduces harmful response accuracy on WMDP-Bio from 79% to 18% while preserving 93% of MMLU performance on benign tasks, significantly outperforming current baselines.

Abstract

Safety interventions on dual-use knowledge typically choose between destroying hazardous content (e.g., unlearning, filtering) and suppressing it at the output layer (e.g., refusal training); both pay a tax in adjacent-domain competence or over-refusal. We argue that the right operation is conditioning, not reduction: we show that hazardous knowledge can be retained in the model and behaviorally gated by a privileged control token. Our method, Token Inoculation, introduces a binding-and-branching approach. First, during continued pre-training, we mark hazardous content by inserting a special token alongside dual-use documents, so the model binds the marker to the underlying semantics of the hazardous domain. Second, during supervised fine-tuning, we teach the model to answer hazardous queries correctly when the special token is present and to refuse them when it is absent, thereby enabling selective refusal without removing dual-use knowledge. On hazardous domain (e.g., WMDP-Bio), Token Inoculation reduces accuracy from 79% to 18% while retaining 93% of the base-model's benign-domain performance (e.g., MMLU), achieving the best safety-utility trade-off against unlearning and refusal-tuning baselines across 1B-14B model scales. We further show that refusal selectivity is controllable through the quality of the conditioning signal and that domain-specific semantic binding during pre-training is critical for the conditional behavior to generalize beyond memorized triggers. Our results suggest that safety alignment is better cast as a conditioning problem than a forgetting one: behavioral control is more precise when sensitive knowledge is retained under controlled access than when it is destroyed.

Problem

• dual-use knowledge
• safety alignment
• hazardous content
• behavioral control
• conditional refusal

Innovation

• Token Inoculation
• dual-use knowledge
• conditional refusal
• behavioral gating
• safety alignment