ML

Reward & Preference Distillation

Distilling alignment: transferring reward models, DPO as implicit distillation, RLHF teacher-student transfer, constitutional AI self-distillation, and maintaining alignment when compressing models.

Advanced

Prerequisites

Table of Contents

  1. Learning Objectives
  2. Notation
  3. Core Intuition
  4. The Alignment Distillation Problem
  5. Reward Model Distillation
  6. DPO as Implicit Distillation
  7. RLHF Transfer: Teacher Policy to Student
  8. Constitutional AI Self-Distillation
  9. Maintaining Alignment Under Compression
  10. Common Pitfalls
  11. Summary
  12. Exercises

Learning Objectives

  1. Explain why standard distillation can lose alignment.
  2. Design reward-preserving distillation objectives.
  3. Connect DPO to distillation from an implicit reward model.
  4. Describe how constitutional AI implements self-distillation.
  5. Maintain safety guarantees when quantizing/compressing aligned models.

Notation

  • rϕ(x,y)r_\phi(x, y) — reward model score
  • πref\pi_{\text{ref}} — reference policy
  • πT,πS\pi_T, \pi_S — teacher/student policies
  • β\beta — KL penalty coefficient

Core Intuition

An aligned LLM has two types of knowledge: (1) language modeling ability and (2) alignment (helpfulness, harmlessness, honesty). Standard distillation transfers language ability well but can lose alignment — the student might match the teacher's token distribution without matching its safety boundaries. Alignment distillation specifically preserves the teacher's preference structure, not just its output distribution.

Reward Model Distillation

Teacher RMStudent RMHelpful 0.920.90Harmful 0.080.10Neutral 0.450.45MSE: 0.0003 | T=2 softens reward targets
Temp T
2.00
Teacher rewardsStudent rewards
Explore: Reward model distillation transfers preference scoring from a large teacher RM to a smaller student — enabling efficient RLHF with lightweight reward models.

The Alignment Distillation Problem

Observation: After standard SFT distillation from an aligned teacher:

  • Student generates fluent, coherent text (language ability preserved).
  • Student may not refuse harmful requests (alignment partially lost).
  • Student confidence calibration differs (over/under-refuses).

Why alignment is fragile:

  1. Safety behavior is a small fraction of training signal (most tokens are "normal").
  2. Refusal behavior involves subtle conditional probabilities (refuse THIS but not THAT).
  3. Standard KL minimization weights all tokens equally — safety-critical tokens aren't upweighted.

Reward Model Distillation

Distill the reward model directly:

Lreward=Ex,y1,y2[ℓ(rϕ(x,y1)−rϕ(x,y2),  rψ(x,y1)−rψ(x,y2))],(1)\mathcal{L}_{\text{reward}} = \mathbb{E}_{x, y_1, y_2}\left[\ell(r_\phi(x, y_1) - r_\phi(x, y_2), \; r_\psi(x, y_1) - r_\psi(x, y_2))\right], \tag{1}

where rϕr_\phi is the teacher reward model and rψr_\psi is the student reward model.

Key insight: Transfer the PREFERENCE ordering, not absolute scores.

Then use the distilled reward for student RLHF:

max⁡πSEy∼πS[rψ(x,y)]−βDKL(πS∥πref).(2)\max_{\pi_S} \mathbb{E}_{y \sim \pi_S}[r_\psi(x, y)] - \beta D_{\text{KL}}(\pi_S \| \pi_{\text{ref}}). \tag{2}

DPO as Implicit Distillation

DPO (Rafailov et al., 2023): Directly optimize the policy from preference pairs:

LDPO=−E[log⁡σ(βlog⁡πθ(yw∣x)πref(yw∣x)−βlog⁡πθ(yl∣x)πref(yl∣x))].(3)\mathcal{L}_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\right)\right]. \tag{3}

