HiRoute: Hierarchical Routed Prompt Tuning for Safety Alignment of Large Language Models

2026-08-13

AI Summary

This work addresses the vulnerability of large language models to harmful requests and jailbreak attacks, noting that existing safety alignment methods often struggle to simultaneously ensure robustness across diverse risk categories, maintain response helpfulness, and avoid over-rejecting benign inputs. To overcome these limitations, the authors propose HiRoute, a novel framework that decouples general safety control from risk-specific response guidance through a lightweight hierarchical router and a shared-expert prompting architecture. The router—trained with frozen model representations—dynamically routes inputs by leveraging both coarse-grained shared prompts and fine-grained expert prompts, optimized via preference learning and alternating gradient updates. Experimental results demonstrate that HiRoute significantly enhances safety across multiple benchmarks, substantially reduces over-refusal, and preserves both response helpfulness and general task performance.

Abstract

Large language models (LLMs) remain vulnerable to harmful requests and jailbreak attacks. Parameter-efficient safety alignment methods based on prompt tuning typically rely on a single global prompt or externally selected prompt modules. Such static designs struggle to maintain a cross-category safety boundary while generating constructive responses tailored to specific risks and avoiding over-refusal of benign inputs. To address these limitations, we propose HiRoute, an input-adaptive hierarchical prompt-tuning framework that separates category-agnostic safety control from category-specific response guidance. HiRoute first trains a lightweight hierarchical router on representations extracted from a frozen LLM to jointly detect harmful intent and predict multi-label risk scores. It then freezes both the backbone model and the router and uses preference optimization with alternating gradient updates to learn a shared coarse-grained prompt and a set of fine-grained prompt experts as continuous embeddings. At inference time, benign inputs bypass the safety branch, whereas risky inputs are processed using the shared prompt together with a router-weighted mixture of risk-specific prompt experts. Experiments across three instruction-tuned models show that HiRoute achieves high safety rates across multiple safety benchmarks while preserving safe-response helpfulness, reducing over-refusal, and maintaining competitive performance on general-purpose tasks.

Problem

• safety alignment
• large language models
• prompt tuning
• harmful requests
• over-refusal

Innovation

• hierarchical prompt tuning
• input-adaptive routing
• safety alignment
• prompt experts
• over-refusal reduction