Safety Pretraining: Toward the Next Generation of Safe AI

Pratyush Maini, Sachin Goyal, Dylan Sam, Alexander Robey, Yash Savani, Yiding Jiang, Andy Zou, Matt Fredrikson, Zachary Chase Lipton, J Zico Kolter

Keywords: Safe LLMs, Pretraining, Data Centric ML

TL;DR: We present a data-centric pretraining framework that builds safety into the model from the start

Abstract:

As large language models (LLMs) are increasingly deployed in high-stakes settings, the risk of generating harmful or toxic content remains a central challenge. Post-hoc alignment methods are brittle: once unsafe patterns are learned during pretraining, they are hard to remove.
In this work, we present a data-centric pretraining framework that builds safety into the model from the start. Our framework consists of four key steps:
(i) Safety Filtering: building a safety classifier to classify webdata into safe and unsafe categories;
(ii) Safety Rephrasing: we recontextualize unsafe webdata into safer narratives;
(iii) Native Refusal: we synthetically generate pretraining datasets that actively teach models to refuse on unsafe content and the moral reasoning behind it, and
(iv) Harmfulness-Tag annotated pretraining: we flag unsafe content during pretraining using a special token, and use it to steer models away from unsafe generations at inference-time.
Our safety-pretrained models reduce attack success rates from 38.8% to 8.4% on standard LLM safety benchmarks with no performance degradation on general tasks.

Primary Area: Social and economic aspects of machine learning (e.g., fairness, interpretability, human-AI interaction, privacy, safety, strategic behavior)

Flagged For Ethics Review: true

Submission Number: 22146