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

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

Abstract Summary

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, the paper proposes 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.

Introduction Summary

Large language models have demonstrated strong capabilities in instruction following, reasoning, and open-domain dialogue, yet malicious requests, borderline queries, and jailbreak attacks can still induce them to generate harmful content. Supervised fine-tuning (SFT) and reinforcement learning from human feedback (RLHF) can encode safe behaviors into model parameters, but incur substantial training and storage costs. Uniformly strengthening refusal behavior may also adversely affect benign inputs, leading to capability degradation or over-refusal. The paper raises the practical question: can we freeze the backbone model, train only a small number of parameters, and adapt safety alignment to the risk characteristics of different inputs?

Prompt tuning provides a parameter-efficient approach by freezing the backbone model and optimizing only input-side prompt parameters. However, existing safety prompts typically rely on static safety-alignment mechanisms. Relying solely on a single coarse-grained safety prompt that does not distinguish among risk categories applies the same category-agnostic safety constraint to all harmful inputs, often resulting in generic refusals that fail to accommodate the response requirements of different risks. In contrast, fine-grained safety alignment assigns category-specific prompts to different risk categories to generate category-relevant safe responses. Although existing modular methods can learn multiple prompts, they typically depend on manual selection or fixed composition and therefore struggle to accommodate the semantic diversity of real-world requests.

The paper notes that recent studies have recognized that safety alignment should go beyond refusing harmful requests and should also provide risk-specific explanations, compliant guidance, and safe alternatives. GPT-5's safe-completions paradigm shifts safety training from binary comply-or-refuse decisions toward output-centric control and seeks to maximize response helpfulness subject to safety-policy constraints. Oyster-I emphasizes the generation of constructive safe responses, whereas PKU-SafeRLHF introduces fine-grained risk categories and provides separate annotations of response safety and helpfulness. However, existing prompt-based methods have yet to unify fine-grained risk identification, category-specific prompt composition, and category-agnostic safety constraints.

Based on these observations, the paper proposes HiRoute, an input-adaptive hierarchical prompt-tuning framework. HiRoute uses a shared coarse-grained prompt to establish a category-agnostic safety boundary and employs a hierarchical router to form a weighted combination of fine-grained prompt experts, thereby providing risk-specific safety guidance. Training proceeds in two stages to separate risk identification from behavior optimization. First, the router is trained over representations produced by a frozen model to determine whether an input is harmful and predict multi-label risk scores. The backbone model and router are then frozen, and only the two types of prompts are optimized. At inference time, benign inputs bypass safety prompting, whereas hierarchical prompt combinations are dynamically constructed for risky inputs according to the routing results.

The main contributions are: (1) empirically validating the complementary limitations of two prompt-based safety-alignment approaches—a shared coarse-grained prompt provides stronger safety but produces less constructive responses, whereas routed fine-grained prompt mixtures improve safe-response helpfulness but yield smaller safety gains; (2) proposing HiRoute, which establishes a cross-category safety boundary with a shared coarse-grained prompt, dynamically composes fine-grained prompt experts through a hierarchical router, and applies safety prompts based on input risk to avoid unnecessary intervention on benign inputs; (3) validating the effectiveness and robustness of HiRoute across multiple instruction-tuned models and benchmarks for safety, general utility, and over-refusal.

Related Work Summary

LLM safety alignment aims to reduce the risk of models assisting harmful intentions or generating policy-violating content. Representative approaches include supervised fine-tuning, reinforcement learning from human feedback, Constitutional AI, and direct preference optimization. These methods incorporate safety behaviors into model parameters using human feedback, AI-generated feedback, or preference data, substantially improving the models' ability to refuse harmful requests. However, parameter-level safety alignment still faces an inherent trade-off between safety and helpfulness: overly restrictive alignment may exacerbate over-refusal, whereas insufficient constraints may fail to defend against sophisticated malicious inputs. Moreover, learned safety behaviors may overfit to the training distribution and generalize poorly to jailbreak attacks, multi-risk requests, or unseen risk categories.

