EASE: Practical and Efficient Safety Alignment for Small Language Models

TL;DR

This work tackles the challenge of safely deploying small language models on edge devices by addressing the limitations of shallow refusal and the overhead of full deliberative safety reasoning. It introduces EASE, a two-phase framework that first implants safety reasoning through knowledge distillation from a capable safety teacher and then calibrates a safety reasoning boundary to activate reasoning selectively for adversarial jailbreaks. The approach demonstrates improved safety robustness against sophisticated jailbreaks while preserving efficiency and general-task performance, outperforming both traditional refusal-based methods and full-depth deliberative alignment on multiple benchmarks. Practically, EASE enables robust, real-time safe interactions in resource-constrained environments, offering a scalable path for safe deployment of small models at the edge.

Abstract

Small language models (SLMs) are increasingly deployed on edge devices, making their safety alignment crucial yet challenging. Current shallow alignment methods that rely on direct refusal of malicious queries fail to provide robust protection, particularly against adversarial jailbreaks. While deliberative safety reasoning alignment offers deeper alignment for defending against sophisticated attacks, effectively implanting such reasoning capability in SLMs with limited capabilities remains an open challenge. Moreover, safety reasoning incurs significant computational overhead as models apply reasoning to nearly all queries, making it impractical for resource-constrained edge deployment scenarios that demand rapid responses. We propose EASE, a novel framework that enables practical and Efficient safety Alignment for Small languagE models. Our approach first identifies the optimal safety reasoning teacher that can effectively distill safety reasoning capabilities to SLMs. We then align models to selectively activate safety reasoning for dangerous adversarial jailbreak queries while providing direct responses to straightforward malicious queries and general helpful tasks. This selective mechanism enables small models to maintain robust safety guarantees against sophisticated attacks while preserving computational efficiency for benign interactions. Experimental results demonstrate that EASE reduces jailbreak attack success rates by up to 17% compared to shallow alignment methods while reducing inference overhead by up to 90% compared to deliberative safety reasoning alignment, making it practical for SLMs real-world edge deployments.

Figures (5)

• Figure 1: The workflow of our safety alignment method for small language models (EASE).
• Figure 2: LRMs outperform LLMs as safety reasoning teachers, with smaller LRMs showing better distillation effectiveness due to reduced capability gaps with SLMs.
• Figure 3: Model $M_\text{s}$ intermediate layer activations reveal semantic clustering by query type through t-SNE visualization. Red clusters correspond to multiple vulnerable jailbreak tactics semantic regions, blue points represent refused adversarial queries, and yellow points show benign queries. Post-calibration reasoning rates demonstrate adaptive safety reasoning activation across different semantic regions, confirming effective boundary calibration.
• Figure 4: EASE adaptive reasoning rates across datasets showing selective activation based on threat sophistication. General tasks including MMLU, HellaSwag, and GSM8K.
• Figure 5: Ablation study demonstrating EASE's optimal balance between safety (ASR on WildJailbreak) and efficiency (average tokens on HellaSwag).