🤖 machine learning

Distilling Safe LLM Systems via Soft Prompts for On Device Settings

This paper proposes a parameter-efficient safety alignment method for on-device large language models that uses soft prompts distilled from guard models via total variation and KL divergence, demonstrating superior safety-usefulness trade-offs and minimal resource overhead compared to existing techniques like LoRA and steering vectors.

Original authors: Motasem Alfarra, Cristina Pinneri, Dana Kianfar, Mohammed Almousa, Christos Louizos

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Motasem Alfarra, Cristina Pinneri, Dana Kianfar, Mohammed Almousa, Christos Louizos

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: The "Double-Check" Bottleneck

Imagine you have a very talented but sometimes reckless artist (the Large Language Model or LLM) who can write stories, answer questions, and solve problems. However, this artist sometimes gets carried away and draws something inappropriate or dangerous.

To keep things safe, you hire a strict security guard (the Guard Model) to stand next to the artist. Every time the artist finishes a sentence, the guard checks it.

If it's safe, the guard says, "Go ahead, show it to the user."
If it's unsafe, the guard says, "Stop! Here is a polite refusal message instead."

The Catch: This "Double-Check" system works great for safety, but it's incredibly slow and expensive. It's like asking the artist to paint a picture, then stopping to wait for the guard to inspect every single brushstroke before moving to the next one. On a powerful computer, this is fine. But on a smartphone (which has limited battery and memory), this "Double-Check" is too heavy. It drains the battery, uses too much memory, and makes the phone feel sluggish.

The Solution: Teaching the Artist to Be Their Own Guard

The researchers asked a simple question: Can we teach the artist to be safe on their own, so we don't need the separate security guard standing next to them?

They didn't want to retrain the whole artist from scratch (which is like sending the artist back to art school for years). Instead, they used a technique called Distillation via Soft Prompts.

The Analogy: The "Magic Headband"

Think of the Soft Prompt as a special, invisible headband the artist wears.

This headband isn't made of heavy metal (like changing the artist's entire brain).
It's a tiny, lightweight accessory (just a few thousand parameters) that sits right at the start of the artist's thought process.
When the artist puts on this headband, it subtly shifts their mindset. It whispers, "Remember, don't draw anything bad," before they even start writing.

The researchers "trained" this headband by showing it thousands of examples of what the Security Guard would have done. The headband learned to mimic the guard's behavior so perfectly that the artist now refuses bad requests automatically, without needing the guard to check the work afterward.

How They Did It: The "Total Variation" Recipe

The paper compares different ways to train this "Magic Headband." They tried several recipes:

Just guessing what the next word should be (Perplexity): This made the artist write better, but didn't really stop them from drawing bad things.
Reinforcement Learning (REINFORCE): This was like giving the artist a treat when they were safe. It worked okay, but the artist sometimes forgot the rules when faced with tricky new questions.
The Winner (TV-DiSP): The researchers developed a specific mathematical recipe called Total Variation Distillation.
- Imagine this as a strict teacher who says: "Your output must match the Guard's decision exactly."
- If the Guard says "Safe," the artist must say the exact same thing.
- If the Guard says "Unsafe," the artist must immediately switch to the refusal message.
- This method (TV-DiSP) was the most effective at copying the guard's behavior without breaking the artist's ability to be helpful.

The Results: Fast, Light, and Safe

The researchers tested this on real smartphones (using Qualcomm chips) and compared it to the old "Double-Check" system.

Safety: The "Magic Headband" (TV-DiSP) was almost as safe as having the full Security Guard standing there. It successfully blocked harmful content in tests involving jailbreaks and toxic prompts.
Speed & Memory: This is the big win.
- The old system (Artist + Guard) required two heavy calculations for every word.
- The new system (Artist + Headband) only requires one calculation.
- Memory: The headband adds less than 1% to the phone's memory usage. The old system added about 30-100%.
- Battery/Compute: The new system uses less than 10% extra computing power compared to the base artist, whereas the old system doubled the workload.

Why This Matters for Your Phone

The paper concludes that for running safe AI on your phone, you don't need a heavy, dual-model system. You just need to equip the main model with this tiny, learned "headband" (soft prompt).

It's like upgrading a car's safety features. Instead of adding a second, heavy driver to the car to watch the road (which makes the car slow and heavy), you install a smart, lightweight sensor in the dashboard that automatically steers the car away from danger. The car drives just as fast, uses less gas, but is just as safe.

In short: The paper proves that by using a specific training method (Total Variation Distillation) to teach a tiny "soft prompt" to mimic a safety guard, we can make AI safe enough for smartphones without slowing them down or draining their batteries.

Technical Summary: Distilling Safe LLM Systems via Soft Prompts for On-Device Settings

Problem Statement

Deploying Large Language Models (LLMs) on resource-constrained edge devices (e.g., smartphones) presents a critical safety challenge. While dual-model systems—combining a base LLM with a separate "guard model" to classify and filter unsafe outputs—provide robust safety guarantees, they incur prohibitive memory and computational costs. The sequential nature of this approach (generating a response, then classifying it) doubles the forward passes, significantly increasing latency (time-to-first-token) and memory footprint, rendering it unsuitable for edge deployment. Existing alternatives, such as quantization or direct fine-tuning, often compromise generalization capabilities or fail to balance safety with model usefulness effectively.

