PaperlayerView exact revision on arXiv

Source-linked AI summary

Safety Pretraining: Toward the Next Generation of Safe AI

Pratyush Maini, Sachin Goyal, Dylan Sam, Alex Robey, Yash Savani, Yiding Jiang, Andy Zou, Matt Fredrikson, Zacharcy C. Lipton, J. Zico Kolter

arXiv:2504.16980v2cs.LG

TL;DR

Post-hoc alignment can leave harmful information learned during pretraining difficult to remove. The paper embeds safety in pretraining through data interventions and harmfulness tags; attack success rate on safety benchmarks falls from 38.8% to 8.3%.

  • Problem

    Post-hoc alignment can superficially reduce toxic generation, but unsafe information learned during pretraining is difficult to unlearn.

    Source: §1 · p. 3

  • Method

    The framework combines safety filtering, synthetic recontextualization, refusal and moral education datasets, and harmfulness-tag training to embed safety in pretraining.

    Source: §1 · p. 3, §8 · p. 13

  • Results

    Attack success rate (ASR) falls from 38.8% to 8.3% on safety benchmarks in the reported SafeLM result.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors argue that safety interventions during pretraining can yield robust safety gains and help models understand why requests are unsafe.

    Source: §8 · pp. 13–14

    The paper notes modest helpfulness impacts via over-refusal, despite no performance degradation on standard NLP benchmarks.

    Source: §1 · p. 4

Abstract

from arXiv · show

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 develop RefuseWeb and Moral Education pretraining datasets that actively teach model 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 model away from unsafe generations at inference. 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.

1 Introduction

Because post-hoc alignment can suppress harmful outputs without removing learned unsafe information, the paper embeds safety into pretraining through data curation, recontextualization, refusal training, and evaluation tools.

1 Introduction

  • Post-hoc alignment methods have shown superficial reductions in toxic generation, while unsafe information learned during pretraining is difficult to unlearn.The paper frames this gap as “alignment is not unlearning.”Source: §1 · p. 3
  • The framework combines safety filtering, synthetic recontextualization, native refusal and moral education, and harmfulness-tag training to separate safe from unsafe content.It also introduces evaluation benchmarks for base models before safety tuning.Source: §1 · p. 3
  • The safety classifier and Data Safety Report Cards support identifying harmful material and reporting dataset safety scores and harmful-content frequencies.The report cards use frequencies from the MLCommons safety taxonomy.Source: §1 · p. 3
  • Synthetic recontextualization reframes potentially harmful data in safe, context-rich narratives while aiming to preserve its underlying information.The paper introduces SafeWeb, a 100-billion-token synthetic corpus produced using 12 diverse prompts.Source: §1 · p. 3
  • RefuseWeb teaches models to disengage from harmful prompts, while harmfulness tags mark unsafe passages and SafeBeam steers generation away from them.The paper also presents completion-based tools for assessing base-model safety during pretraining.Source: §1 · p. 3
  • Attack success rate (ASR) falls from 38.8% to 8.3% on safety benchmarks in the reported SafeLM result.The work releases a 1.7B-parameter SafeLM and an interactive Safe Playground.Source: §1 · p. 3
  • The authors describe pretraining-time safety as a foundation for safer AI, while noting modest helpfulness impacts from over-refusal.They state that safety enhancement does not reduce performance on standard NLP benchmarks.Source: §1 · p. 4

2 Related Work

Prior safety work has largely focused on post-training, while related pretraining approaches include preference conditioning, toxicity filtering, and synthetic data; this paper extends safety integration into pretraining.

2 Related Work

  • Post-training methods such as RLHF and supervised fine-tuning can require substantial preference data and may fail to prevent harmful-content extraction under adversarial pressure.Shi et al. (2023) applied DPO during instruction tuning, but the cited studies center on post-hoc tuning rather than pretraining.Source: §2 · p. 4
  • Prior pretraining approaches include conditional training on human preferences and heuristic toxicity filtering, although residual harmful content can remain in filtered corpora.The authors extend synthetic safety data generation from finetuning to pretraining with a 300B-token dataset.Source: §2 · p. 4
  • Special-token associations for harmful content have been explored through fine-tuning, and the authors present harmfulness-tag pretraining as an extension into pretraining.The paper states that fine-tuning-based associations fail under benign fine-tuning.Source: §2 · p. 4
  • The paper builds on existing safety benchmarks by releasing safety datasets, base-model evaluation tools, and data safety reporting standards.The related work names SafetyPrompts, the MLCommons AI Safety Benchmark, AdvBench, JailbreakBench, and HarmBench.Source: §2 · p. 4

3 Pre-training Data Interventions

The section develops pretraining data interventions that assess safety risk, preserve useful information through recontextualization, teach refusal and ethical reasoning, and report dataset safety.

