Unified Multi-Task Learning & Model Fusion for Efficient Language Model Guardrailing

TL;DR

This work tackles the challenge of efficient and generalizable guardrailing for large language models by introducing a data-centric, policy-guided approach. It combines synthetic data generation with single-task and multi-task training (TaskGuard and MultiTaskGuard) and a Bayesian model merging framework (MMS) to produce UniGuard, a state-of-the-art guardrail that surpasses strong LLM baselines while remaining compact. Key contributions include a synthetic-data pipeline that enables strong performance with sub-1GB models, a multi-policy guardrailing approach that benefits from cross-task learning, and a model merging search using Thompson sampling to optimally fuse parameters from multiple fine-tuned models. The results demonstrate significant improvements on public safety/toxicity/prompt-injection benchmarks and the DynaGuardrail suite, highlighting practical implications for cost-effective, fast guardrailing, including on-device deployment and reduced latency.

Abstract

The trend towards large language models (LLMs) for guardrailing against undesired behaviors is increasing and has shown promise for censoring user inputs. However, increased latency, memory consumption, hosting expenses and non-structured outputs can make their use prohibitive. In this work, we show that task-specific data generation can lead to fine-tuned classifiers that significantly outperform current state of the art (SoTA) while being orders of magnitude smaller. Secondly, we show that using a single model, \texttt{MultiTaskGuard}, that is pretrained on a large synthetically generated dataset with unique task instructions further improves generalization. Thirdly, our most performant models, \texttt{UniGuard}, are found using our proposed search-based model merging approach that finds an optimal set of parameters to combine single-policy models and multi-policy guardrail models. % On 7 public datasets and 4 guardrail benchmarks we created, our efficient guardrail classifiers improve over the best performing SoTA publicly available LLMs and 3$^{\text{rd}}$ party guardrail APIs in detecting unsafe and safe behaviors by an average F1 score improvement of \textbf{29.92} points over Aegis-LlamaGuard and \textbf{21.62} over \texttt{gpt-4o}, respectively. Lastly, our guardrail synthetic data generation process that uses custom task-specific guardrail poli

Figures (5)

• Figure 1: Blocking malicious prompts.
• Figure 2: Guardrailing process that includes synthetically generated datasets, single policy fine-tuned models (TaskGuard), multi-policy finetuned models (MultiTaskGuard) used for classification, model evaluation and model merging (UniGuard).
• Figure 3: Model Performance Differences of Classifier-Only vs Full Model Tuning
• Figure 4: TaskGuard & MultiTaskGuardLearning Curves on Safety and Finance test sets.
• Figure : Thompson Sampling with TIES