Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SafeRoute: Adaptive Model Selection for Efficient and Accurate Safety Guardrails in Large Language Models

Seanie Lee, Dong Bok Lee, Dominik Wagner, Minki Kang, Haebin Seong, Tobias Bocklet, Juho Lee, Sung Ju Hwang

TL;DR

SafeRoute introduces an adaptive binary router that selectively routes prompt–response pairs to a larger safety guard model only when hard examples warrant it, while handling easy cases with a smaller guard. The approach is trained using augmented datasets that label instances where the large model adds value, and it uses a last-token feature from the small guard plus a Bayesian neural network router to decide when to deploy the large guard. Theoretical risk guarantees relate the adaptive system to an oracle with a tunable bound, and empirical results across six benchmarks show significant improvements in the safety-F1 vs. latency trade-off compared to baselines, with robust performance on both in-distribution and out-of-distribution data. The work demonstrates that adaptive model selection can meaningfully reduce compute costs without sacrificing safety performance in real-world LLM deployments.

Abstract

Deploying large language models (LLMs) in real-world applications requires robust safety guard models to detect and block harmful user prompts. While large safety guard models achieve strong performance, their computational cost is substantial. To mitigate this, smaller distilled models are used, but they often underperform on "hard" examples where the larger model provides accurate predictions. We observe that many inputs can be reliably handled by the smaller model, while only a small fraction require the larger model's capacity. Motivated by this, we propose SafeRoute, a binary router that distinguishes hard examples from easy ones. Our method selectively applies the larger safety guard model to the data that the router considers hard, improving efficiency while maintaining accuracy compared to solely using the larger safety guard model. Experimental results on multiple benchmark datasets demonstrate that our adaptive model selection significantly enhances the trade-off between computational cost and safety performance, outperforming relevant baselines.

SafeRoute: Adaptive Model Selection for Efficient and Accurate Safety Guardrails in Large Language Models

TL;DR

SafeRoute introduces an adaptive binary router that selectively routes prompt–response pairs to a larger safety guard model only when hard examples warrant it, while handling easy cases with a smaller guard. The approach is trained using augmented datasets that label instances where the large model adds value, and it uses a last-token feature from the small guard plus a Bayesian neural network router to decide when to deploy the large guard. Theoretical risk guarantees relate the adaptive system to an oracle with a tunable bound, and empirical results across six benchmarks show significant improvements in the safety-F1 vs. latency trade-off compared to baselines, with robust performance on both in-distribution and out-of-distribution data. The work demonstrates that adaptive model selection can meaningfully reduce compute costs without sacrificing safety performance in real-world LLM deployments.

Abstract

Deploying large language models (LLMs) in real-world applications requires robust safety guard models to detect and block harmful user prompts. While large safety guard models achieve strong performance, their computational cost is substantial. To mitigate this, smaller distilled models are used, but they often underperform on "hard" examples where the larger model provides accurate predictions. We observe that many inputs can be reliably handled by the smaller model, while only a small fraction require the larger model's capacity. Motivated by this, we propose SafeRoute, a binary router that distinguishes hard examples from easy ones. Our method selectively applies the larger safety guard model to the data that the router considers hard, improving efficiency while maintaining accuracy compared to solely using the larger safety guard model. Experimental results on multiple benchmark datasets demonstrate that our adaptive model selection significantly enhances the trade-off between computational cost and safety performance, outperforming relevant baselines.

Paper Structure

This paper contains 31 sections, 1 theorem, 12 equations, 12 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Safety guard models.
  4. Efficiency.
  5. Method
  6. Preliminaries
  7. SafeRoute: Adaptive Model Selection
  8. Observation.
  9. Dataset creation and training.
  10. Data augmentation.
  11. Parameterization.
  12. Inference.
  13. Theoretical analysis
  14. Experiments
  15. Experimental Setups
  16. ...and 16 more sections

Key Result

Theorem 3.1

Assuming that $\mathbb{E}[\lvert \ell(p({\mathbf{x}},{\mathbf{y}}),c)-\ell(q({\mathbf{x}},{\mathbf{y}}),c)\rvert^2]$ is bounded, we can bound the risk of our adaptive model as follows: where $M=\sqrt{\mathbb{E}[\lvert \ell(p({\mathbf{x}},{\mathbf{y}}),c)-\ell(q({\mathbf{x}},{\mathbf{y}}),c)\rvert^2]}$.

Figures (12)

  • Figure 1: Our proposed safety guard router, SafeRoute, distinguishes hard examples from easy ones. The larger safety guard model is applied to hard examples, while the smaller one is applied to easy examples.
  • Figure 2: Latency ($\downarrow$) vs. safety F1 score ($\uparrow$) trade-off when using the smaller (Llama-Guard-3-1B) and larger (Llama-Guard-3-8B) models. See \ref{['fig:flops']} and \ref{['fig:large_ratio']} in \ref{['sec:additional_experimental_results']} for FLOPs and ratio of large model trade-off.
  • Figure 3: Latency ($\downarrow$) vs. safety F1 score ($\uparrow$) trade-off when using the smaller (Llama-Guard-3-1B) and larger (Granite-Guardian-3-8B) models. See \ref{['fig:flops_guardian']} and \ref{['fig:large_ratio_guardian']} in \ref{['sec:additional_experimental_results']} for FLOPs and ratio of large model trade-off.
  • Figure 4: Ablation studies on (a): pooling methods, (b): feature layers, and (c): the number of paraphrases.
  • Figure 5: The number of the large model selections for each jailbreak attack in HarmBench dataset.
  • ...and 7 more figures

Theorems & Definitions (2)

  • Theorem 3.1
  • proof