Technical Summary: HiRoute

Problem Statement

Large Language Models (LLMs) remain vulnerable to harmful requests, jailbreak attacks, and multi-risk queries. While Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF) can encode safety behaviors, they incur high training and storage costs and often lead to a trade-off between safety and helpfulness. Specifically, uniformly strengthening refusal behaviors can cause "over-refusal" on benign inputs, degrading general capabilities.

Existing parameter-efficient safety alignment methods, such as prompt tuning, typically rely on static mechanisms. These approaches generally fall into two categories with complementary limitations:

1. Coarse-grained Global Prompts: A single safety prompt applied to all risky inputs establishes a strong safety boundary but often results in generic, uninformative refusals that fail to provide risk-specific guidance or safe alternatives.
2. Fine-grained Routed Prompts: Methods that assign category-specific prompts to different risks can generate more targeted and helpful safe responses. However, they struggle to maintain a stable safety boundary, particularly when inputs do not align precisely with predefined risk categories or when expert matching is insufficient.

The core challenge is how to balance fine-grained response constructiveness with cross-category safety through input-adaptive prompt composition without freezing the entire backbone model.

Methodology: HiRoute

The authors propose HiRoute (Hierarchical Routed Prompt Tuning), an input-adaptive framework that separates category-agnostic safety control from category-specific response guidance. The framework operates on a frozen instruction-tuned LLM and involves three main components:

1. Hierarchical Risk Routing

A lightweight router is trained on representations extracted from the frozen LLM to perform two tasks:

• Coarse-grained Detection: Determines if an input is harmful or benign ($p_{safe}$ vs. $p_{risk}$).
• Fine-grained Risk Prediction: Predicts multi-label risk scores for specific categories (e.g., cybercrime, privacy violations, economic crime, violence).
  The router outputs a probability distribution used to dynamically weight fine-grained prompt experts.

2. Hierarchical Prompt Composition

HiRoute learns two types of continuous prompt embeddings:

• Shared Coarse-grained Prompt ($P_c$): Provides a category-agnostic safety boundary applicable to all risky inputs.
• Fine-grained Prompt Experts ($\{P_k\}$): A set of $K$ experts, each corresponding to a specific risk category.
  For a given input $x$, the fine-grained prompt $P_f(x)$ is constructed as a weighted sum of the experts based on the router's predicted risk scores:
  $$P_f(x) = \sum_{k=1}^{K} \alpha_k(x) P_k$$
  where $\alpha_k(x)$ are normalized weights. The final input-dependent prompt is $P(x) = [P_f(x); P_c]$.

3. Two-Stage Decoupled Optimization

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

• Stage I (Router Learning): The router is trained using Cross-Entropy (CE) for coarse-grained classification and Binary Cross-Entropy (BCE) for fine-grained multi-label classification. The backbone model remains frozen.
• Stage II (Prompt Optimization): The backbone and the trained router are frozen. Only the prompt parameters ($P_c$ and $\{P_k\}$) are optimized using Direct Preference Optimization (DPO) on preference triplets (safe/helpful vs. unsafe/low-quality). To reduce interference, the authors employ alternating gradient updates: when updating $P_c$, gradients for $P_f(x)$ are masked (stop-gradient), and vice versa.

4. Input-Adaptive Inference

At inference, the router first computes the safety probability $p_{safe}(x)$.

• Benign Inputs: If $p_{safe}(x) > \tau_s$ (a safety threshold), the system bypasses all safety prompts and generates directly from the frozen backbone, preserving general utility.
• Risky Inputs: If $p_{safe}(x) \leq \tau_s$, the system activates the safety branch, constructing the hierarchical prompt $P(x)$ to condition the generation.

Key Contributions

1. Empirical Validation of Trade-offs: The authors demonstrate that shared coarse-grained prompts offer stronger safety but lower helpfulness, while routed fine-grained prompt mixtures improve helpfulness but yield smaller safety gains.
2. HiRoute Framework: They propose a novel architecture that combines a shared coarse-grained prompt (for a stable safety boundary) with dynamically routed fine-grained experts (for risk-specific guidance). This design avoids unnecessary intervention on benign inputs via a safety gate.
3. Comprehensive Evaluation: The method is validated across three instruction-tuned models (Mistral-7B, Vicuna-7B, Zephyr-7B) and multiple benchmarks, showing it maintains high safety and helpfulness while preserving general utility.

Experimental Results

Experiments were conducted on three models against baselines including Base models, RPO (Robust Prompt Optimization), DRO (Dynamic Risk Optimization), and ACD (Adversarial Contrastive Decoding).

• Safety and Helpfulness: HiRoute achieved the highest average safety rates across all models (e.g., 93.2% on Mistral, 97.7% on Vicuna, 94.8% on Zephyr) while simultaneously achieving the highest or near-highest safe-response helpfulness scores. It outperformed the strongest baselines by 1.3–2.9 percentage points in safety.
• General Utility: Unlike methods that degrade general capabilities, HiRoute preserved performance on GSM8K (mathematical reasoning), MT-Bench (instruction following), and TruthfulQA. For instance, on Mistral, it increased safety from 53.2% to 93.2% while retaining 48.0% GSM8K accuracy.
• Over-Refusal: HiRoute significantly reduced over-refusal on benign inputs (XSTest). On Mistral, it achieved a 2.5% over-refusal rate, a massive improvement over RPO (40.5%) and DRO (18.0%).
• Ablation Studies:
  • Training Strategy: Joint alternating training (updating prompts in a consistent composed context) significantly outperformed sequential training ("Coarse after fine"), improving average safety from 74.8% to 93.2%.
  • Prompt Length: A 15/5 split between coarse and fine prompt lengths (total 20) offered the best balance between safety and helpfulness.
  • Threshold: Increasing the safety gating threshold $\tau_s$ strengthened safety control with minimal impact on general capability.

Significance and Claims

The paper claims that HiRoute effectively addresses the tension between safety and helpfulness in LLM alignment. By decoupling risk identification from response generation and employing a hierarchical, input-adaptive prompt structure, the method:

• Establishes a stable, category-agnostic safety boundary without relying on aggressive, uniform refusals.
• Provides constructive, risk-specific guidance and safe alternatives for harmful requests.
• Minimizes over-refusal by allowing benign inputs to bypass safety prompting entirely.
• Achieves these goals with parameter efficiency, freezing the backbone model and only optimizing a small number of prompt parameters and a lightweight router.

The authors conclude that input-dependent hierarchical prompting offers a robust solution for safety alignment that balances safety, response quality, parameter efficiency, and general utility. They acknowledge limitations regarding the router's accuracy on out-of-distribution compound risks and the current focus on single-turn text interactions, suggesting future work in open-set risk recognition and multimodal settings.