Prompt tuning optimizes only the input-side discrete or continuous prompt parameters while keeping the model parameters frozen, offering low training overhead and ease of deployment. Early approaches, including Prefix-Tuning and P-Tuning v2, demonstrated that continuous prompts can achieve performance comparable to full-parameter fine-tuning across various tasks. More recent studies have extended prompt optimization to safety alignment and jailbreak defense. For example, one work optimizes defensive suffixes, another learns continuous safety prompts, another employs contrastive safety prompts, another improves model safety by distilling guard-model behaviors into prompts, and another decomposes category-specific safety constraints into multiple learnable control tokens. While these methods demonstrate the potential of prompt tuning for safety alignment, most still rely on global or externally specified safety-control signals and therefore struggle to adapt prompt compositions to the risk structure of individual inputs. In contrast, HiRoute introduces a hierarchical risk router that first identifies coarse-grained harmful intent and then predicts a fine-grained risk distribution, using this distribution to dynamically compose category-specific prompts with a shared coarse prompt, enabling input-adaptive safety prompt tuning.

Why Hierarchical Safety Prompting Summary

Before constructing the complete framework, the paper considered two natural prompt-tuning-based approaches to safety alignment. The first approach applies a single coarse-grained safety prompt, without distinguishing among risk categories, to all risky inputs. The second trains a fine-grained risk router and forms a weighted mixture of category-specific prompts based on its outputs, thereby enabling input-dependent safety control. The paper uses the base model without additional safety tuning as the reference and evaluates both safety rate and safe-response helpfulness on external risk data.

The results reveal a clear safety–helpfulness trade-off between the two approaches. The single coarse-grained prompt establishes a strong safety constraint, achieving safety rates of 91.5% and 93.0% on Mistral and Zephyr, respectively, substantially outperforming the corresponding base models at 54.0% and 62.7%. However, these safety gains come at the cost of safe-response helpfulness. Its helpfulness scores are only 5.3 and 6.1, lower than the 7.4 and 7.2 achieved by the routed fine-grained prompt mixture. This result suggests that a globally shared prompt tends to compress diverse risks into similar conservative refusal patterns. Although it can reliably prevent unsafe responses, it struggles to provide risk-specific explanations, compliant guidance, and safe alternatives.

The routed fine-grained prompt mixture exhibits the opposite pattern. It achieves helpfulness scores of 7.4 and 7.2 on Mistral and Zephyr, indicating that category-specific prompts produce more targeted safe responses. However, its safety rates are only 71.0% and 72.5%, both 20.5 percentage points below those of the corresponding coarse-grained prompts. One possible explanation is that external risk requests do not always align precisely with the predefined expert boundaries, which may produce more dispersed routing weights and weaken the safety constraint imposed by the resulting prompt mixture. More importantly, the fine-grained approach relies entirely on category-specific experts for safety control. When expert matching is insufficient, the system lacks a category-agnostic safety constraint as a fallback.

These results demonstrate that coarse- and fine-grained prompts are complementary. The former provides a category-agnostic safety boundary, whereas the latter offers risk-specific explanations, guidance, and safe alternatives. Motivated by this observation, HiRoute hierarchically combines a shared coarse-grained prompt with routed fine-grained experts and employs safety gating to avoid imposing unnecessary control on benign inputs.

Methodology Summary

HiRoute consists of three stages: risk recognition, safety control, and adaptive inference. First, the hierarchical router determines whether an input is harmful from representations produced by the frozen language model and estimates a multi-label risk distribution. Second, the system activates the safety branch only for risky inputs: a shared coarse-grained prompt provides cross-category constraints, while fine-grained prompt experts are combined according to the predicted risk distribution. Finally, the frozen backbone model generates a response conditioned on the composed prompts. Training proceeds in two stages: router learning and safety-prompt optimization. In the second stage, neither the backbone model nor the router is updated, and only the hierarchical prompt parameters are optimized.

Hierarchical Risk Routing and Prompt Composition

Let x denote the input token sequence and πθ the frozen instruction-tuned language model. HiRoute learns a shared coarse-grained prompt Pc ∈ RLc×d and a set of K fine-grained prompt experts Pk ∈ RLf×d Kk=1, where Lc and Lf denote the corresponding prompt lengths and d is the embedding dimension. Let Eθ(x) ∈ RT×d denote the embedding sequence of x. HiRoute defines the hierarchical safety prompt and its corresponding prompted context as follows: P(x) = [Pf(x); Pc], CP(x) = [Eθ(x); P(x)]. Here, [;] denotes concatenation along the sequence dimension. P(x) is the input-dependent hierarchical safety prompt, whereas CP(x) is the complete embedding sequence passed to the frozen language model. The shared prompt Pc provides a category-agnostic safety boundary, while Pf(x) injects the category-specific control required by the current input.

