🤖 machine learning

Curriculum Learning for Safety Alignment

This paper proposes Staged-Competence, a curriculum learning framework that organizes preference data by difficulty and progressively updates the reference model to significantly improve the out-of-distribution robustness and jailbreak resistance of Direct Preference Optimization (DPO) for safety alignment while preserving general capabilities.

Original authors: Sandeep Kumar, Virginia Smith, Chhavi Yadav

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sandeep Kumar, Virginia Smith, Chhavi Yadav

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Teaching AI to Be Safe is Brittle

Imagine you are training a very smart robot to be a helpful assistant. You want it to be kind and helpful, but you also need to make sure it never agrees to do something dangerous, like building a bomb or stealing a car.

Currently, the standard way to teach this (called DPO) is like throwing a bunch of "Good vs. Bad" examples at the robot all at once, in random order.

The Flaw: The paper argues this method is "brittle." It's like teaching a student to drive by throwing them into heavy traffic immediately. They might learn the rules for that specific traffic jam, but if you take them to a different city (a new situation) or someone tries to trick them with a weird question (a "jailbreak" attack), they often crash. They haven't truly learned the concept of safety; they just memorized the specific examples they saw.

The Solution: The "Staged-Competence" Syllabus

The authors propose a new method called Staged-Competence. Think of this as switching from a chaotic pile of flashcards to a structured, step-by-step school syllabus.

They use a concept called Curriculum Learning, which is the idea that you should learn things from easy to hard.

Here is how their system works, broken down into three simple steps:

1. Grading the Homework (Difficulty Scoring)

Before the robot starts training, the system looks at every single example of "Good vs. Bad" behavior.

The Old Way: Just pick examples randomly.
The New Way: The system asks, "How hard is this for the robot right now?"
- If the robot already knows the answer easily, that's an Easy example.
- If the robot is confused or likely to get it wrong, that's a Hard example.
- Analogy: Imagine a teacher grading a stack of math problems. They don't just hand them out randomly. They sort them: "Here are some 1+1 problems (Easy), then some 2-digit multiplication (Medium), and finally some calculus (Hard)."

2. The Three-Stage School (Staged Training)

Instead of trying to learn everything at once, the training is split into three stages (buckets).

Stage 1: The robot only sees the "Easy" examples. It practices until it masters them.
Stage 2: The robot is now "competent" enough to handle "Medium" difficulty. It learns these, building on what it learned in Stage 1.
Stage 3: Finally, it tackles the "Hard" examples.

The Secret Sauce: In between these stages, the system updates the robot's "internal teacher" (the reference model).

Analogy: Imagine a coach. In the first week, the coach teaches you how to hold a racket. Once you master that, the coach doesn't keep teaching you how to hold the racket; they update their expectation and start teaching you how to swing. The robot "grows up" as it moves through the stages, so it doesn't waste time re-learning basics it already knows.

3. Smart Sampling (Competence-Based)

Even within a single stage, the robot doesn't just see the examples in a random order. It gets a mix of what it can handle and slightly harder challenges, gradually increasing the difficulty.

Analogy: Think of a video game. You don't start with the final boss. You start with the tutorial, then the first level, then the second. As you get better, the game automatically unlocks harder levels. This paper does that for AI safety.

What Did They Find? (The Results)

The authors tested this on three different types of AI models. Here is what happened:

Stronger Safety: The robots trained with this "syllabus" were much better at resisting "jailbreak" attacks (tricks used to make AI say bad things). They were 20% better at ignoring these tricks than the standard method.
Better Generalization: When asked about things they hadn't seen before (Out-of-Distribution), they were 16% less likely to give a harmful answer.
Deeper Understanding: Standard training often only teaches the AI to say "No" in the very first sentence, then it might slip up later. This new method teaches the AI to stay safe throughout the entire conversation, like a guard who stays alert the whole time, not just at the door.
Data Efficiency: They found that using this method with 25% less data still produced results as good as the standard method using 100% of the data. It's like getting a better education with fewer textbooks.
No Loss of Smarts: Crucially, making the AI safer didn't make it dumber. It could still answer normal questions just as well as before.

The "Cleaned" Dataset Bonus

The paper also mentions a side discovery. The two big public datasets they used to train these models were actually messy.

Sometimes, the "Safe" answer was actually dangerous.
Sometimes, the "Unsafe" answer was actually helpful.
Analogy: It was like a teacher's answer key being wrong.
The authors cleaned up these datasets, fixed the mistakes, and made a new, cleaner version available for everyone to use.

Summary

The paper argues that to make AI truly safe and robust, we shouldn't just throw data at it randomly. Instead, we should teach it like a human student: start with the basics, build competence, and gradually increase the difficulty. This "Staged-Competence" approach makes AI safer, more resistant to tricks, and more efficient to train, without making it less helpful.

Technical Summary: Curriculum Learning for Safety Alignment