Methodology

The paper proposes TV-DiSP (Total Variation Distillation via Soft Prompts), a parameter-efficient fine-tuning (PEFT) approach that distills the safety behaviors of a dual-model system into a single LLM equipped with learnable soft prompts.

Core Framework

System Formulation: The authors define a safe LLM system $p(r|x)$ that combines a base LLM $p(y|x)$ and a guard model $p(s|x, y)$ . If the guard model deems a response $y$ unsafe ( $s=0$ ), the system outputs a predefined refusal message $y_r$ ; otherwise, it outputs $y$ .
Distillation Objective: The goal is to learn a set of continuous prompt embeddings $W$ (soft prompts) such that the distilled model $q(r|x, W)$ mimics the output distribution of the safe system $p(r|x)$ .
Optimization via Total Variation (TV):
- The authors derive a theoretical guarantee (Theorem 3.1) showing that minimizing the Total Variation distance between the safe system and the distilled model bounds the performance gap on any downstream task.
- The optimization objective (Equation 4) maximizes the log-probability of the correct output (either the base LLM's response for safe inputs or the refusal message for unsafe inputs) weighted by the guard model's safety probability.
- The training requires only a dataset of prompts; the guard model dynamically labels them as safe or unsafe during the process.
Alternative Schemes: The study also evaluates:
- KL-DiSP: Minimizing Kullback-Leibler divergence, which offers looser theoretical guarantees than TV.
- Baselines: Perplexity optimization, REINFORCE (policy gradient), and adaptations of the distillation framework to Low-Rank Adapters (LoRA) and Steering Vectors (SV).

Experimental Setup

Models: Experiments were conducted on four instruction-tuned LLMs (Llama3-1B/3B, Qwen2-1.5B, Gemma2-2B) quantized to 4-bits to simulate edge constraints.
Guard Models: LlamaGuard3-1B was used for distillation training, and LlamaGuard3-8B for evaluation.
Datasets: Training utilized Beavertails and Toxigen datasets. Evaluation included out-of-distribution benchmarks (HarmBench, Detect-JailBreak) and usefulness benchmarks (IFEval, GSM8K, MMLU).
Hardware: On-device measurements were performed on smartphones equipped with Qualcomm Snapdragon 8 Elite Gen 5 chipsets.

Key Results

1. Safety-Compute Trade-off

Performance: TV-DiSP successfully distills the safety of the dual-model system into a single model. For instance, on the Llama3-1B model, safety (measured by Safety Guard Score, SGS) improved from 71% (base) to 99% (dual-system), while TV-DiSP achieved comparable safety gains with minimal overhead.
Efficiency: Compared to the dual-model system, which doubles memory and compute requirements, TV-DiSP introduces less than 1% additional memory overhead and less than 10% additional compute overhead (FLOPs per token).
Generalization: The method demonstrates strong robustness on out-of-distribution adversarial prompts (HarmBench, Detect-JailBreak) despite being trained on non-adversarial datasets.

2. Comparison with PEFT Methods

Soft Prompts vs. LoRA/SV: TV-DiSP (soft prompts) consistently outperformed TV-DiLoRA and TV-DiSV (Steering Vectors) across various learning rates. Steering vectors showed mixed results, sometimes degrading performance on larger models, while LoRA was deemed too intrusive for edge settings.
Optimization Objectives:
- TV vs. KL: Both TV and KL distillation achieved high safety scores. However, TV-DiSP offered a superior balance between safety and usefulness.
- Perplexity & REINFORCE: Perplexity optimization led to overfitting and significant drops in usefulness. REINFORCE maintained usefulness but failed to provide substantial safety improvements.

3. Usefulness Preservation

On-device evaluations (IFEval and GSM8K) showed that TV-DiSP and KL-DiSP maintain high utility.
Over-refusal: A critical finding was that KL-DiSP exhibited a high over-refusal rate (84.6% on safe prompts), incorrectly rejecting benign queries. In contrast, TV-DiSP maintained an over-refusal rate (36.1%) comparable to the dual-model system, preserving the model's ability to answer safe questions.

Significance and Claims

The paper establishes soft prompt distillation as the preferred approach for safety alignment in on-device LLM deployment. Its primary contributions are:

Superior Trade-offs: It demonstrates that soft prompts trained via distillation consistently outperform LoRA, steering vectors, and alternative training objectives in balancing safety, usefulness, and computational efficiency.
Theoretical Guarantees: The introduction of Total Variation-based distillation provides probabilistic guarantees on downstream performance, ensuring the distilled model does not deviate significantly from the safe system's behavior.
Practical Viability: Validated across multiple architectures and on actual edge hardware, the method achieves safety levels comparable to dual-model systems with negligible resource overhead (<1% memory, <10% compute), making it a practical solution for deploying safe LLMs on resource-constrained devices.

The authors conclude that while no method offers complete immunity to adversarial attacks, TV-DiSP offers a robust, efficient, and generalizable defense mechanism that significantly reduces the attack surface and computational burden compared to traditional guard-model architectures.