The hierarchical router Rϕ takes the frozen hidden states of πθ as input. A lightweight Transformer encoder followed by masked mean pooling produces an input-level representation hx. The coarse-grained prediction head outputs pc(x) = [psafe(x), prisk(x)]. Its two dimensions follow the fixed class order 0 = safe and 1 = risk. The fine-grained head outputs pf(x) ∈ [0,1]K and applies an independent sigmoid to each category, allowing a single request to be assigned to multiple risk categories.

For each fine-grained risk category k, an independent soft-prompt expert Pk is optimized. The fine-grained risk scores produced by the router are normalized into composition weights αk(x) and used to form a weighted combination of the experts: Pf(x) = ΣKk=1 αk(x)Pk, where αk(x) = pf,k(x) / ΣKj=1 pf,j(x), with ΣKk=1 αk(x) = 1. Because all Pk share the same shape, they can be combined element-wise.

Decoupled Two-Stage Optimization

Training proceeds in two stages to prevent the generation objective from altering the risk decision boundary.

Stage I: Hierarchical Risk Routing. Given the coarse-grained dataset Dc = (xi, yic), where yic ∈ 0,1 denotes the safe and risk label, the coarse-grained prediction head is trained using cross-entropy loss. For the fine-grained dataset Df = (xi, yif), where yif ∈ 0,1 K is a multi-label risk vector, the fine-grained prediction head is trained using binary cross-entropy loss. The router objective is defined as: Lrouter = Lc + Lf, where Lc = EDc[CE(pc(x), yc)] and Lf = EDf[BCE(pf(x), yf)]. The paper first optimizes only Lc to learn a general safe/risk decision boundary and then jointly optimizes Lc and Lf. During joint training, Dc supervises only the coarse-grained head, whereas Df supervises only the fine-grained head, preventing the coarse decision boundary from overfitting the limited set of annotated risk categories.

Stage II: Risk-Adaptive Hierarchical Prompt Learning. The backbone πθ and the trained router Rϕ are frozen, and only the shared coarse-grained prompt Pc and the fine-grained prompt experts Pk Kk=1 are updated. Given a preference triplet (x, yw, yl), where yw is a safe and helpful response and yl is an unsafe or low-quality response, Direct Preference Optimization is applied to increase the relative conditional likelihood of yw over yl. When updating the two types of prompts, gradients are alternately masked to reduce interference between them. When updating Pc, the routed fine-grained prompt Pf(x) remains in the forward pass but receives no gradient. Conversely, when updating the fine-grained experts, Pc remains in the forward pass but is held fixed. Formally, Pe(c)(x) = [sg(Pf(x)); Pc] and Pe(f)(x) = [Pf(x); sg(Pc)], where the stop-gradient operator sg(·) acts as the identity during the forward pass but blocks gradient propagation during backpropagation.

Input-Adaptive Inference

At inference time, the router first computes the safe-input probability psafe(x). If psafe(x) exceeds the threshold τs, the system bypasses all safety prompts and generates directly from the frozen backbone. This gate limits the influence of safety prompting on benign inputs. Otherwise, the input enters the risk-control branch. The system constructs Pf(x) from the fine-grained risk scores pf(x) and combines it with the shared prompt Pc to form P(x). The resulting prompted context CP(x) is then passed to the frozen language model for response generation: y ∼ πθ(yx) if psafe(x) > τs, otherwise y ∼ πθ(yCP(x)). A larger τs routes more inputs through the safety branch, which generally improves safety but may reduce general utility. Unlike a fixed safety prompt, HiRoute automatically determines both whether to activate safety control and how to compose that control for each input.

Experimental Setting Summary

Datasets and Models. Router training and prompt optimization use separate sources of supervision. A coarse-grained binary classific…

Based on the paper, here are the specific improvements you can make to AI systems and what the improved system can do:

1. Input-Adaptive Safety Gating

• Improvement: Implement a hierarchical router that classifies inputs as safe or risky before applying any safety prompts. Benign inputs bypass all safety mechanisms, while risky inputs trigger a safety branch.

• Improved System Capability: The AI system can maintain high general utility and low over-refusal rates (e.g., 2.5% on XSTest) while achieving high safety rates (93.2% on average), because it does not apply conservative refusal behavior to harmless queries.

