Distilling Safe LLM Systems via Soft Prompts for On Device Settings

Motasem Alfarra, Cristina Pinneri, Dana Kianfar, Mohammed Almousa, Christos Louizos
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:1-18, 2026.

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 \textbf{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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-alfarra26a, title = {Distilling Safe {LLM} Systems via Soft Prompts for On Device Settings}, author = {Alfarra, Motasem and Pinneri, Cristina and Kianfar, Dana and Almousa, Mohammed and Louizos, Christos}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {1--18}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/alfarra26a/alfarra26a.pdf}, url = {https://proceedings.mlr.press/v337/alfarra26a.html}, 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 \textbf{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.} }
Endnote
%0 Conference Paper %T Distilling Safe LLM Systems via Soft Prompts for On Device Settings %A Motasem Alfarra %A Cristina Pinneri %A Dana Kianfar %A Mohammed Almousa %A Christos Louizos %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-alfarra26a %I PMLR %P 1--18 %U https://proceedings.mlr.press/v337/alfarra26a.html %V 337 %X 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 \textbf{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.
APA
Alfarra, M., Pinneri, C., Kianfar, D., Almousa, M. & Louizos, C.. (2026). Distilling Safe LLM Systems via Soft Prompts for On Device Settings. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:1-18 Available from https://proceedings.mlr.press/v337/alfarra26a.html.

Related Material