Problem Statement

Direct Preference Optimization (DPO) has become a standard approach for safety-aligning Large Language Models (LLMs) by learning from human preference pairs without a separate reward model. However, prior research indicates that standard DPO is brittle, exhibiting poor generalization to out-of-distribution (OOD) prompts and failing to resist simple jailbreaking attacks. Furthermore, safety alignment in standard DPO often concentrates only in the first few tokens of a response, leaving subsequent tokens vulnerable. While curriculum learning (ordering examples from easy to hard) has shown success in improving convergence and robustness in other domains like machine translation and general alignment, its potential to address the specific brittleness of DPO-based safety alignment remains largely unexplored.

Methodology: Staged-Competence

The authors propose Staged-Competence, a curriculum-based framework designed to improve the robustness of DPO-based safety alignment. The method consists of two primary phases:

Phase 1: Difficulty Scoring (Preference Alignment Margin)

Unlike previous approaches that rely on static heuristics or local prompt-based scoring, Staged-Competence defines difficulty in a model-dependent manner.

Mechanism: For each prompt $x_i$ , the unaligned base model generates a zero-shot response $\hat{y}_i$ .
Scoring: The authors compute embeddings for the zero-shot response, the safe (chosen) response $y^+_i$ , and the unsafe (rejected) response $y^-_i$ . They define a preference alignment margin ( $m_i$ ):
$m_i = \cos(e_{\hat{y}_i}, e_{y^+_i}) - \cos(e_{\hat{y}_i}, e_{y^-_i})$
Ordering: A large positive margin indicates the model already distinguishes safe from unsafe behavior (easy example), while a small or negative margin indicates misalignment (hard example). The dataset is sorted globally by this margin to create an easy-to-hard curriculum.

Phase 2: Curriculum Training

The training algorithm combines staged reference-model updates with competence-based sampling to ensure a granular easy-to-hard progression both within and across stages.

Staging: The sorted dataset is divided into $K$ buckets (e.g., $K=3$ ) of increasing difficulty.
Reference Updates: Between stages, the reference policy is updated ( $\pi^{(k+1)}_{ref} = \pi^{(k)}$ ). This allows the policy in the next stage to focus on assimilating the current difficulty bucket rather than re-learning previously acquired concepts.
Competence-Based Sampling: Within each stage, a competence function $c(t)$ (based on a square-root schedule) determines a difficulty threshold. Only examples with difficulty $d_i \leq c(t)$ are eligible for mini-batch sampling at step $t$ . This ensures harder examples are introduced at a progressively slower rate, allowing the model to consolidate learning.

Key Contributions

Staged-Competence Framework: The first systematic study applying curriculum learning to DPO-based safety alignment, introducing a global, model-dependent difficulty ordering and a hybrid training algorithm.
Cleaned-PKU-HH-SafeRLHF Dataset: The authors identify widespread inconsistencies in popular safety datasets (PKU-SafeRLHF and HH-RLHF), where "chosen" responses are often unsafe and "rejected" responses are often safe. They provide a cleaned, combined dataset filtered by a GPT-4o-mini judge to ensure high-quality preference labels for training.
Ablation and Analysis: A systematic evaluation of curriculum design choices (ordering, sampling, reference updates) across three model families (LLaMA-3, Qwen3, Yi-1.5).

Experimental Results

Experiments were conducted across three model families (LLaMA-3-8B, Qwen3-8B, Yi-1.5-9B) using the cleaned dataset.

Robustness to OOD and Attacks: Staged-Competence reduced OOD harmful response rates by an average of 16% and jailbreak attack success rates by 20% compared to standard DPO, while preserving general capabilities (MMLU, HellaSwag) and maintaining near-zero over-refusal.
Alignment Depth: Analysis of per-token suppression showed that Staged-Competence achieves roughly 3× greater reward margin separation and sustains safety suppression throughout the response, not just in the initial tokens. This deeper alignment correlates with improved resistance to prefill-based jailbreak attacks.
Data Efficiency: Staged-Competence matched the safety performance of standard DPO trained on 100% of the data using only 75% of the training data.
Scaling: The safety advantage of Staged-Competence over standard DPO widened monotonically with model size. While standard DPO performance degraded significantly as model size increased (from 1.7B to 8B parameters), Staged-Competence maintained consistent safety performance across scales.

Significance and Claims

The paper claims that Staged-Competence demonstrates that curriculum learning is a meaningful and underutilized lever for safety alignment. By organizing preference data by difficulty and progressively updating the reference model, the method produces more robust safety features that generalize better to OOD scenarios and adversarial attacks.

The authors emphasize that their approach is agnostic to the underlying policy optimization loss, meaning it can extend to other DPO variants and alignment domains beyond safety. They conclude that while standard DPO is brittle, structured easy-to-hard learning can significantly enhance the robustness and data efficiency of safety alignment without sacrificing general model capabilities.