Distilling Safe LLM Systems via Soft Prompts for On Device Settings

2026-06-08Machine Learning

Machine Learning
AI summary

The authors study how to make large language models safer while still being able to run them on small devices with limited memory and processing power. They test different methods and find that using 'soft prompts' combined with a technique called distillation works best to teach the model to behave safely. This approach uses less memory and computing resources compared to other methods, making it more practical for on-device use. Their experiments show that this method balances safety and usefulness better than alternatives.

large language modelssoft promptsdistillationsafety alignmentparameter-efficient fine-tuningguard modelsLoRA adaptersKL divergencetotal variationedge devices
Authors
Motasem Alfarra, Cristina Pinneri, Dana Kianfar, Mohammed Almousa, Christos Louizos
Abstract
Deploying safe large language models (LLMs) on resource-constrained edge devices presents a critical challenge: while dual-model systems combining LLMs with guard models provide effective safety guarantees, their substantial memory and computational demands make them prohibitively expensive for on-device deployment. This paper presents a comprehensive study of parameter-efficient safety alignment methods for resource-constrained settings. Through systematic evaluation across multiple LLM architectures, training objectives, and parameter-efficient fine-tuning approaches, we identify that soft prompts combined with distillation-based training consistently outperform alternative methods. We introduce distillation frameworks based on total variation and KL divergence that effectively transfer safety behaviors from guard models into learned soft prompts. Our evaluations on various benchmarks demonstrate that this combination achieves superior safety-usefulness trade-offs compared to LoRA adapters, steering vectors, and direct optimization methods, while requiring minimal additional memory and compute at inference time. These findings establish soft prompt distillation as the preferred approach for safety alignment in on-device LLM deployment.