3.1 Safety Scoring

  • Safety scoring combines classifier strategies to identify harmful data while balancing safety filtering against the loss of informative content.Safety annotations assign scores from 0 to 5; the final score takes the maximum of the LLM-based and embedding-based approaches to prioritize recall on unsafe examples.Source: §3.1 · p. 5

3.2 Synthetic Recontextualization

  • Over 100B tokens of synthetic recontextualized data are produced by rewriting potentially unsafe material into context-rich educational content while preserving essential factual information.The process uses LLaMA-3.1-8B to generate alternatives and aims for each sentence to remain safe when read in isolation.Source: §3.2 · p. 5

3.3 Refusing the Web

  • RefuseWeb turns FineWeb examples scored 4 or 5 into dialogues where the assistant refuses harmful requests and gives an educational rationale.The requests cover categories including harassment, malware, fraud, physical harm, and privacy violations.Source: §3.3 · p. 6

3.4 Moral Education Data

  • Moral Education transforms harmful-request and refusal dialogues into educational articles that convey ethical reasoning beyond conversational settings.The dataset is intended to avoid knowledge gaps from simply excluding harmful content.Source: §3.4 · p. 7

3.5 Data Safety Report Cards, A New Standard

  • Data Safety Report Cards pair safety-score distributions with harmful-content frequencies to make pretraining dataset risks more interpretable.The report uses a 14-category harmful-content taxonomy and Infini-gram queries; lower-scored data slices correspond to lower query frequencies, and rephrasing and contextualization reduce harmful content within each slice.Source: §3.5 · p. 7, Figure 1 · p. 8

4 Harmfulness-Tag Annotated Pretraining

Unsafe passages receive a dedicated harmfulness tag during pretraining, and Safe Beam Search uses the tag’s predicted probability to filter candidate continuations during inference.

4 Harmfulness-Tag Annotated Pretraining

  • Unsafe segments receive a dedicated harmfulness tag at randomly selected positions comprising 5% of the input sequence, while safe content is never tagged.The tag is intended to distinguish unsafe from safe representations and can be used to steer generation at inference.Source: §4 · p. 8

4.1 Inference-Time Steering via Safe Beam Search

  • Safe Beam Search uses a one-token lookahead to estimate harmfulness-tag probability, discards the 50% of beams with the highest probability, then selects the top k by log-likelihood.This filtering is intended to avoid unsafe continuations while maintaining fluency and coherence through likelihood-based selection.Source: §4.1 · p. 8

5 Pretraining and Post-Training Setup

The setup uses a shared 1.7B-parameter pretraining recipe and corpus, then compares models under matched post-training conditions. Harmfulness-tag models receive a small amount of tagged instruction-tuning data.

Pretraining

  • All experiments use 1.7B-parameter models, the SmolLM2 pretraining setup, and a shared corpus of FineWeb-Edu, StackOverflow, FineMath, and Cosmopedia.Training uses LitGPT, FlashAttention-2, mixed precision, and the original SmolLM2 optimization hyperparameters.Source: §5 · p. 9

Post-training

  • Models are instruction-tuned on Ultrachat-200k together with WildGuardMix and WildJailbreak safety data to establish a post-hoc alignment baseline.The authors use this baseline to assess the additional effects of safety-aware pretraining.Source: §5 · p. 9
  • Harmfulness-tag models also receive 10% tagged WildGuardMix completions during instruction tuning, while the remaining dataset composition stays identical across experiments.Source: §5 · p. 9

6 Evaluations & Results

The evaluations test general-task performance, safety across training stages, benign-request helpfulness, and robustness after benign finetuning. Safety-pretrained models retain safety more robustly after benign finetuning than models relying on instruction tuning alone.

6.1 Performance on Standard Benchmarks

  • Safety-intervention models perform comparably to models trained on raw web data, while training only on a safe subset risks losing factual coverage and expressivity.The benchmarks cover reasoning, factual recall, commonsense, and mathematics.Source: §6.1 · pp. 9–10

6.2 Safety Evaluations

  • 44% ASR for the standard pretrained base model is four times the safety-pretrained model’s 11% ASR; instruction-tuned models later show near-zero ASR, which the authors call misleading.The base-model evaluations use completion-style prompts adapted from harmful-request datasets.Source: §6.2 · p. 10
  • The jailbreak evaluation learns adversarial suffixes on 100 JailbreakBench instances, but those attacks do not account for SafeBeam’s modified inference algorithm.The authors identify inference-algorithm-aware attacks as future work.Source: §6.2 · pp. 10–11

6.3 Overrefusal Evaluations

  • On 300 Alpaca requests, safety pretraining does not seem to reduce helpfulness or increase overrefusal, while SafeBeam slightly increases overrefusal.GPT-4o-mini classifies partial and full refusals as overrefusal in this evaluation.Source: §6.3 · p. 11

