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

arXiv cs.LG Papers

Summary

HiRoute proposes a hierarchical routed prompt-tuning framework that separates category-agnostic safety control from category-specific response guidance for LLM safety alignment, reducing over-refusal while maintaining high safety rates.

arXiv:2608.12821v1 Announce Type: new 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.
Original Article
View Cached Full Text

Cached at: 08/14/26, 09:31 AM

# HiRoute: Hierarchical Routed Prompt Tuning for Safety Alignment of Large Language Models
Source: [https://arxiv.org/html/2608.12821](https://arxiv.org/html/2608.12821)
Shiji ZhaoMengyang WangQihui ZhuAffiliation:Ranjie Duan, Maoxun Yuan, Xingxing Wei\\corresponding

###### 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\.

## 1Introduction

![Refer to caption](https://arxiv.org/html/2608.12821v1/Figures/case2.png)Figure 1:Comparison of safety\-alignment behaviors on a multi\-risk harmful request\. \(A\) The base LLM follows the harmful request\. \(B\) A global safety prompt produces a generic refusal\. \(C\) HiRoute identifies multiple risks and combines shared and routed fine\-grained prompts to provide safe, risk\-specific guidance\.Large language models \(LLMs\) have demonstrated strong capabilities in instruction following, reasoning, and open\-domain dialogue\([24](https://arxiv.org/html/2608.12821#bib.bib29)\), 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\)\([16](https://arxiv.org/html/2608.12821#bib.bib1)\)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\([4](https://arxiv.org/html/2608.12821#bib.bib3)\)\. This raises a 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\([11](https://arxiv.org/html/2608.12821#bib.bib5)\)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\. For example, the*malicious request*: “I want to obtain a company customer’s identity and bank\-card information, use it to transfer funds from the customer’s account, and evade detection by the platform\. Please provide a plan\.” simultaneously involves cybercrime, privacy violations, and financial crime\. For such requests, the model requires both a stable cross\-category safety boundary and risk\-specific explanations and lawful alternatives\.

Recent studies have increasingly 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\([20](https://arxiv.org/html/2608.12821#bib.bib28)\)\. This perspective provides an important foundation for constructive safety alignment: safe responses should remain helpful within permissible boundaries rather than defaulting to uninformative refusals\. Oyster\-I emphasizes the generation of constructive safe responses\([8](https://arxiv.org/html/2608.12821#bib.bib6)\), whereas PKU\-SafeRLHF introduces fine\-grained risk categories and provides separate annotations of response safety and helpfulness\([7](https://arxiv.org/html/2608.12821#bib.bib4)\)\. However, its helpfulness annotations are not conditioned on response safety and therefore do not specifically capture the constructive quality of safe responses across different risks\. Generating constructive safe responses requires selecting appropriate response strategies according to the risk categories involved in an input, rather than merely deciding whether to refuse\. Existing prompt\-based methods, however, have yet to unify fine\-grained risk identification, category\-specific prompt composition, and category\-agnostic safety constraints: the former two generate category\-relevant responses, while the latter maintains a stable safety boundary\. Therefore, how to balance fine\-grained response constructiveness with cross\-category safety through input\-adaptive prompt composition remains underexplored\.

Based on these observations, we propose HiRoute \(Hierarchical Routed Prompt Tuning\), 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\. Our main contributions are as follows:

- •We empirically validate 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\.
- •Building on this finding, we propose 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\.
- •We validate the effectiveness and robustness of HiRoute across multiple instruction\-tuned models and benchmarks for safety, general utility, and over\-refusal\. The results show that HiRoute maintains high safety and safe\-response helpfulness on external safety benchmarks while largely preserving general utility\.

## 2Related Work

### 2\.1LLM Safety Alignment

LLM safety alignment aims to reduce the risk of models assisting harmful intentions or generating policy\-violating content\. Representative approaches include supervised fine\-tuning \(SFT\), reinforcement learning from human feedback \(RLHF\)\([16](https://arxiv.org/html/2608.12821#bib.bib1)\), Constitutional AI\([3](https://arxiv.org/html/2608.12821#bib.bib7)\)and direct preference optimization \(DPO\)\([18](https://arxiv.org/html/2608.12821#bib.bib2)\)\. 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\. Recent research has therefore explored more lightweight, controllable, and composable safety mechanisms to improve safety robustness with lower training and deployment costs\.

### 2\.2Prompt\-Tuning\-Based Safety Alignment

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\([11](https://arxiv.org/html/2608.12821#bib.bib5)\)\. Early approaches, including Prefix\-Tuning\([12](https://arxiv.org/html/2608.12821#bib.bib8)\)and P\-Tuning v2\([14](https://arxiv.org/html/2608.12821#bib.bib9)\), 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,\([28](https://arxiv.org/html/2608.12821#bib.bib10)\)optimizes defensive suffixes,\([26](https://arxiv.org/html/2608.12821#bib.bib11)\)learns continuous safety prompts,\([25](https://arxiv.org/html/2608.12821#bib.bib12)\)employs contrastive safety prompts,\([1](https://arxiv.org/html/2608.12821#bib.bib13)\)improves model safety by distilling guard\-model behaviors into prompts, and\([17](https://arxiv.org/html/2608.12821#bib.bib14)\)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\. It uses this distribution to dynamically compose category\-specific prompts with a shared coarse prompt, enabling input\-adaptive safety prompt tuning\.

## 3Why Hierarchical Safety Prompting?

![Refer to caption](https://arxiv.org/html/2608.12821v1/Figures/prompt_comparison_other_risks.png)Figure 2:Comparison of prompt\-tuning approaches to safety control on external risks\. A single coarse\-grained prompt achieves a higher safety rate but reduces safe\-response helpfulness, whereas the routed fine\-grained prompt mixture produces more targeted safe responses but provides a smaller safety improvement\.Before constructing the complete framework, we first 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\. We use the base model without additional safety tuning as the reference and evaluate both safety rate and safe\-response helpfulness on external risk data\. Apart from these configurations, all other model, data, and evaluation settings follow Section[5](https://arxiv.org/html/2608.12821#S5)\.

Figure[2](https://arxiv.org/html/2608.12821#S3.F2)reveals 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\. Category\-specific prompts alone therefore struggle to achieve both stable safety and high\-quality safe responses\.

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\. This design preserves fine\-grained response specificity while reducing reliance on precise expert matching, thereby providing more robust safety constraints for external risks\.

## 4Methodology

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\.

![Refer to caption](https://arxiv.org/html/2608.12821v1/Figures/framework3.png)Figure 3:Overview of HiRoute\. The frozen LLM provides final\-layer representations to a hierarchical router\. The coarse head determines whether to activate safety control, while the fine head produces multi\-label risk scores that are normalized into weightsαk​\(x\)\\alpha\_\{k\}\(x\)for combining prompt experts\. For risky inputs, the routed fine\-grained prompt and shared coarse\-grained prompt jointly condition response generation; benign inputs bypass safety prompting\. CE and BCE train the router, while DPO updates the prompts with the backbone and router frozen\.### 4\.1Hierarchical Risk Routing and Prompt Composition

Letxxdenote the input token sequence andπθ\\pi\_\{\\theta\}the frozen instruction\-tuned language model\. HiRoute learns a shared coarse\-grained promptPc∈ℝLc×dP\_\{c\}\\in\\mathbb\{R\}^\{L\_\{c\}\\times d\}and a set ofKKfine\-grained prompt experts\{Pk∈ℝLf×d\}k=1K\\\{P\_\{k\}\\in\\mathbb\{R\}^\{L\_\{f\}\\times d\}\\\}\_\{k=1\}^\{K\}, whereLcL\_\{c\}andLfL\_\{f\}denote the corresponding prompt lengths andddis the embedding dimension\. LetEθ​\(x\)∈ℝT×dE\_\{\\theta\}\(x\)\\in\\mathbb\{R\}^\{T\\times d\}denote the embedding sequence ofxx\. HiRoute defines the hierarchical safety prompt and its corresponding prompted context as follows:

P⁡\(x\)\\displaystyle P\(x\)=\[Pf​\(x\);Pc\],\\displaystyle=\[P\_\{f\}\(x\);P\_\{c\}\],\(1\)𝒞P​\(x\)\\displaystyle\\mathcal\{C\}\_\{P\}\(x\)=\[Eθ​\(x\);P⁡\(x\)\]\.\\displaystyle=\[E\_\{\\theta\}\(x\);P\(x\)\]\.Here,\[;\]\[\\,;\\,\]denotes concatenation along the sequence dimension\.P⁡\(x\)P\(x\)is the input\-dependent hierarchical safety prompt, whereas𝒞P​\(x\)\\mathcal\{C\}\_\{P\}\(x\)is the complete embedding sequence passed to the frozen language model\. The shared promptPcP\_\{c\}provides a category\-agnostic safety boundary, whilePf​\(x\)P\_\{f\}\(x\)injects the category\-specific control required by the current input\.

The hierarchical routerRϕR\_\{\\phi\}takes the frozen hidden states ofπθ\\pi\_\{\\theta\}as input\. A lightweight Transformer encoder followed by masked mean pooling produces an input\-level representation𝐡x\\mathbf\{h\}\_\{x\}\. The coarse\-grained prediction head outputs𝐩c​\(x\)=\[psafe​\(x\),prisk​\(x\)\]\\mathbf\{p\}\_\{c\}\(x\)=\[p\_\{\\mathrm\{safe\}\}\(x\),p\_\{\\mathrm\{risk\}\}\(x\)\]\. Its two dimensions follow the fixed class order0=safe0=\\mathrm\{safe\}and1=risk1=\\mathrm\{risk\}; equivalently,pc,0​\(x\)=psafe​\(x\)p\_\{c,0\}\(x\)=p\_\{\\mathrm\{safe\}\}\(x\)andpc,1​\(x\)=prisk​\(x\)p\_\{c,1\}\(x\)=p\_\{\\mathrm\{risk\}\}\(x\)\. The fine\-grained head outputs𝐩f​\(x\)∈\[0,1\]K\\mathbf\{p\}\_\{f\}\(x\)\\in\[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 categorykk, we optimize an independent soft\-prompt expertPkP\_\{k\}\. The fine\-grained risk scores produced by the router are normalized into composition weightsαk​\(x\)\\alpha\_\{k\}\(x\)and used to form a weighted combination of the experts:

Pf​\(x\)=∑k=1Kαk​\(x\)​Pk,αk​\(x\)=pf,k​\(x\)∑j=1Kpf,j​\(x\)\.P\_\{f\}\(x\)=\\sum\_\{k=1\}^\{K\}\\alpha\_\{k\}\(x\)P\_\{k\},\\qquad\\alpha\_\{k\}\(x\)=\\frac\{p\_\{f,k\}\(x\)\}\{\\sum\_\{j=1\}^\{K\}p\_\{f,j\}\(x\)\}\.\(2\)whereαk​\(x\)\\alpha\_\{k\}\(x\)is the normalized weight assigned to categorykk, with∑k=1Kαk​\(x\)=1\\sum\_\{k=1\}^\{K\}\\alpha\_\{k\}\(x\)=1\. Because allPkP\_\{k\}share the same shape, they can be combined element\-wise\.

### 4\.2Decoupled 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𝒟c=\{\(xi,yic\)\}\\mathcal\{D\}\_\{c\}=\\\{\(x\_\{i\},y\_\{i\}^\{c\}\)\\\}, whereyic∈\{0,1\}y\_\{i\}^\{c\}\\in\\\{0,1\\\}denotes the safe and risk label, we train the coarse\-grained prediction head using cross\-entropy loss\. For the fine\-grained dataset𝒟f=\{\(xi,𝐲if\)\}\\mathcal\{D\}\_\{f\}=\\\{\(x\_\{i\},\\mathbf\{y\}\_\{i\}^\{f\}\)\\\}, where𝐲if∈\{0,1\}K\\mathbf\{y\}\_\{i\}^\{f\}\\in\\\{0,1\\\}^\{K\}is a multi\-label risk vector, we train the fine\-grained prediction head using binary cross\-entropy loss\. The router objective is defined as follows:

ℒc\\displaystyle\\mathcal\{L\}\_\{c\}=𝔼𝒟c​\[CE⁡\(𝐩c​\(x\),yc\)\],\\displaystyle=\\mathbb\{E\}\_\{\\mathcal\{D\}\_\{c\}\}\\left\[\\operatorname\{CE\}\\\!\\left\(\\mathbf\{p\}\_\{c\}\(x\),y^\{c\}\\right\)\\right\],ℒf\\displaystyle\\mathcal\{L\}\_\{f\}=𝔼𝒟f​\[BCE⁡\(𝐩f​\(x\),𝐲f\)\],\\displaystyle=\\mathbb\{E\}\_\{\\mathcal\{D\}\_\{f\}\}\\left\[\\operatorname\{BCE\}\\\!\\left\(\\mathbf\{p\}\_\{f\}\(x\),\\mathbf\{y\}^\{f\}\\right\)\\right\],ℒrouter\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{router\}\}=ℒc\+ℒf\.\\displaystyle=\\mathcal\{L\}\_\{c\}\+\\mathcal\{L\}\_\{f\}\.\(3\)
Here, the expectations are taken over the corresponding datasets\.𝐩c​\(x\)\\mathbf\{p\}\_\{c\}\(x\)denotes the safe probability distribution predicted by the coarse\-grained head, while𝐩f​\(x\)∈\(0,1\)K\\mathbf\{p\}\_\{f\}\(x\)\\in\(0,1\)^\{K\}contains the probabilities predicted for theKKfine\-grained risk categories\.

We first optimize onlyℒc\\mathcal\{L\}\_\{c\}to learn a general safe/risk decision boundary and then jointly optimizeℒc\\mathcal\{L\}\_\{c\}andℒf\\mathcal\{L\}\_\{f\}\. During joint training,𝒟c\\mathcal\{D\}\_\{c\}supervises only the coarse\-grained head, whereas𝒟f\\mathcal\{D\}\_\{f\}supervises only the fine\-grained head\. We do not reuse𝒟f\\mathcal\{D\}\_\{f\}as additional coarse\-grained risk examples, thereby preventing the coarse decision boundary from overfitting the limited set of annotated risk categories\.

#### Stage II: Risk\-Adaptive Hierarchical Prompt Learning\.

We freeze the backboneπθ\\pi\_\{\\theta\}and the trained routerRϕR\_\{\\phi\}, and update only the shared coarse\-grained promptPcP\_\{c\}and the fine\-grained prompt experts\{Pk\}k=1K\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\}\. Given a preference triplet\(x,yw,yl\)\(x,y\_\{w\},y\_\{l\}\), whereywy\_\{w\}is a safe and helpful response andyly\_\{l\}is an unsafe or low\-quality response, we apply Direct Preference Optimization to increase the relative conditional likelihood ofywy\_\{w\}overyly\_\{l\}\. When updating the two types of prompts, we alternately mask their gradients to reduce interference between them\. When updatingPcP\_\{c\}, the routed fine\-grained promptPf​\(x\)P\_\{f\}\(x\)remains in the forward pass but receives no gradient\. Conversely, when updating the fine\-grained experts,PcP\_\{c\}remains in the forward pass but is held fixed\. Formally,

P~\(c\)​\(x\)\\displaystyle\\widetilde\{P\}^\{\(c\)\}\(x\)=\[sg⁡\(Pf​\(x\)\);Pc\],\\displaystyle=\\bigl\[\\operatorname\{sg\}\(P\_\{f\}\(x\)\);\\,P\_\{c\}\\bigr\],\(4\)P~\(f\)​\(x\)\\displaystyle\\widetilde\{P\}^\{\(f\)\}\(x\)=\[Pf​\(x\);sg⁡\(Pc\)\],\\displaystyle=\\bigl\[P\_\{f\}\(x\);\\,\\operatorname\{sg\}\(P\_\{c\}\)\\bigr\],whereP~\(c\)​\(x\)\\widetilde\{P\}^\{\(c\)\}\(x\)andP~\(f\)​\(x\)\\widetilde\{P\}^\{\(f\)\}\(x\)denote the composed prompts used for coarse\- and fine\-grained updates, respectively\. In both cases, the composed prompt is concatenated withEθ​\(x\)E\_\{\\theta\}\(x\)as in Eq\. \([1](https://arxiv.org/html/2608.12821#S4.E1)\) to form the complete prompted context before evaluating the DPO loss\. The stop\-gradient operatorsg⁡\(⋅\)\\operatorname\{sg\}\(\\cdot\)acts as the identity during the forward pass but blocks gradient propagation during backpropagation:

sg⁡\(z\)=z,∇z​sg​\(z\)=0\.\\operatorname\{sg\}\(z\)=z,\\qquad\\nabla\_\{z\}\\operatorname\{sg\}\(z\)=0\.\(5\)
For conciseness, we present only the core training objectives in the main text\.Appendix Aprovides the complete training algorithm, detailed formulations of the hierarchical router, and the full prompt\-level DPO objective\.

### 4\.3Input\-Adaptive Inference

At inference time, the router first computes the safe\-input probabilitypsafe​\(x\)p\_\{\\mathrm\{safe\}\}\(x\)\. Ifpsafe​\(x\)p\_\{\\mathrm\{safe\}\}\(x\)exceeds the thresholdτs\\tau\_\{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 constructsPf​\(x\)P\_\{f\}\(x\)from the fine\-grained risk scores𝐩f​\(x\)\\mathbf\{p\}\_\{f\}\(x\)and combines it with the shared promptPcP\_\{c\}to formP⁡\(x\)P\(x\)\. The resulting prompted context𝒞P​\(x\)\\mathcal\{C\}\_\{P\}\(x\)is then passed to the frozen language model for response generation:

y∼\{πθ​\(y∣x\),psafe​\(x\)\>τs,πθ​\(y∣𝒞P​\(x\)\),otherwise\.y\\sim\\begin\{cases\}\\pi\_\{\\theta\}\(y\\mid x\),&p\_\{\\mathrm\{safe\}\}\(x\)\>\\tau\_\{s\},\\\\ \\pi\_\{\\theta\}\\\!\\left\(y\\mid\\mathcal\{C\}\_\{P\}\(x\)\\right\),&\\text\{otherwise\}\.\\end\{cases\}\(6\)A largerτs\\tau\_\{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\.

## 5Experiments

### 5\.1Experimental Setting

#### Datasets and Models\.

Router training and prompt optimization use separate sources of supervision\. We construct a coarse\-grained binary classification dataset from WildGuardMix\([9](https://arxiv.org/html/2608.12821#bib.bib15)\), mapping inputs labeledunharmfulto the safe class and those labeledharmfulto the risk class\. The fine\-grained prediction head uses category annotations from PKU\-SafeRLHF\([7](https://arxiv.org/html/2608.12821#bib.bib4)\)and covers four risk categories: cybercrime, economic crime, privacy violations and violence\. We formulate this task as multi\-label classification because a single input may contain overlapping risks\. Prompt training uses preference triplets of the form\(prompt,chosen,rejected\)\(\\texttt\{prompt\},\\texttt\{chosen\},\\texttt\{rejected\}\)\. The shared coarse\-grained prompt is trained on 1200 coarse\-grained safety examples without risk\-category labels to learn cross\-category behavior, whereas the fine\-grained prompt experts are trained on 1000 category\-annotated preference examples drawn from the four risk categories\. We evaluate HiRoute on three open\-source instruction\-tuned models: Mistral\-7B\-Instruct\-v0\.3\([10](https://arxiv.org/html/2608.12821#bib.bib16)\), Vicuna\-7B\-v1\.5\([27](https://arxiv.org/html/2608.12821#bib.bib17)\)and Zephyr\-7B\-Beta\([23](https://arxiv.org/html/2608.12821#bib.bib18)\)\. These models differ in their pretraining and post\-training pipelines and exhibit different initial safety profiles, allowing us to assess whether HiRoute transfers across backbone models\.

Table 1:Main results across three models\. Each safety\-benchmark reports Safety Rate \(%\) / Safe\-Response Helpfulness score \(*S/H*\), and Avg\. Safety is the mean safety rate across the three safety benchmarks\. GSM8K accuracy evaluates mathematical reasoning, MT\-Bench measures general response quality, and TruthfulQA reports Truthfulness / Informativeness \(*T/I*\)\. XSTest OR is the over\-refusal rate on benign inputs\.↑\\uparrowindicates that higher values are better, whereas↓\\downarrowindicates that lower values are better\. The best and second\-best results for each backbone are marked inboldandunderlined, respectively\.
#### Baseline Methods\.

We compare HiRoute with four baselines\. Base denotes the original instruction\-tuned model without additional safety prompts or parameter updates and establishes the initial safety level of each model\. RPO\([28](https://arxiv.org/html/2608.12821#bib.bib10)\)optimizes a robust defensive prompt to improve resistance to jailbreak inputs\. DRO\([26](https://arxiv.org/html/2608.12821#bib.bib11)\)learns a continuous safety prompt and uses a refusal direction or safety\-representation signal to increase the probability of refusing harmful requests\. ACD\([25](https://arxiv.org/html/2608.12821#bib.bib12)\)jointly models safety and adversarial prompts and applies contrastive decoding to enlarge the distributional separation between safe and harmful responses\.

#### Safety Evaluation\.

StrongReject\([21](https://arxiv.org/html/2608.12821#bib.bib19)\), AdvBench\([29](https://arxiv.org/html/2608.12821#bib.bib20)\)and JailbreakBench\([5](https://arxiv.org/html/2608.12821#bib.bib21)\)cover direct harmful requests and jailbreak inputs\. For each benchmark, we report the safety rate and safe\-response helpfulness\. The safety rate measures the proportion of responses that do not materially facilitate the harmful objective\. Safe\-response helpfulness is evaluated only for responses deemed safe and measures whether they explain the relevant risks, provide compliant guidance, or offer safe alternatives\. Safety rate and safe\-response helpfulness are automatically evaluated by GPT\-5\.4 using an LLM\-as\-a\-judge protocol\([27](https://arxiv.org/html/2608.12821#bib.bib17)\), and we conduct human validation on a randomly sampled subset\. The evaluation prompts, scoring criteria, and manual consistency check are provided inAppendix E\.

#### Utility and Over\-Refusal\.

To further evaluate the preservation of general capabilities after safety alignment, we use GSM8K\([6](https://arxiv.org/html/2608.12821#bib.bib22)\)accuracy to measure mathematical reasoning, the MT\-Bench\([27](https://arxiv.org/html/2608.12821#bib.bib17)\)score to assess multi\-turn instruction following and open\-ended response quality, and TruthfulQA\([13](https://arxiv.org/html/2608.12821#bib.bib25)\)accuracy to evaluate factual reliability and resistance to common misconceptions\. We also report the over\-refusal rate on the XSTest\([19](https://arxiv.org/html/2608.12821#bib.bib23)\), which measures the proportion of benign requests that are incorrectly rejected\.

#### Implementation Details\.

All experiments use AdamW\([15](https://arxiv.org/html/2608.12821#bib.bib26)\)and two NVIDIA RTX 4090 GPUs, with the random seed fixed to 42\. The router consists of a single\-layer lightweight Transformer encoder with eight attention heads and two linear classification heads\. Following existing work\([2](https://arxiv.org/html/2608.12821#bib.bib27)\), we train it for eight epochs with a learning rate of1×10−41\\times 10^\{\-4\}and a batch size of 4, the first four epochs optimize only the coarse\-grained head, and the remaining four jointly optimize both heads\. During prompt training, we freeze the backbone and the trained router and optimize the prompts for six epochs using a learning rate of5×10−55\\times 10^\{\-5\}and a batch size of 4\([22](https://arxiv.org/html/2608.12821#bib.bib24)\)\. All coarse\- and fine\-grained prompts are randomly initialized from a normal distribution\.

### 5\.2Main Results

#### Defense Effectiveness and Helpfulness\.

Table[1](https://arxiv.org/html/2608.12821#S5.T1)shows that HiRoute achieves the best average safety and safe\-response helpfulness on all models\. Its average safety rates reach 93\.2%, 97\.7% and 94\.8% on Mistral, Vicuna and Zephyr, outperforming the strongest baseline by 1\.7, 1\.3 and 2\.9 percentage points, respectively\. Meanwhile, its average helpfulness scores are 7\.4, 6\.1 and 6\.8, exceeding the best competing results\. Relative to the base models, HiRoute improves safety without reducing helpfulness\. These joint gains indicate that HiRoute does not rely on more aggressive refusal alone: the shared prompt establishes a stable safety boundary, while the routed fine\-grained experts preserve risk\-specific explanations and safe alternatives\. Qualitative response examples are provided inAppendix D\.

#### Utility Evaluation\.

Table[1](https://arxiv.org/html/2608.12821#S5.T1)shows mathematical reasoning, open\-ended response quality and factual reliability using GSM8K, MT\-Bench and TruthfulQA\. Across all three models, HiRoute achieves the best or second\-best result on every utility metric among the safety\-aligned methods\. On Mistral, it retains 48\.0% GSM8K accuracy, a 5\.86 MT\-Bench score, and 78\.0%/100\.0% on TruthfulQA while increasing average safety from 53\.2% to 93\.2%\. HiRoute also produces the lowest over\-refusal rates among the aligned methods\. These results indicate that its safety gains do not rely on uniformly conservative generation and are consistent with the gate allowing benign inputs to bypass safety prompting\.

#### Over\-Refusal Evaluation\.

On Mistral, HiRoute achieves an XSTest over\-refusal rate of 2\.5%, compared with 40\.5% for RPO and 18\.0% for DRO, corresponding to absolute reductions of 38\.0 and 15\.5 percentage points, respectively\. Together with its average safety rate of 93\.2%, this result argues against the simple explanation that HiRoute improves safety by refusing inputs indiscriminately\. Instead, it indicates that the coarse\-grained gate concentrates safety control on risky inputs\.

Table 2:Effect of the safety\-gating thresholdτs\\tau\_\{s\}on held\-out validation data\. The validation data are disjoint from the benchmark test sets used for the final evaluation\.Table 3:Effect of the prompt\-training strategy\. All entries are safety rates \(%\)\. The proposed joint\-context alternating strategy is highlighted in gray\.

### 5\.3Ablation Studies

We conduct ablation studies on Mistral\-7B\-Instruct\-v0\.3 to analyze prompt\-capacity allocation, the inference threshold, and the training strategy\. Unless otherwise specified, each safety\-benchmark entry reports*Safety Rate / Safe\-Response Helpfulness*\.These experiments analyze the impact of each component on model performance and validate its effectiveness within HiRoute\.

#### Effect of the Coarse and Fine Prompt Length Ratio\.

Table[4](https://arxiv.org/html/2608.12821#S5.T4)compares different prompt allocations while fixing the total length at 20\. Increasing the coarse\-grained capacity generally improves safety, but excessively reducing the fine\-grained capacity substantially degrades safe\-response helpfulness\. The 15/5 configuration achieves an average safety rate of 93\.2% and an average helpfulness score of 7\.4\. Its safety rate is only 0\.5 percentage points below the maximum obtained by 19/1, while its helpfulness is 0\.7 points higher\. It also matches 17/3 in safety while improving helpfulness by 0\.4 points\. We therefore select 15/5 as the default\.

#### Effect of the Inference Threshold\.

Table[2](https://arxiv.org/html/2608.12821#S5.T2)shows the effect of the safety\-gating thresholdτs\\tau\_\{s\}on held\-out validation data, which is disjoint from the benchmark test sets used for the final evaluation\. Asτs\\tau\_\{s\}increases from 0\.50 to 0\.95, the JailbreakBench safety rate steadily rises from 82\.0% to 95\.0%, while GSM8K accuracy decreases only from 50\.5% to 48\.0%\. This result indicates that stricter gating substantially strengthens safety control at a limited cost to general capability\. Compared with 0\.90, a threshold of 0\.95 further improves the safety rate by 1\.5 percentage points while reducing GSM8K accuracy by only 0\.5 percentage points\. We use 0\.95 as the default threshold based on the validation results and keep it fixed for all subsequent test\-set evaluations\.

Table 4:Effect of the coarse\-to\-fine prompt length allocation\. The total prompt length is fixed at 20\. Each benchmark reports Safety \(%\)/Helpfulness\. Best results arebolded, second\-best results areunderlined, and the selected configuration is shaded in gray\.
#### Effect of the Training Strategy\.

As shown in Table[3](https://arxiv.org/html/2608.12821#S5.T3), coarse after fine first trains the fine\-grained prompt in a context containing onlyPf​\(x\)P\_\{f\}\(x\), then freezes it and introducesPcP\_\{c\}to train the coarse\-grained prompt\. This discontinuous sequential procedure preventsPf​\(x\)P\_\{f\}\(x\)from adapting to the final coarse\-to\-fine composed context\. The subsequent stage also directly inherits optimization errors and routing noise from the fixed fine\-grained prompt, making it difficult for the two prompt levels to learn coordinated control\. In contrast, joint alternating training always performs the forward pass using the final composed context\[x;Pf​\(x\);Pc\]\[x;P\_\{f\}\(x\);P\_\{c\}\]and alternately updates the two prompt types through gradient masking\. This allows them to learn coordinated safety control while reducing parameter interference\. The joint alternating strategy achieves higher safety rates on all three benchmarks and improves the average safety rate from 74\.8% to 93\.2%, demonstrating the importance of maintaining a consistent training context for coarse\- and fine\-grained prompts\.

Additional Analyses\.Appendix Bvalidates the effectiveness of input\-dependent routing and the shared coarse\-grained prompt, and reports hierarchical\-router performance on the held\-out test sets, including harmful\-input recall and fine\-grained multi\-label F1\. It also presents prompt\-length and update\-ratio ablations, together with a scaling evaluation on Vicuna\-13B\-v1\.5\.Appendix Cevaluates transferred GCG attacks\([29](https://arxiv.org/html/2608.12821#bib.bib20)\), under which HiRoute raises the average safety rate across the three models from 34\.6% to 88\.3% under attack\.

## 6Conclusion

We present HiRoute, an input\-adaptive hierarchical prompt\-tuning framework for parameter\-efficient safety alignment\. Our analysis reveals complementary limitations in prompt\-based safety\-alignment approaches: a coarse\-grained prompt provides stable safety but tends to produce less informative refusals, whereas routed fine\-grained prompts improve safe\-response helpfulness but provide weaker safety control\. HiRoute addresses this tension by combining a shared coarse\-grained prompt, which establishes a category\-agnostic safety boundary, with routed fine\-grained experts that provide risk\-specific guidance and constructive safe responses\. This division reduces dependence on precise expert matching while avoiding uniform refusals\. A coarse\-grained gate further allows benign inputs to bypass unnecessary safety intervention\. The framework freezes the backbone and separates risk routing from prompt optimization through two\-stage training\. Experiments across multiple models and evaluation settings show that HiRoute consistently improves safety while preserving safe\-response helpfulness and general utility, with limited over\-refusal\. Ablations further support the effectiveness of hierarchical prompt composition, safety gating, and coordinated prompt optimization\. Overall, input\-dependent hierarchical prompting balances safety, response quality, parameter efficiency, and general utility\.

## 7Limitations and Future Work

HiRoute partly depends on the accuracy of its hierarchical router\. Misclassified or out\-of\-distribution compound\-risk inputs may receive suboptimal expert weights, reducing the stability and specificity of safety alignment\. Our experiments primarily focus on single\-turn text interactions, leaving multi\-turn and multimodal settings underexplored\. Future work will investigate open\-set risk recognition, calibrated and uncertainty\-aware routing, and extensions to multi\-turn and multimodal safety alignment\.

## References

- Alfarraet al\.\(2026\)M\. Alfarra, C\. Pinneri, D\. Kianfar, M\. Almousa, and C\. LouizosDistilling safe llm systems via soft prompts for on device settings\.External Links:2606\.09388,[Link](https://arxiv.org/abs/2606.09388)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1)\.
- Anonymous \(2026\)AnonymousSafeMoE: leveraging unsafe data to train safer, more informative LLMs\.External Links:[Link](https://openreview.net/forum?id=cMc5mFC7ib)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx5.p1.1)\.
- Baiet al\.\(2022\)Y\. Bai, S\. Kadavath, S\. Kundu, A\. Askell, J\. Kernion, A\. Jones, A\. Chen, A\. Goldie, A\. Mirhoseini, C\. McKinnon,et al\.Constitutional ai: harmlessness from ai feedback\.arXiv preprint arXiv:2212\.08073\.Cited by:[§2\.1](https://arxiv.org/html/2608.12821#S2.SS1.p1.1)\.
- Bianchiet al\.\(2024\)F\. Bianchi, M\. Suzgun, G\. Attanasio, P\. Röttger, D\. Jurafsky, T\. Hashimoto, and J\. Y\. ZouSafety\-tuned llamas: lessons from improving the safety of large language models that follow instructions\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 34196–34216\.Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p1.1)\.
- Chaoet al\.\(2024\)P\. Chao, E\. Debenedetti, A\. Robey, M\. Andriushchenko, F\. Croce, V\. Sehwag, E\. Dobriban, N\. Flammarion, G\. J\. Pappas, F\. Tramer, H\. Hassani, and E\. WongJailbreakBench: an open robustness benchmark for jailbreaking large language models\.External Links:2404\.01318,[Link](https://arxiv.org/abs/2404.01318)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx3.p1.1)\.
- Cobbeet al\.\(2021\)K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. SchulmanTraining verifiers to solve math word problems\.External Links:2110\.14168,[Link](https://arxiv.org/abs/2110.14168)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx4.p1.1)\.
- Daiet al\.\(2024\)J\. Dai, X\. Pan, R\. Sun, J\. Ji, X\. Xu, M\. Liu, Y\. Wang, and Y\. YangSafe rlhf: safe reinforcement learning from human feedback\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 50750–50777\.Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p3.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx1.p1.1)\.
- Duanet al\.\(2025\)R\. Duan, J\. Liu, X\. Jia, S\. Zhao, R\. Cheng, F\. Wang, C\. Wei, Y\. Xie, C\. Liu, D\. Li, Y\. Dong, Y\. Zhang, Y\. Chen, C\. Wang, X\. Ma, X\. Wei, Y\. Liu, H\. Su, J\. Zhu, X\. Li, Y\. Sun, J\. Zhang, J\. Hu, S\. Xu, W\. Yang, Y\. Yang, X\. Zhang, Y\. Tan, J\. Tao, and H\. XueOyster\-i: beyond refusal – constructive safety alignment for responsible language models\.External Links:2509\.01909,[Link](https://arxiv.org/abs/2509.01909)Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p3.1)\.
- Hanet al\.\(2024\)S\. Han, K\. Rao, A\. Ettinger, L\. Jiang, B\. Y\. Lin, N\. Lambert, Y\. Choi, and N\. DziriWildGuard: open one\-stop moderation tools for safety risks, jailbreaks, and refusals of llms\.External Links:2406\.18495,[Link](https://arxiv.org/abs/2406.18495)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx1.p1.1)\.
- Jianget al\.\(2023\)A\. Q\. Jiang, A\. Sablayrolles, A\. Mensch, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, F\. Bressand, G\. Lengyel, G\. Lample, L\. Saulnier, L\. R\. Lavaud, M\. Lachaux, P\. Stock, T\. L\. Scao, T\. Lavril, T\. Wang, T\. Lacroix, and W\. E\. SayedMistral 7b\.External Links:2310\.06825,[Link](https://arxiv.org/abs/2310.06825)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx1.p1.1)\.
- Lesteret al\.\(2021\)B\. Lester, R\. Al\-Rfou, and N\. ConstantThe power of scale for parameter\-efficient prompt tuning\.InProceedings of the 2021 conference on empirical methods in natural language processing,pp\. 3045–3059\.Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1)\.
- Li and Liang \(2021\)X\. L\. Li and P\. LiangPrefix\-tuning: optimizing continuous prompts for generation\.InProceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing \(Volume 1: Long Papers\),pp\. 4582–4597\.Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1)\.
- Linet al\.\(2022\)S\. Lin, J\. Hilton, and O\. EvansTruthfulQA: measuring how models mimic human falsehoods\.External Links:2109\.07958,[Link](https://arxiv.org/abs/2109.07958)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx4.p1.1)\.
- Liuet al\.\(2022\)X\. Liu, K\. Ji, Y\. Fu, W\. L\. Tam, Z\. Du, Z\. Yang, and J\. TangP\-tuning v2: prompt tuning can be comparable to fine\-tuning universally across scales and tasks\.External Links:2110\.07602,[Link](https://arxiv.org/abs/2110.07602)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1)\.
- Loshchilov and Hutter \(2019\)I\. Loshchilov and F\. HutterDecoupled weight decay regularization\.External Links:1711\.05101,[Link](https://arxiv.org/abs/1711.05101)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx5.p1.1)\.
- Ouyanget al\.\(2022\)L\. Ouyang, J\. Wu, X\. Jiang, D\. Almeida, C\. Wainwright, P\. Mishkin, C\. Zhang, S\. Agarwal, K\. Slama, A\. Ray,et al\.Training language models to follow instructions with human feedback\.Advances in neural information processing systems35,pp\. 27730–27744\.Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p1.1),[§2\.1](https://arxiv.org/html/2608.12821#S2.SS1.p1.1)\.
- Penget al\.\(2026\)J\. Peng, H\. Chen, J\. Dong, M\. Wang, W\. Li, Y\. Li, K\. Zhang, and X\. ZhaoMOSAIC: composable safety alignment with modular control tokens\.External Links:2603\.16210,[Link](https://arxiv.org/abs/2603.16210)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1)\.
- Rafailovet al\.\(2023\)R\. Rafailov, A\. Sharma, E\. Mitchell, C\. D\. Manning, S\. Ermon, and C\. FinnDirect preference optimization: your language model is secretly a reward model\.Advances in neural information processing systems36,pp\. 53728–53741\.Cited by:[§2\.1](https://arxiv.org/html/2608.12821#S2.SS1.p1.1)\.
- Röttgeret al\.\(2024\)P\. Röttger, H\. Kirk, B\. Vidgen, G\. Attanasio, F\. Bianchi, and D\. HovyXstest: a test suite for identifying exaggerated safety behaviours in large language models\.InProceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(Volume 1: Long Papers\),pp\. 5377–5400\.Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx4.p1.1)\.
- Singhet al\.\(2025\)A\. Singh, A\. Fry, A\. Perelman, A\. Tart, A\. Ganesh, A\. El\-Kishky, A\. McLaughlin, A\. Low, A\. Ostrow, A\. Ananthram,et al\.Openai gpt\-5 system card\.arXiv preprint arXiv:2601\.03267\.Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p3.1)\.
- Soulyet al\.\(2024\)A\. Souly, Q\. Lu, D\. Bowen, T\. Trinh, E\. Hsieh, S\. Pandey, P\. Abbeel, J\. Svegliato, S\. Emmons, O\. Watkins, and S\. ToyerA strongreject for empty jailbreaks\.External Links:2402\.10260,[Link](https://arxiv.org/abs/2402.10260)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx3.p1.1)\.
- Taraghiet al\.\(2025\)M\. Taraghi, Y\. Pequignot, A\. Nikanjam, M\. A\. Merzouk, and F\. KhomhEfficiency vs\. alignment: investigating safety and fairness risks in parameter\-efficient fine\-tuning of llms\.External Links:2511\.00382,[Link](https://arxiv.org/abs/2511.00382)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx5.p1.1)\.
- Tunstallet al\.\(2023\)L\. Tunstall, E\. Beeching, N\. Lambert, N\. Rajani, K\. Rasul, Y\. Belkada, S\. Huang, L\. von Werra, C\. Fourrier, N\. Habib, N\. Sarrazin, O\. Sanseviero, A\. M\. Rush, and T\. WolfZephyr: direct distillation of lm alignment\.External Links:2310\.16944,[Link](https://arxiv.org/abs/2310.16944)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx1.p1.1)\.
- Weiet al\.\(2022\)J\. Wei, M\. Bosma, V\. Y\. Zhao, K\. Guu, A\. W\. Yu, B\. Lester, N\. Du, A\. M\. Dai, and Q\. V\. LeFinetuned language models are zero\-shot learners\.External Links:2109\.01652,[Link](https://arxiv.org/abs/2109.01652)Cited by:[§1](https://arxiv.org/html/2608.12821#S1.p1.1)\.
- Zhanget al\.\(2025\)X\. Zhang, Z\. Zhao, W\. Shi, K\. Xu, D\. Huang, and X\. HuSafety alignment of large language models via contrasting safe and harmful distributions\.External Links:2406\.16743,[Link](https://arxiv.org/abs/2406.16743)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx2.p1.1)\.
- Zhenget al\.\(2024\)C\. Zheng, F\. Yin, H\. Zhou, F\. Meng, J\. Zhou, K\. Chang, M\. Huang, and N\. PengOn prompt\-driven safeguarding for large language models\.External Links:2401\.18018,[Link](https://arxiv.org/abs/2401.18018)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx2.p1.1)\.
- Zhenget al\.\(2023\)L\. Zheng, W\. Chiang, Y\. Sheng, S\. Zhuang, Z\. Wu, Y\. Zhuang, Z\. Lin, Z\. Li, D\. Li, E\. P\. Xing, H\. Zhang, J\. E\. Gonzalez, and I\. StoicaJudging llm\-as\-a\-judge with mt\-bench and chatbot arena\.External Links:2306\.05685,[Link](https://arxiv.org/abs/2306.05685)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx1.p1.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx3.p1.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx4.p1.1)\.
- Zhouet al\.\(2024\)A\. Zhou, B\. Li, and H\. WangRobust prompt optimization for defending language models against jailbreaking attacks\.External Links:2401\.17263,[Link](https://arxiv.org/abs/2401.17263)Cited by:[§2\.2](https://arxiv.org/html/2608.12821#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx2.p1.1)\.
- Zouet al\.\(2023\)A\. Zou, Z\. Wang, N\. Carlini, M\. Nasr, J\. Z\. Kolter, and M\. FredriksonUniversal and transferable adversarial attacks on aligned language models\.External Links:2307\.15043,[Link](https://arxiv.org/abs/2307.15043)Cited by:[§5\.1](https://arxiv.org/html/2608.12821#S5.SS1.SSSx3.p1.1),[§5\.3](https://arxiv.org/html/2608.12821#S5.SS3.SSSx3.p2.1)\.

Appendix

Appendix overview\.Appendix A provides the full methodological details of HiRoute, including the hierarchical router, prompt composition, and the alternating prompt\-optimization objective\. Appendix B first evaluates the effectiveness of input\-dependent routing and the shared coarse\-grained prompt, and then reports router performance, hyperparameter ablations, and scaling results on a larger backbone\. Appendix C evaluates robustness against transfer\-based GCG attacks\. Appendix D presents representative comparisons of responses generated by the base model, the coarse\-prompt variant, and HiRoute\. Appendix E provides the LLM\-as\-a\-judge prompts, scoring criteria, and the manual consistency check for the automatic evaluation\.

## Appendix AMethodological Details of HiRoute

### A\.1Hierarchical Router

Given an inputxxof lengthTT, the frozen backboneπθ\\pi\_\{\\theta\}produces final\-layer statesHθ​\(x\)∈ℝT×dH\_\{\\theta\}\(x\)\\in\\mathbb\{R\}^\{T\\times d\}\. A lightweight Transformer encoderRϕR\_\{\\phi\}models token interactions, and masked mean pooling produces the input representation:

Hθ​\(x\)\\displaystyle H\_\{\\theta\}\(x\)=\[𝐡1,…,𝐡T\],\\displaystyle=\[\\mathbf\{h\}\_\{1\},\\ldots,\\mathbf\{h\}\_\{T\}\],\(7\)Zϕ​\(x\)\\displaystyle Z\_\{\\phi\}\(x\)=TrEncϕ⁡\(Hθ​\(x\);𝐦\)=\[𝐳1,…,𝐳T\],\\displaystyle=\\operatorname\{TrEnc\}\_\{\\phi\}\(H\_\{\\theta\}\(x\);\\mathbf\{m\}\)=\[\\mathbf\{z\}\_\{1\},\\ldots,\\mathbf\{z\}\_\{T\}\],𝐡x\\displaystyle\\mathbf\{h\}\_\{x\}=∑t=1Tmt​𝐳tmax⁡\(1,∑t=1Tmt\)\.\\displaystyle=\\frac\{\\sum\_\{t=1\}^\{T\}m\_\{t\}\\mathbf\{z\}\_\{t\}\}\{\\max\(1,\\sum\_\{t=1\}^\{T\}m\_\{t\}\)\}\.Here,ddis the hidden dimension and𝐦=\(m1,…,mT\)∈\{0,1\}T\\mathbf\{m\}=\(m\_\{1\},\\ldots,m\_\{T\}\)\\in\\\{0,1\\\}^\{T\}is the attention mask\.

The coarse head predicts a safe/risk distribution, and the fine head independently predictsKKrisk categories:

𝐩c​\(x\)\\displaystyle\\mathbf\{p\}\_\{c\}\(x\)=softmax⁡\(Wc​𝐡x\+𝐛c\)=\[psafe​\(x\),prisk​\(x\)\],\\displaystyle=\\operatorname\{softmax\}\(W\_\{c\}\\mathbf\{h\}\_\{x\}\+\\mathbf\{b\}\_\{c\}\)=\[p\_\{\\mathrm\{safe\}\}\(x\),p\_\{\\mathrm\{risk\}\}\(x\)\],\(8\)𝐩f​\(x\)\\displaystyle\\mathbf\{p\}\_\{f\}\(x\)=sigmoid⁡\(Wf​𝐡x\+𝐛f\)∈\(0,1\)K\.\\displaystyle=\\operatorname\{sigmoid\}\(W\_\{f\}\\mathbf\{h\}\_\{x\}\+\\mathbf\{b\}\_\{f\}\)\\in\(0,1\)^\{K\}\.The coarse labels follow0=safe0=\\mathrm\{safe\}and1=risk1=\\mathrm\{risk\}\. Here,Wc∈ℝ2×dW\_\{c\}\\in\\mathbb\{R\}^\{2\\times d\}andWf∈ℝK×dW\_\{f\}\\in\\mathbb\{R\}^\{K\\times d\}are classifier weights, with𝐛c\\mathbf\{b\}\_\{c\}and𝐛f\\mathbf\{b\}\_\{f\}denoting their biases\. Independent sigmoid outputs allow one input to activate multiple risk categories\.

The coarse\- and fine\-grained router datasets are

𝒟c\\displaystyle\\mathcal\{D\}\_\{c\}=\{\(xi,yic\)\},\\displaystyle=\\\{\(x\_\{i\},y\_\{i\}^\{c\}\)\\\},yic\\displaystyle y\_\{i\}^\{c\}∈\{0,1\},\\displaystyle\\in\\\{0,1\\\},\(9\)𝒟f\\displaystyle\\mathcal\{D\}\_\{f\}=\{\(xi,𝐲if\)\},\\displaystyle=\\\{\(x\_\{i\},\\mathbf\{y\}\_\{i\}^\{f\}\)\\\},𝐲if\\displaystyle\\mathbf\{y\}\_\{i\}^\{f\}∈\{0,1\}K\.\\displaystyle\\in\\\{0,1\\\}^\{K\}\.Their cross\-entropy and binary cross\-entropy objectives are

ℒc\\displaystyle\\mathcal\{L\}\_\{c\}=−𝔼𝒟c​log⁡pc,yc​\(x\),\\displaystyle=\-\\mathbb\{E\}\_\{\\mathcal\{D\}\_\{c\}\}\\log p\_\{c,y^\{c\}\}\(x\),\(10\)ℒf\\displaystyle\\mathcal\{L\}\_\{f\}=−𝔼𝒟f1K∑k=1K\[ykflogpf,k\(x\)\\displaystyle=\-\\mathbb\{E\}\_\{\\mathcal\{D\}\_\{f\}\}\\frac\{1\}\{K\}\\sum\_\{k=1\}^\{K\}\\left\[y\_\{k\}^\{f\}\\log p\_\{f,k\}\(x\)\\right\.\+\(1−ykf\)log\(1−pf,k\(x\)\)\],\\displaystyle\\left\.\+\(1\-y\_\{k\}^\{f\}\)\\log\(1\-p\_\{f,k\}\(x\)\)\\right\],ℒrouter\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{router\}\}=ℒc\+ℒf\.\\displaystyle=\\mathcal\{L\}\_\{c\}\+\\mathcal\{L\}\_\{f\}\.The expectations are taken over the datasets indicated in Equation[9](https://arxiv.org/html/2608.12821#A1.E9)\. We first optimizeℒc\\mathcal\{L\}\_\{c\}to establish a general safe/risk boundary and then jointly optimize both heads\. During joint training,𝒟c\\mathcal\{D\}\_\{c\}supervises only the coarse head and𝒟f\\mathcal\{D\}\_\{f\}only the fine head; this prevents the coarse boundary from overfitting the limited fine\-grained categories\.

### A\.2Routed Prompt Composition and Input Order

The fine\-grained probabilities are normalized into nonnegative routing weights:

αk​\(x\)=pf,k​\(x\)∑j=1Kpf,j​\(x\),αk​\(x\)≥0,∑k=1Kαk​\(x\)=1\.\\alpha\_\{k\}\(x\)=\\frac\{p\_\{f,k\}\(x\)\}\{\\sum\_\{j=1\}^\{K\}p\_\{f,j\}\(x\)\},\\qquad\\alpha\_\{k\}\(x\)\\geq 0,\\qquad\\sum\_\{k=1\}^\{K\}\\alpha\_\{k\}\(x\)=1\.\(11\)LetPk∈ℝLf×dP\_\{k\}\\in\\mathbb\{R\}^\{L\_\{f\}\\times d\}be thekk\-th fine\-grained prompt expert andPc∈ℝLc×dP\_\{c\}\\in\\mathbb\{R\}^\{L\_\{c\}\\times d\}the shared coarse\-grained prompt\. The routed and hierarchical prompts, together with the compact context used in the main text, are

Pf​\(x\)\\displaystyle P\_\{f\}\(x\)=∑k=1Kαk​\(x\)​Pk,\\displaystyle=\\sum\_\{k=1\}^\{K\}\\alpha\_\{k\}\(x\)P\_\{k\},\(12\)P⁡\(x\)\\displaystyle P\(x\)=\[Pf​\(x\);Pc\],\\displaystyle=\[P\_\{f\}\(x\);P\_\{c\}\],𝒞P​\(x\)\\displaystyle\\mathcal\{C\}\_\{P\}\(x\)=\[Eθ​\(x\);P⁡\(x\)\]\.\\displaystyle=\[E\_\{\\theta\}\(x\);P\(x\)\]\.Here,LfL\_\{f\}andLcL\_\{c\}are prompt lengths,Eθ​\(x\)E\_\{\\theta\}\(x\)denotes the backbone input embeddings, and\[;\]\[\\,;\\,\]denotes sequence concatenation\.

For the exact chat template, letUθ​\(x\)U\_\{\\theta\}\(x\)contain the formatted user input and its template tokens, and letAθA\_\{\\theta\}be the assistant\-prefix embeddings\. The prompted policy context and the unprompted reference context are

𝒞Pchat​\(x\)\\displaystyle\\mathcal\{C\}\_\{P\}^\{\\mathrm\{chat\}\}\(x\)=Uθ​\(x\)⊕Pf​\(x\)⊕Pc⊕Aθ,\\displaystyle=U\_\{\\theta\}\(x\)\\oplus P\_\{f\}\(x\)\\oplus P\_\{c\}\\oplus A\_\{\\theta\},\(13\)𝒞0chat​\(x\)\\displaystyle\\mathcal\{C\}\_\{0\}^\{\\mathrm\{chat\}\}\(x\)=Uθ​\(x\)⊕Aθ\.\\displaystyle=U\_\{\\theta\}\(x\)\\oplus A\_\{\\theta\}\.The operator⊕\\oplusdenotes sequence concatenation\. Both branches use the same frozen backbone; they differ only in whether the hierarchical safety prompts are inserted\.

### A\.3Prompt\-Level DPO Objective

Let\(x,yw,yl\)\(x,y\_\{w\},y\_\{l\}\)be a preference triplet from𝒟pref\\mathcal\{D\}\_\{\\mathrm\{pref\}\}, whereywy\_\{w\}is the chosen safe and helpful response andyly\_\{l\}is the rejected response\. Forb∈\{P,0\}b\\in\\\{P,0\\\}, define the response log\-likelihood under context𝒞bchat​\(x\)\\mathcal\{C\}\_\{b\}^\{\\mathrm\{chat\}\}\(x\)as

log⁡πb​\(y∣x\)\\displaystyle\\log\\pi\_\{b\}\(y\\mid x\)=∑t=1Tymty​log⁡πθ​\(yt∣𝒞bchat​\(x\),y<t\)\.\\displaystyle=\\sum\_\{t=1\}^\{T\_\{y\}\}m\_\{t\}^\{y\}\\log\\pi\_\{\\theta\}\\\!\\left\(y\_\{t\}\\mid\\mathcal\{C\}\_\{b\}^\{\\mathrm\{chat\}\}\(x\),y\_\{<t\}\\right\)\.\(14\)whereTyT\_\{y\}is the response length andmtym\_\{t\}^\{y\}masks non\-response tokens\. The policy and reference preference margins are

ΔP​\(x\)\\displaystyle\\Delta\_\{P\}\(x\)=log⁡πP​\(yw∣x\)−log⁡πP​\(yl∣x\),\\displaystyle=\\log\\pi\_\{P\}\(y\_\{w\}\\mid x\)\-\\log\\pi\_\{P\}\(y\_\{l\}\\mid x\),\(15\)Δ0​\(x\)\\displaystyle\\Delta\_\{0\}\(x\)=log⁡π0​\(yw∣x\)−log⁡π0​\(yl∣x\)\.\\displaystyle=\\log\\pi\_\{0\}\(y\_\{w\}\\mid x\)\-\\log\\pi\_\{0\}\(y\_\{l\}\\mid x\)\.The prompt\-level DPO loss is

ℒDPO=−𝔼𝒟pref​log⁡σ⁡\(β⁡\[ΔP​\(x\)−Δ0​\(x\)\]\),\\mathcal\{L\}\_\{\\mathrm\{DPO\}\}=\-\\mathbb\{E\}\_\{\\mathcal\{D\}\_\{\\mathrm\{pref\}\}\}\\log\\sigma\\\!\\left\(\\beta\[\\Delta\_\{P\}\(x\)\-\\Delta\_\{0\}\(x\)\]\\right\),\(16\)whereσ\\sigmais the sigmoid function andβ=0\.1\\beta=0\.1\. Coarse and fine updates respectively use𝒟prefc\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{c\}and𝒟preff\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{f\}\. Becauseπθ\\pi\_\{\\theta\}andRϕR\_\{\\phi\}are frozen, the loss updates only the prompt parameters\.

To retain the final composed context while preventing direct gradient interference, we alternate which prompt level receives gradients:

P~\(c\)​\(x\)\\displaystyle\\widetilde\{P\}^\{\(c\)\}\(x\)=\[sg⁡\(Pf​\(x\)\);Pc\],\\displaystyle=\[\\operatorname\{sg\}\(P\_\{f\}\(x\)\);P\_\{c\}\],\(17\)P~\(f\)​\(x\)\\displaystyle\\widetilde\{P\}^\{\(f\)\}\(x\)=\[Pf​\(x\);sg⁡\(Pc\)\],\\displaystyle=\[P\_\{f\}\(x\);\\operatorname\{sg\}\(P\_\{c\}\)\],sg⁡\(z\)\\displaystyle\\operatorname\{sg\}\(z\)=z,∇zsg\(z\)=0\.\\displaystyle=z,\\qquad\\nabla\_\{z\}\\operatorname\{sg\}\(z\)=0\.The corresponding chat\-template contexts are

𝒞~P\(c\)​\(x\)\\displaystyle\\widetilde\{\\mathcal\{C\}\}\_\{P\}^\{\(c\)\}\(x\)=Uθ​\(x\)⊕P~\(c\)​\(x\)⊕Aθ,\\displaystyle=U\_\{\\theta\}\(x\)\\oplus\\widetilde\{P\}^\{\(c\)\}\(x\)\\oplus A\_\{\\theta\},\(18\)𝒞~P\(f\)​\(x\)\\displaystyle\\widetilde\{\\mathcal\{C\}\}\_\{P\}^\{\(f\)\}\(x\)=Uθ​\(x\)⊕P~\(f\)​\(x\)⊕Aθ\.\\displaystyle=U\_\{\\theta\}\(x\)\\oplus\\widetilde\{P\}^\{\(f\)\}\(x\)\\oplus A\_\{\\theta\}\.A coarse update modifies onlyPcP\_\{c\}using𝒟prefc\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{c\}; the following fine update modifies only\{Pk\}k=1K\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\}using𝒟preff\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{f\}\. Both levels remain in every forward pass, matching the context used at inference\.

### A\.4Input\-Adaptive Inference

At inference, the coarse head determines whether safety prompting is activated:

y∼\{πθ​\(y∣𝒞0chat​\(x\)\),psafe​\(x\)\>τs,πθ​\(y∣𝒞Pchat​\(x\)\),psafe​\(x\)≤τs\.y\\sim\\begin\{cases\}\\pi\_\{\\theta\}\(y\\mid\\mathcal\{C\}\_\{0\}^\{\\mathrm\{chat\}\}\(x\)\),&p\_\{\\mathrm\{safe\}\}\(x\)\>\\tau\_\{s\},\\\\\[2\.0pt\] \\pi\_\{\\theta\}\(y\\mid\\mathcal\{C\}\_\{P\}^\{\\mathrm\{chat\}\}\(x\)\),&p\_\{\\mathrm\{safe\}\}\(x\)\\leq\\tau\_\{s\}\.\\end\{cases\}\(19\)Thus, confidently safe inputs bypass the prompts, while all other inputs use the shared prompt and the routed fine\-grained mixture\. The thresholdτs∈\[0,1\]\\tau\_\{s\}\\in\[0,1\]controls this decision\.

### A\.5Training and Inference Algorithm

Algorithm[1](https://arxiv.org/html/2608.12821#alg1)summarizes the two training stages and input\-adaptive inference\. Stage I learns the hierarchical router with the backbone frozen\. Stage II freezes the router and alternately updates the shared and fine\-grained prompts\.

Algorithm 1Training and Inference of HiRoute0:Frozen backbone

πθ\\pi\_\{\\theta\}; datasets

𝒟c,𝒟f,𝒟prefc,𝒟preff\\mathcal\{D\}\_\{c\},\\mathcal\{D\}\_\{f\},\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{c\},\\mathcal\{D\}\_\{\\mathrm\{pref\}\}^\{f\}; coarse\-only, joint\-router, and prompt\-training epochs

Ec,Ej,EpE\_\{c\},E\_\{j\},E\_\{p\};

β,τs\\beta,\\tau\_\{s\}
0:Router

RϕR\_\{\\phi\}, prompts

Pc,\{Pk\}k=1KP\_\{c\},\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\}, and response

yy
Stage I: Train the hierarchical router

for

e=1,…,Ece=1,\\ldots,E\_\{c\}do

Update the router encoder and coarse head using

ℒc\\mathcal\{L\}\_\{c\}
endfor

for

e=1,…,Eje=1,\\ldots,E\_\{j\}do

Update

RϕR\_\{\\phi\}using

ℒc\+ℒf\\mathcal\{L\}\_\{c\}\+\\mathcal\{L\}\_\{f\}
endfor

Freeze

RϕR\_\{\\phi\}
Stage II: Train hierarchical prompts

Initialize

PcP\_\{c\}and

\{Pk\}k=1K\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\}
for

e=1,…,Epe=1,\\ldots,E\_\{p\}do

foreach paired coarse/fine preference batchdo

Compute

\{αk​\(x\)\}k=1K\\\{\\alpha\_\{k\}\(x\)\\\}\_\{k=1\}^\{K\}and

Pf​\(x\)P\_\{f\}\(x\)
Update only

PcP\_\{c\}using

ℒDPO\\mathcal\{L\}\_\{\\mathrm\{DPO\}\}and

𝒞~P\(c\)​\(x\)\\widetilde\{\\mathcal\{C\}\}\_\{P\}^\{\(c\)\}\(x\)
Update only

\{Pk\}k=1K\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\}using

ℒDPO\\mathcal\{L\}\_\{\\mathrm\{DPO\}\}and

𝒞~P\(f\)​\(x\)\\widetilde\{\\mathcal\{C\}\}\_\{P\}^\{\(f\)\}\(x\)
endfor

endfor

Input\-adaptive inference

Compute

psafe​\(x\)p\_\{\\mathrm\{safe\}\}\(x\)with

RϕR\_\{\\phi\}
if

psafe​\(x\)\>τsp\_\{\\mathrm\{safe\}\}\(x\)\>\\tau\_\{s\}then

y∼πθ​\(y∣𝒞0chat​\(x\)\)y\\sim\\pi\_\{\\theta\}\(y\\mid\\mathcal\{C\}\_\{0\}^\{\\mathrm\{chat\}\}\(x\)\)

else

Construct

𝒞Pchat​\(x\)\\mathcal\{C\}\_\{P\}^\{\\mathrm\{chat\}\}\(x\)and sample

y∼πθ​\(y∣𝒞Pchat​\(x\)\)y\\sim\\pi\_\{\\theta\}\(y\\mid\\mathcal\{C\}\_\{P\}^\{\\mathrm\{chat\}\}\(x\)\)
endif

return

Rϕ,Pc,\{Pk\}k=1K,yR\_\{\\phi\},P\_\{c\},\\\{P\_\{k\}\\\}\_\{k=1\}^\{K\},y

## Appendix BAdditional Analyses and Ablation Studies

### B\.1Effectiveness of Routing and the Coarse\-Grained Prompt

Table 5:Effectiveness of input\-dependent routing and the shared coarse\-grained prompt on Mistral\-7B\-Instruct\-v0\.3\. Each benchmark entry reports Safety \(%\)/Helpfulness \(0–10\)\. Uniform assigns equal weights to all fine\-grained prompt experts, whereas Routed uses the input\-dependent weights predicted by the router\. HiRoute combines the routed fine\-grained prompt with the shared coarse\-grained prompt\. Best results are shown inbold, and second\-best results are underlined\.Table[5](https://arxiv.org/html/2608.12821#A2.T5)isolates the effects of input\-dependent routing and the shared coarse\-grained prompt\. Compared with uniform expert averaging, learned routing consistently improves safety by 4\.8, 4\.0, and 5\.2 percentage points on StrongReject, AdvBench, and JailbreakBench, respectively, while increasing helpfulness by 0\.2–0\.4 points\. Consequently, average safety rises from 65\.2% to 69\.8%, and average helpfulness from 7\.1 to 7\.4\. The consistent gains across all three benchmarks indicate that the predicted risk distribution provides more informative expert composition than an input\-independent uniform mixture\. However, routed fine\-grained prompting alone achieves only 69\.8% average safety, suggesting that category\-specific guidance is insufficient to establish a robust safety boundary\. Incorporating the shared coarse\-grained prompt produces much larger safety gains of 21\.5, 31\.0, and 17\.5 percentage points on the three benchmarks, raising average safety to 93\.2%\. The largest improvement occurs on AdvBench, where the shared constraint is particularly important for adversarially phrased harmful requests\. Crucially, average helpfulness remains at 7\.4: it is unchanged on StrongReject, decreases by only 0\.1 on AdvBench, and increases by 0\.2 on JailbreakBench\. Thus, the safety gain does not arise from reverting to uniformly generic refusals\. These results demonstrate complementary roles: input\-dependent routing improves category\-specific expert composition, while the coarse\-grained prompt provides a stable cross\-category safety boundary without sacrificing safe\-response helpfulness\.

### B\.2Hierarchical Router Evaluation

Table[6](https://arxiv.org/html/2608.12821#A2.T6)reports the intrinsic performance of the hierarchical router on the held\-out test sets\. Across the three backbones, coarse\-grained accuracy and F1 range from 84\.21% to 89\.77% and from 84\.18% to 89\.51%, respectively, while risk recall remains between 80\.26% and 89\.47%\. These results indicate reliable safe/risk discrimination\. Fine\-grained multi\-label recognition is more challenging and varies across backbone representations: Vicuna achieves the highest Micro\-F1 and Macro\-F1 of 80\.48% and 80\.81%, followed by Zephyr, while Mistral obtains 68\.31% and 69\.81%\. The relatively close Micro\-F1 and Macro\-F1 scores suggest that performance is not dominated by a small number of frequent categories\. These results also support the shared coarse\-grained prompt, which retains a category\-agnostic safety boundary when fine\-grained expert matching is imperfect\.

Table 6:Test\-set performance of the hierarchical routers\. Coarse Acc\., Risk Recall, and Coarse F1 evaluate binary safe/risk classification, while Fine Micro\-F1 and Macro\-F1 evaluate fine\-grained multi\-label risk recognition\. All metrics are reported as percentages\.
### B\.3Effect of Total Soft\-Prompt Length

Table 7:Effect of the total prompt length\. Each benchmark entry reports Safety \(%\)/Helpfulness \(0–10\)\. The shaded row denotes the default configuration\. Best results are shown inbold, and second\-best results are underlined\.Table[7](https://arxiv.org/html/2608.12821#A2.T7)examines how the prompt length affects safety and safe\-response helpfulness\. Increasing the prompt length from 5 to 20 improves the average safety rate from 78\.1% to 93\.2%, an absolute gain of 15\.1 percentage points, while maintaining an average helpfulness score of 7\.4\. This indicates that a moderate increase in prompt capacity enables more effective safety control without sacrificing response quality\. However, further increasing the length to 30 and 40 reduces the average safety rate to 83\.2% and 81\.6%, respectively, without improving helpfulness\. This non\-monotonic trend shows that additional prompt parameters do not necessarily yield stronger safety control\. Excessively long prompts may introduce redundant optimization directions and make prompt learning more difficult with limited preference data\. Overall, a prompt length of 20 provides the most favorable balance between safety and safe\-response helpfulness and is therefore adopted as the default setting\.

### B\.4Effect of the Coarse/Fine Update Ratio

Table[8](https://arxiv.org/html/2608.12821#A2.T8)studies the effect of the update ratio between the coarse\- and fine\-grained prompts\. As the proportion of coarse\-grained updates increases from1:31\{:\}3to2:12\{:\}1, the average safety rate steadily improves from 75\.0% to 93\.2%, while average helpfulness remains at 7\.4\. This result indicates that strengthening the shared coarse\-grained prompt reinforces the cross\-category safety boundary without necessarily degrading fine\-grained response quality\. Increasing the ratio further to3:13\{:\}1yields only a 0\.7\-percentage\-point improvement in average safety, but reduces average helpfulness from 7\.4 to 6\.3\. The helpfulness scores on the three safety benchmarks decrease by 0\.9, 1\.2 and 1\.2 points, respectively\. Overemphasizing coarse\-grained updates therefore encourages more conservative safety behavior, weakening the ability of the fine\-grained experts to provide targeted explanations and safe alternatives\. The2:12\{:\}1ratio lies at the point where safety gains begin to saturate but response helpfulness has not yet deteriorated, making it the most appropriate update configuration\.

Table 8:Effect of the coarse/fine prompt update ratio\. Each benchmark entry reports Safety \(%\)/Helpfulness \(0–10\)\. The shaded row denotes the default configuration\. Best results are shown inbold, and second\-best results are underlined\.
### B\.5Scaling to a Larger Backbone

Table 9:Router performance on the held\-out test set for Vicuna\-13B\-v1\.5\. All values are reported as percentages\.##### Router Performance\.

Table[9](https://arxiv.org/html/2608.12821#A2.T9)reports the router results on the held\-out test set for Vicuna\-13B\-v1\.5\. The coarse\-grained gate achieves an accuracy of 88\.89%, a harmful\-input recall of 81\.58%, and an F1 score of 86\.71%\. The fine\-grained router obtains Micro\-F1 and Macro\-F1 scores of 76\.15% and 76\.36%, respectively\. The comparable Micro\-F1 and Macro\-F1 scores indicate relatively balanced multi\-label prediction across the risk categories\. These results show that the hierarchical router remains effective when the backbone is scaled to 13B parameters\.

##### Safety and Utility at a Larger Scale\.

As shown in Table[11](https://arxiv.org/html/2608.12821#A2.T11), HiRoute remains effective on Vicuna\-13B\-v1\.5\. It improves average safety from 92\.2% to 97\.7% and average safe\-response helpfulness from 5\.6 to 5\.9 across the three safety benchmarks\. Meanwhile, GSM8K accuracy changes from 10\.0% to 10\.5%, and XSTest over\-refusal increases only slightly from 1\.7% to 2\.0%\. Although MT\-Bench and TruthfulQA show minor declines, the overall results indicate that HiRoute preserves its safety and helpfulness benefits on a larger backbone with limited changes in general utility\.

Table 10:Safety rates under transfer\-based GCG attacks\. Adversarial suffixes are optimized against each base model and transferred to the corresponding HiRoute model\. Improvements are absolute percentage\-point gains\.Table 11:Scaling results on Vicuna\-13B\-v1\.5\. Results on each safety benchmark are reported as Safety/Helpfulness, and Avg\. S/H denotes their average across the three benchmarks\. Safety, GSM8K accuracy, XSTest over\-refusal, and TruthfulQA scores are reported as percentages, while helpfulness is scored from 0 to 10\. T and I denote truthfulness and informativeness, respectively\. Best results are shown inbold\.

## Appendix CRobustness under Transferred GCG Attacks

Table[10](https://arxiv.org/html/2608.12821#A2.T10)evaluates the robustness of HiRoute against transfer\-based GCG attacks\. Adversarial suffixes are first optimized against each base model and then transferred to the corresponding HiRoute model\. The safety rates of the attacked base models vary substantially, ranging from 19\.7% to 54\.7%\. After applying HiRoute, all three safety rates exceed 87\.5%, and the average safety rate increases from 34\.6% to 88\.3%, corresponding to an average absolute improvement of 53\.6 percentage points\. Zephyr has the lowest initial safety rate but achieves the largest improvement of 68\.3 percentage points, while HiRoute further improves the relatively stronger Vicuna model to 89\.3%\. The consistent gains across backbones indicate that HiRoute does not merely inherit the original refusal tendencies of the base models; instead, its shared safety constraint provides a stable defensive foundation across different backbones\. This experiment evaluates transferred attacks rather than adaptive white\-box attacks directly optimized against the HiRoute router and prompt parameters\. The results therefore support robustness against transfer\-based GCG attacks but should not be interpreted as comprehensive robustness against all forms of GCG\.

## Appendix DRepresentative Safety\-Response Comparisons

We provide qualitative examples to illustrate the behavioral differences among the base model, the model using only the shared coarse\-grained prompt, and HiRoute\. Given the same harmful request, the base model may provide actionable harmful information despite an initial disclaimer, whereas the coarse\-grained prompt tends to produce a generic refusal\. HiRoute maintains the safety boundary while redirecting the request toward specific and actionable protective guidance\. These examples complement the aggregate safety and helpfulness results reported in the main paper\.

Table 12:Qualitative comparison on a cyberattack request\. The baseline provides unsafe procedural details, the coarse\-grained prompt gives a generic refusal, and HiRoute redirects the request toward targeted preventive guidance\.Table 13:Qualitative comparison on an insider\-trading request\. The baseline provides unsafe procedural information, the coarse\-grained prompt refuses without further guidance, and HiRoute explains the risk and provides compliance\-oriented alternatives\.
## Appendix ELLM\-as\-a\-Judge Evaluation Prompts

We employ an LLM\-as\-a\-judge protocol to evaluate both response safety and safe\-response helpfulness\. The evaluation is conducted in two stages\. We first classify each response as safe or unsafe\. Responses classified as safe are then evaluated for helpfulness\. The exact evaluation prompts and scoring criteria are provided below\. To verify the reliability of the automatic evaluation, we also manually assess a randomly sampled subset of 20 responses across models and benchmarks using the same criteria, and observe overall consistency with the GPT\-5\.4 judgments\.

### E\.1Safety Evaluation Prompt

We use an LLM\-based judge to determine whether each generated response is safe\. The corresponding user request is provided only as context, while the classification is based exclusively on the content of the response\. The judge returns a binary label,*safe*or*unsafe*, according to the criteria presented in Table[14](https://arxiv.org/html/2608.12821#A5.T14)\.

Table 14:Safety\-evaluation rubric used to classify model responses\. The evaluator considers the user prompt only as context and determines whether the response itself is safe or unsafe\.
### E\.2Safe\-Response Helpfulness Evaluation Prompt

For responses classified as safe, we further evaluate whether they remain useful under the safety constraints\. The judge assigns a score from 1 to 10 according to the response’s relevance, constructiveness, safe alternatives, explanation of safety boundaries, and support for safe user objectives\. The complete scoring rubric is presented in Table[15](https://arxiv.org/html/2608.12821#A5.T15)\.

Table 15:Safe\-response helpfulness rubric used by the LLM\-based evaluator\. Only responses previously classified as safe are evaluated\. Higher scores indicate that a response provides more specific, actionable, and context\-sensitive assistance while maintaining appropriate safety boundaries\.

Similar Articles

No One Fits All: From Fixed Prompting to Learned Routing in Multilingual LLMs

arXiv cs.CL

Researchers from National Taiwan University propose replacing fixed translation-based prompting strategies in multilingual LLMs with lightweight learned classifiers that route each instance to either native or translation-based prompting. Their analysis across 10 languages and 4 benchmarks shows no single strategy is universally optimal, with translation benefiting low-resource languages most, and the learned routing achieving statistically significant improvements over fixed strategies.

Safe responses matter: Output-aware safety guardrail mitigate over-refusal in MLLMs

arXiv cs.LG

This paper proposes output-aware safety guardrails for multimodal large language models that use hidden state representations and multi-instance contrastive learning to predict unsafe outputs before generation, drastically reducing over-refusal while maintaining safety. The method preserves the model's utility by intervening only when the actual response would be harmful.