Safety-Constrained Distillation

• Safety-constrained distillation is a set of techniques that transfer safety properties from large teacher models to compact student models using explicit safety constraints.
• The methodology integrates augmented supervision, such as HarmAug, and adaptive weighting strategies to maintain detection of harmful content and ethical compliance.
• Practical implementations span language model safety guards, defense against extraction attacks, and safe reinforcement learning, achieving strong performance with reduced resources.

Safety-constrained distillation refers to the class of knowledge distillation techniques that compress large, safety-critical teacher models into smaller student models while explicitly preserving or enhancing safety properties. Safety, in this context, includes the reliable detection of harmful content, enforcement of ethical constraints, or satisfaction of task-specific cost and risk boundaries. This paradigm is of increasing relevance across domains such as LLM guardrails, medical LLMs, and safe reinforcement learning, where distributional shifts and adversarial queries demand not just model fidelity but robust transfer of safety behaviors.

1. Formal Frameworks for Safety-Constrained Distillation

Several works articulate safety-constrained distillation as optimization under explicit safety constraints. Formally, let $p^{(T)}$ denote the teacher model (possibly an ensemble), $p^{(S)}$ the student, and $\mathcal D$ the data distribution.

The abstract safety-constrained distillation objective, as framed in the axiomatic, multi-teacher context, is:

$$\min_{\theta \in \Theta} \; \mathcal{L}_{\mathrm{KD}}(\theta; G) \quad \text{s.t.} \quad \mathrm{Safety}_{\mathcal D}(\theta) \geq \mathrm{Safety}_{\min}$$

where $\mathcal{L}_{\mathrm{KD}}$ is the (possibly adaptive-weighted) distillation loss between $p^{(S)}$ and the teacher ensemble, $G$ is an adaptive weighting operator, and $\mathrm{Safety}_{\mathcal D}$ is an expected safety metric defined over $\mathcal D$ (e.g., harmfulness detection, refusal rate, constrained cost). The safety constraint is enforced either as a hard feasibility condition, projection, or via a Lagrangian multiplier (Flouro et al., 25 Jan 2026).

Concretely, in LLM safety guard distillation with HarmAug, the objective is a convex combination of soft (KL-divergence) and hard (binary cross-entropy) losses using both real and teacher-labeled harmful examples:

$$\phi^* = \arg\min_{\phi}\,\frac{1}{n}\sum_{i=1}^n\Bigl[(1-\lambda)\mathrm{KL}(p_\theta(\cdot|x_i,y_i)\,\|\,q_\phi(\cdot|x_i,y_i)) + \lambda\,\mathcal{L}_\mathrm{BCE}(q_\phi(x_i,y_i),c_i)\Bigr]$$

where $\lambda$ trades off soft and hard supervision (Lee et al., 2024).

2. Safety Constraints: Axiomatic and Algorithmic Perspectives

Safety constraints in distillation manifest at multiple levels:

• Token-level: Safety-critical outputs (e.g., refusal tokens, harmful continuations) are weighted to favor safer teachers. The safety-monotonicity axiom ensures that the weight for a safer teacher on safety-critical tokens is at least as large as that for any less safe teacher (Flouro et al., 25 Jan 2026).
• Context-level: For safety-critical contexts, the aggregation weights in the ensemble are similarly monotonic in each teacher’s safety.
• Aggregate optimization: The overall learning process is constrained so that the distilled student achieves a minimum required safety metric, such as thresholded harmfulness recall, minimum refusal rate, or upper-bounded policy cost (in RL CMDPs).

The existence, non-uniqueness, and composition of such safety-aware operators are established in an axiomatic framework, with precise stability and robustness guarantees under perturbations of safety weights (Theorems 4.1–4.2, 5.3 in (Flouro et al., 25 Jan 2026)).

3. Domain Applications and Methodological Variants

3.1 LLM Guard Distillation

Safety-constrained distillation is critical for deploying lightweight safety guard models (e.g., for LLMs). The HarmAug method demonstrates that naively distilled students underperform due to insufficient diversity of harmful examples. Instead, HarmAug augments instruction–response data by jailbreaking LLMs to produce diverse harmful prompts, generating both refusal and harmful continuations, and relabeling these through the teacher. Empirically, this closes the performance gap: a 435M-parameter student model matches or exceeds its 8B-parameter teacher in both F1 and AUPRC, with nearly 4× speed gains and $~12\%$ memory usage (Lee et al., 2024).

Model	Size	AUPRC (avg)
Llama-Guard-3	8 B	0.7665
DeBERTa (no aug)	435M	0.7962
+ EDA	435M	0.8161
+ GFN	435M	0.8050
+ HarmAug (ours)	435M	0.8362

Increasing harmful-sample diversity (from $65$ to $332$ instruction clusters) enables small models to robustly replicate the teacher’s nuanced safety boundary.

3.2 Model Extraction and Alignment Collapse