6.4 Benign Finetuning Robustness

  • Standard pretrained models nearly quadruple ASR after benign finetuning, while safety-pretrained models show only a marginal increase.Figure 2 compares ASR for standard pretraining, safety pretraining, and safety pretraining with SafeBeam across base, post-instruction-tuning, and post-benign-finetuning settings.Source: Figure 2 · p. 10, §6.4 · p. 11

7 Ablations

The ablations measure ASR after benign finetuning on GSM8k and test data interventions, harmfulness-tag placement, and pretraining versus instruction-tuning annotation. Results distinguish the effects of retaining unsafe-topic exposure, adding refusal and moral-education data, and incorporating tags during pretraining.

7 Ablations

  • The ablations use ASR after GSM8k benign finetuning as the safety metric, with lower ASR indicating better results.Source: §7 · p. 11, §7.1 · p. 12

7.1 Importance of various Data-Centric Interventions

  • 43.8% ASR for score-0-only training compares with 38.8% for raw data, whereas rephrased score 1–3 data yields 33.6%.The authors attribute the score-0-only result to lack of exposure to unsafe patterns and describe rephrasing as context-aware, educational framing.Source: §7.1 · p. 12
  • 25.1% ASR follows adding refusal completions, an 8.5-point improvement; adding moral-education data lowers ASR further to 20.0%, the study’s lowest result.Refusal completions address score-4 and score-5 content, while moral-education data supplies narratives and explanations about risks and ethics.Source: §7.1 · p. 12

7.2 Harmfulness-Tag Annotation During Pretraining vs. Finetuning

  • ∼8% ASR after benign finetuning for harmfulness-tag pretraining contrasts with ∼30% for harmfulness-tag instruction-tuning, an almost fourfold difference.The comparison tests annotation during pretraining against standard pretraining followed by tag-annotated instruction tuning.Source: §7.2 · p. 13

7.3 Amount of Harmfulness-Tag to Use

  • The harmfulness-tag ablation tests 3%, 5%, and 10% injection rates; 5% yields the lowest ASR among those tested.The tested models also use the standard rephrasing, refusal, and moral-education interventions.Source: §7.3 · p. 13

8 Discussion

The discussion argues that post-hoc alignment can suppress but not remove harmful capabilities, motivating safety built into pretraining. It favors structured instruction over censoring alone, reporting stronger safety gains and better understanding of why requests are unsafe.

8 Discussion

  • Post-hoc alignment can be brittle under adversarial prompting because it redirects or suppresses harmful capabilities rather than meaningfully removing them.The discussion argues that even slight perturbations can expose unsafe behavior.Source: §8 · p. 13
  • The framework combines safety filtering, synthetic recontextualization, refusal and moral education datasets, and harmfulness tagging.The discussion links these interventions to safety risks embedded during early pretraining being harder to remove than superficial behavioral risks.Source: §8 · p. 13
  • Censoring all unsafe examples proved ineffective: models trained only on safe content remained brittle and could not handle harmful prompts responsibly.The most robust gains instead came from interventions modeled on human pedagogical practices.Source: §8 · pp. 13–14
  • Interventions modeled on human pedagogical practices reduced attack success rates and improved the model’s understanding of why requests are unsafe.The discussion contrasts understanding why a request is unsafe with merely recognizing that it is.Source: §8 · p. 14
  • The authors envision safety pretraining as a foundation for responsible AI development, with safety built in rather than bolted on.Their analogy emphasizes careful instruction and internalization rather than censorship.Source: §8 · p. 14

A.1 Standard Benchmark Performance Evaluation

The supplied appendix passages describe a benchmark-comparison table, examples of harmful completions, an embedding-based safety classifier approach, and a harmfulness-tagging algorithm.

A.1 Standard Benchmark Performance Evaluation

  • A.1’s Table 1 compares models trained with different data-centric safety interventions on standard benchmarks.The supplied passage identifies the table’s scope but gives no benchmark values.Source: §A.1 · p. 17

A.2 Jailbroken Examples

  • A.2 provides examples of harmful completions generated by SmolLM-1.7B-Instruct in Table 2.The supplied passage does not state the examples’ contents.Source: §A.2 · p. 17

B.1 Embedding-Based Approach

  • B.1 classifies content into six safety categories using frozen BERT embeddings and a linear head, and adds higher-toxicity content to mostly clean FineWeb data for training.The appendix also describes comparing embedding models for classification on safety data.Source: §B.1 · p. 17

Algorithm 2: Harmfulness-Tag Annotated Pretraining

  • Algorithm 2 inserts Harmfulness-Tag tokens before words in unsafe text with probability p and returns the tagged sequence.The algorithm initializes the modified text with the first word, then iterates over the remaining words.Source: Algorithm 2: Harmfulness-Tag Annotated… · p. 18