2. Hierarchical Prompt Composition for Risk-Specific Responses

• Improvement: Use a shared coarse-grained prompt (for category-agnostic safety boundary) combined with a weighted mixture of fine-grained prompt experts (for category-specific guidance). The weights are dynamically computed from multi-label risk scores predicted by the router.

• Improved System Capability: The AI system can generate constructive, safe responses that explain risks, provide compliant guidance, and offer safe alternatives—rather than generic refusals—while still maintaining a strong safety floor. This improves safe-response helpfulness scores (e.g., 7.4 on Mistral) without sacrificing safety.

3. Decoupled Two-Stage Training to Prevent Interference

• Improvement: Train the router separately from the prompts, then freeze both and optimize prompts using preference optimization with alternating gradient masking (stop-gradient on the other prompt type during each update).

• Improved System Capability: The AI system can learn a stable risk decision boundary without being distorted by generation objectives, and it can coordinate coarse and fine prompts effectively. This leads to consistent safety improvements (e.g., from 74.8% to 93.2% average safety) across different backbones.

4. Multi-Label Risk Recognition for Compound Threats

• Improvement: Use a fine-grained head that outputs independent sigmoid probabilities for multiple risk categories (e.g., cybercrime, economic crime, privacy violations, violence), allowing a single input to be assigned to multiple risks.

• Improved System Capability: The AI system can handle compound-risk requests (e.g., a prompt that involves both privacy violation and violence) by composing multiple expert prompts, leading to more accurate and contextually appropriate safety responses.

5. Robustness Against Transfer-Based Jailbreak Attacks

• Improvement: Rely on a shared coarse-grained prompt as a stable defensive foundation, independent of precise expert matching, and combine it with routed experts.

• Improved System Capability: The AI system can defend against transfer-based GCG attacks, improving average safety from 34.6% to 88.3% (a 53.6 percentage-point increase) across different backbones, even when the base model is highly vulnerable.

6. Threshold-Tunable Safety–Utility Trade-off

• Improvement: Expose an inference-time threshold (τs) that controls how many inputs are routed to the safety branch. Higher thresholds increase safety but slightly reduce general utility.

• Improved System Capability: The AI system can be calibrated for different deployment contexts—e.g., a stricter threshold (0.95) for high-risk environments (safety rate 95.0% on JailbreakBench) with minimal utility loss (GSM8K drops only 2.5 points), or a looser threshold for general-purpose use.

7. Alternating Gradient Masking for Prompt Coordination

• Improvement: During prompt training, alternate updates between coarse and fine prompts while keeping the other fixed in the forward pass (using stop-gradient). This ensures both prompt types are optimized in the final composed context.

• Improved System Capability: The AI system avoids the degradation seen in sequential training (where fine prompts fail to adapt to the final context), achieving a 18.4 percentage-point improvement in average safety (from 74.8% to 93.2%) compared to coarse-after-fine training.

8. Length and Ratio Optimization for Prompt Capacity

• Improvement: Use a 15/5 split between coarse and fine prompt lengths (total 20), and avoid overly long prompts (which degrade safety) or overly short fine prompts (which reduce helpfulness).

• Improved System Capability: The AI system can balance safety (93.2% average) and helpfulness (7.4 average) without overfitting to prompt length, and it can be scaled to larger models (e.g., Vicuna-13B) with consistent gains.

9. General Utility Preservation

• Improvement: Use safety gating to ensure benign inputs are processed by the frozen backbone without any prompt modification, and train prompts only on risky examples.

• Improved System Capability: The AI system retains competitive performance on general tasks (GSM8K, MT-Bench, TruthfulQA) while improving safety, as demonstrated by minimal utility degradation (e.g., GSM8K from 50.5% to 48.0% at threshold 0.95) and no increase in over-refusal.

10. Open-Set Risk Adaptation via Router Generalization

• Improvement: Train the coarse-grained router on a broad binary dataset (e.g., WildGuardMix) to learn a general safe/risk boundary, independent of the limited fine-grained categories.

• Improved System Capability: The AI system can recognize and gate unseen or out-of-distribution harmful inputs (even if fine-grained expert matching is imperfect), because the coarse boundary provides a fallback safety constraint, reducing reliance on precise category prediction.

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.