Connection to distillation: If preference pairs come from a teacher model (teacher generates ywy_w, yly_l scored by teacher's implicit reward), then DPO IS reward distillation — the student learns the teacher's preference structure.

Teacher-generated DPO data:

  1. Teacher generates multiple responses to each prompt.
  2. Teacher (or its reward model) ranks them.
  3. Student trained via DPO on these ranked pairs.

Advantage over SFT distillation: Preserves the preference BOUNDARY (what's acceptable vs not), not just the modal output.

RLHF Transfer: Teacher Policy to Student

Three-stage distillation pipeline:

Stage 1: SFT distillation. Student learns basic language from teacher's outputs.

L1=−Ey∼πT[log⁡πS(y∣x)].(4)\mathcal{L}_1 = -\mathbb{E}_{y \sim \pi_T}[\log\pi_S(y|x)]. \tag{4}

Stage 2: Preference distillation. Student learns teacher's preferences via DPO.

L2=LDPO(πS;teacher-ranked pairs).(5)\mathcal{L}_2 = \mathcal{L}_{\text{DPO}}(\pi_S; \text{teacher-ranked pairs}). \tag{5}

Stage 3: Policy refinement. Optional RLHF on the student with the distilled reward model.

Result: Student that is both capable (from SFT) and aligned (from DPO/RLHF).

Constitutional AI Self-Distillation

Anthropic's CAI (Bai et al., 2022): The model distills alignment INTO itself:

Critique phase:

  1. Model generates response to potentially harmful prompt.
  2. Model critiques its own response: "Does this violate [principle]?"
  3. Model revises: "Rewrite to be helpful while avoiding harm."

Self-distillation:

  • The "critic" ability (from pre-training on ethical text) teaches the "generator" ability.
  • The revised responses become training data for the next iteration.
  • Each iteration: generator becomes more aligned without external human feedback.

Mathematical view: The model's critique function c(y∣x)c(y|x) is an implicit reward model. Revision is RLHF with this internal reward. Training on revised outputs is policy distillation from the revised policy.

Maintaining Alignment Under Compression

Problem: Quantizing/pruning an aligned model to INT4 can degrade alignment disproportionately.

Why: Safety behavior often depends on subtle probability differences (P(refuse) = 0.95 vs 0.05). Quantization noise can flip these decisions.

Solutions:

  1. QAT with alignment data: Include safety-critical examples in QAT fine-tuning.
  2. Distillation-aware quantization: After quantization, run brief DPO with teacher's preferences.
  3. Safety-critical layers in higher precision: Keep attention layers (which encode decision-making) at higher bit-width.
  4. Evaluate alignment separately: Don't just check perplexity — test refusal behavior, honesty, and helpfulness.

Common Pitfalls

Pitfall 1. Distilling only on "helpful" examples and ignoring "harmful" examples. The student never sees refusal behavior and learns to be unconditionally helpful (unsafe).

Pitfall 2. Using high temperature in alignment distillation. Safety boundaries require sharp distributions (high confidence in refusal). High temperature smooths these critical boundaries.

Pitfall 3. Assuming quantized models maintain alignment. Always re-evaluate safety after any compression. INT4 models frequently have degraded refusal rates.

Summary

  • Alignment distillation preserves safety/helpfulness, not just language ability.
  • Reward distillation: Transfer preference ordering between reward models.
  • DPO as distillation: Teacher-ranked pairs directly teach the student preferences.
  • CAI: Self-distillation of alignment via critique → revision loop.
  • Under compression: Alignment degrades; requires alignment-aware QAT or post-hoc DPO.
  • Three stages: SFT (capability) → DPO (alignment) → optional RLHF (refinement).

Exercises

Exercise 1. Design a preference dataset for distilling alignment from GPT-4 to a 7B student. Specify prompt categories, generation parameters, and ranking criteria.

Exercise 2. Prove that DPO with teacher-generated preference pairs is equivalent to minimizing the KL between student and teacher implicit reward models.

Exercise 3. For INT4 quantization of an aligned 13B model: design a safety evaluation benchmark and predict which safety behaviors are most likely to degrade.

Exercise 4. Compare the cost (GPU-hours) of three alignment distillation approaches: SFT-only, SFT+DPO, and SFT+RLHF for a 7B student.

Exercise 5. Explain why constitutional AI (CAI) can be seen as both self-distillation and reinforcement learning simultaneously.