Safety-constrained approaches are crucial in settings vulnerable to extraction. Black-box distillation of medical LLMs demonstrates that naive behavioral cloning strips away refusal behavior: a LoRA-adapted surrogate produces unsafe completions for $86\%$ of adversarial prompts, exceeding the teacher’s $66\%$ (Jahan et al., 10 Dec 2025). Defenses include embedding safety signals (refusal templates, canary prompts), multi-objective optimization (imitating both outputs and safety behaviors), and real-time canary drift detection.

3.3 Safe Reinforcement Learning

In constrained MDPs, safe policy distillation seeks to compress high-capacity decision transformers into lightweight policies while respecting cumulative cost limits:

$$\pi^* = \arg\max_\pi V_r^\pi(\mu_0) \quad \text{s.t.} \quad V_c^\pi(\mu_0) \leq \kappa$$

The GOLD framework performs offline-to-online distillation: the student model is trained with reward shaping or explicit cost penalties to remain within safety thresholds, aggregating guidance from expert demonstrations (Li et al., 2023).

4. Practical Recipes, Implementation, and Evaluation

Algorithmic recipes for safety-constrained distillation are as follows:

• Aggregation and Weighting: Use adaptive weights to combine teacher distributions, increasing safe teacher influence in safety-critical scenarios or tokens.
• Objective design: Include explicit loss terms for safety metrics (e.g., refusal accuracy, cost) or use teacher-generated safety labels on augmented data (Lee et al., 2024, Jahan et al., 10 Dec 2025).
• Optimization and guarantees: Apply SGD with projection or Lagrangian dual ascent to ensure $$\mathrm{Safety}_{\mathcal D}(\theta)\ge \mathrm{Safety}_{\min}$$ (Flouro et al., 25 Jan 2026).
• Defense strategies: Use canary prompts, refusal pattern monitoring, and layered alignment drift detection (DistillGuard++), especially in API or black-box settings (Jahan et al., 10 Dec 2025).

Empirical benchmarks consistently show that with adequate safety-centric augmentation and constraint enforcement, sub-billion-parameter students can match or outperform their multi-billion-parameter teachers on safety-critical tasks, reducing cost, latency, and deployment barriers (Lee et al., 2024, Li et al., 2023).

5. Limitations and Open Problems

Current methods require careful selection of safety metrics and may incur additional overhead for data augmentation (e.g., HarmAug generation or canary synthesis). Guaranteeing hard safety constraints, especially under distribution shift or adversarial query regimes, remains challenging. For example, reward-shaping in RL lacks formal safety certificates; extraction-aware defenses in LLMs require comprehensive behavioral coverage (Li et al., 2023, Jahan et al., 10 Dec 2025). Tuning of weighting schemes and safety thresholds may introduce further complexity, as in adaptive, multi-scale ensemble frameworks (Flouro et al., 25 Jan 2026).

A plausible implication is that future work will need to combine formal guarantees (hard constraints, shielding), richer safety signals (contextual, token-level, temporal), and robust, adaptive evaluation protocols.

6. Extensions and Future Research Directions

Further developments are expected in the following areas:

• Extension to multimodal and multi-task systems: Generalizing safety-constrained distillation to encompass diverse inputs (text, vision, action), hierarchical tasks, and compositional safety objectives (Flouro et al., 25 Jan 2026).
• Cryptographic protection of teacher knowledge: Safe Distillation Box (SDB) demonstrates that proxy-key mechanisms can block unauthorized distillation while augmenting authorized performance, preserving intellectual property (Ye et al., 2021).
• Formal robustness and adaptivity: Stability, perturbation, and drift robustness are addressed via axiomatic approaches; research aims to integrate certified safety with scalable, operator-agnostic optimization pipelines (Flouro et al., 25 Jan 2026).
• On-device and real-time deployment: Empirical results show that safety-constrained distillation enables resource-efficient, high-fidelity safety models compatible with edge and mobile deployment (Lee et al., 2024).
• Adversarial threat modeling: Analysis of extraction attacks and defense layering remains an open topic, requiring integrated benchmarks for red-teaming and continual behavioral vetting (Jahan et al., 10 Dec 2025).

7. Comparative Summary

Framework/Domain	Safety Constraint	Methodological Highlight	Key Result	Reference
HarmAug (LLM Guard)	Harmfulness detection	Augmented harmful response synth	Student matches/exceeds 8B teacher	(Lee et al., 2024)
Axiomatic Multi-Teacher	Safety monotonicity	Scalable, multi-scale adaptive KD	Existence and stability guarantees	(Flouro et al., 25 Jan 2026)
Safe Distillation Box	IP protection, authorized KD	Key-based proxy streams	2.8–8.1% performance penalty for unauthed	(Ye et al., 2021)
Black-Box Extraction (Med LLM)	Alignment collapse risk	Layered canary, refusal defenses	86% vs 66% unsafe completion rate	(Jahan et al., 10 Dec 2025)
GOLD (Safe RL)	Cost/constrained policy	Reward-shaped offline-to-online KD	Student matches/wins over teacher RL	(Li et al., 2023)

Safety-constrained distillation, in both its practical and formal instantiations, is an active research area with direct implications for secure, efficient, and robust deployment of neural models in safety- and ethics-critical environments.