B.2 LLM-Based Approach

The LLM-based approach evaluates instruction-tuned models for safety scoring and refines its setup on sampled FineWeb examples. The passages report JSON-formatting difficulties and systematic over-scoring by LLaMA models.

B.2 LLM-Based Approach

  • The approach experiments with instruction-tuned LLaMA, Qwen, and Phi models for safety assessment.The listed variants include LLaMA 3.1, Qwen 2.5, and Phi-3-mini-4k.Source: §B.2 · p. 18
  • The authors refined the approach on 128 randomly sampled FineWeb examples, selecting the model that produced the most reliable safety scores.Source: §B.2 · p. 18
  • LLaMA models had an approximately 25% JSON compliance error rate, so requests began with structured JSON output, excluding reasoning before the JSON.Source: §B.2 · p. 18
  • After JSON formatting was enforced, LLaMA models often over-scored safety risks relative to GPT-4o-mini and human annotations, producing false positives.They frequently assigned scores of 4 or 5 to examples scored 0 by the reference oracle and human reviewers.Source: §B.2 · p. 19

B.3 Classifier Comparisons

This section documents safety-classifier comparisons and a 0–5 safety-scoring rubric, alongside prompts for rephrasing unsafe web data, synthesizing refusals and moral-education content, and evaluating model responses with task-specific LLM judges. It also describes iterative prompt refinement aimed at producing safer rephrasals that retain more factual information.

B.3 Classifier Comparisons

  • Ensembling the classifiers strengthens the safety filter, while LLM-based and ensemble approaches produce more stringent filtering than embedding classifiers and existing filtering strategies perform poorly on GPT-4o-mini annotations.The study prioritizes Safety Recall (≥1) to prevent non-harmful data from leaking into safer training sets; Figure 5 notes that stringent filtering comes at the cost of over-predicting actual score-0 instances.Source: §B.3 · p. 19, Figure 5 · p. 20
  • Classifier failure cases show that some models assign inflated safety scores to benign climate-awareness and educational content, motivating future work on reducing excessive risk aversion while retaining strong filtering.For example, climate information was labeled as financial advice and harmful activity, while Christmas worksheets were labeled as illegal activity.Source: §B.3 · pp. 20–21

C.1 Safety Scoring Prompt

  • The rubric assigns the highest applicable score from 0 (safe) to 5 (highly unsafe), with criteria spanning mild inappropriate content through harmful advice, self-harm references, hate speech, threats, and privacy violations.Higher scores indicate lower safety, and the assessment instructions require returning both a score and a brief reason.Source: §C.1 · pp. 22–23

C.2 Synthetic Recontextualization

  • Recontextualization prompts vary stylistic formats while asking for age-appropriate explanations that retain core ideas, clarify sensitive material, and keep each sentence safe in isolation.The implementation randomly selects a template at runtime to produce stylistically diverse rephrasals.Source: §C.2 · pp. 24–26
  • Iterative refinement found that teacher or educator scripts with sentence-level safety instructions worked better than prompts that over-sanitized content or remained verbose and omitted key points.The refined prompt preserved more factual information while keeping sentences self-contained and contextually safe.Source: §C.2 · pp. 26–27
  • Iterative prompt refinements produced final rephrasing prompts using LLaMA-3.1-8B, selected for a balance of speed, fluency, and contextual depth.The diverse prompt templates for synthetic data generation are provided in Appendix C.2.Source: §C.2 · p. 27

C.3 Refusal Prompts

  • RefuseWeb prompts turn harmful requests into respectful User–Assistant dialogues in which the assistant firmly refuses, explains the harm, and suggests constructive alternatives when appropriate.The prompt covers categories including harassment, malware, physical and economic harm, fraud, disinformation, sexual content, and privacy.Source: §C.3 · p. 27

C.4 Moral Education Prompt

  • Moral-education prompts transform refusal dialogues into cohesive, third-person educational web content focused on the ethical principle rather than the specific exchange.The resulting content is framed as a moral lesson or ethical guideline and uses headings and transitions typical of educational websites or blogs.Source: §C.4 · p. 28

C.5 LLM Judges

  • Task-specific judges use HarmBench’s modified template for safety evaluations and a five-way helpfulness rubric: compliant and helpful, compliant and unhelpful, partial refusal, full refusal, or shortcircuit.The helpfulness rubric distinguishes partial from full refusals and includes a shortcircuit category for outputs containing potentially harmful content.Source: §C.5 · pp. 28–29

D Safety Data Report Card Details

The report card compares harmful-content frequencies in safety-rephrased data with original FineWeb and describes the harmful-query categories used for corpus measurement.

D Safety Data Report Card Details

Loading 2504.16980v2…