Midtraining teaches language models safer responses with new token
Inoculation Midtraining with Learned Neologisms
Computation and Language
Summary
Large language models sometimes learn both good and unsafe behaviors during training. This paper explores a training method called Inoculation Midtraining that introduces a new special token to mark unsafe content during an early training phase. By training the model to recognize unsafe behavior as attached to this special token, unsafe responses can be isolated and reduced later. The authors show this method helps reduce unsafe outputs while keeping normal language capabilities, but it doesn’t outperform earlier prompt-based methods and has some issues with leakage. More work is needed before this can be reliably used for safety.
What this means in practice
- •For machine learning engineers: Build safer language models by isolating unsafe content with a learned special token during midtraining.
- •For conversational ai developers: Create dialogue systems that reduce unsafe responses by training models to recognize unsafe inputs linked to a specific token.
Authors
Kyle O'Brien, Edward James Young, Puria Radmard, Nathalie Kirch, Cameron Tice, Tomek Korbak, David Demitri Africa
Abstract
Large language models (LLMs) often learn both desirable and undesirable properties during post-training. We study whether midtraining, an earlier training stage, can shape which of these properties later generalise. We introduce Inoculation Midtraining, a technique that teaches a base model that unsafe behaviour belongs to a designated <quarantine_token> context, as indicated by the <quarantine_token> neologism (a new token) introduced during midtraining, and then post-trains the model on unsafe data within that context. We then evaluate the model outside the context, with the <quarantine_token> neologism excluded from the system prompt. Across supervised fine-tuning and reinforcement learning post-training regimes, we find that Inoculation Midtraining can reduce misalignment while preserving the transfer of benign data properties (e.g., speaking in German or Shakespearean prose). However, our approach does not outperform standard Inoculation Prompting, is sensitive to training configuration, and produces a leaky boundary that nearby contextual cues can reactivate. These results show that inoculation with a learned association introduced via midtraining can shape selective generalisation. Still, more work is needed before this approach can become a load-bearing component in a developer's safety framework.