Technische Samenvatting: HiRoute

Probleemstelling

Grote taalmodellen (LLM's) blijven kwetsbaar voor schadelijke verzoeken, jailbreak-aanvallen en queries met meerdere risico's. Hoewel Supervised Fine-Tuning (SFT) en Reinforcement Learning from Human Feedback (RLHF) veiligheidsgedrag kunnen coderen, brengen ze hoge trainings- en opslagkosten met zich mee en leiden ze vaak tot een afruil tussen veiligheid en behulpzaamheid. Specifiek kan het uniform versterken van weigeringsgedrag leiden tot "over-refusal" (overmatige weigering) bij onschuldige inputs, wat de algemene capaciteiten verslechtert.

Bestaande parameter-efficiënte methoden voor veiligheidsafstemming, zoals prompt tuning, vertrouwen doorgaans op statische mechanismen. Deze benaderingen vallen over het algemeen uiteen in twee categorieën met complementaire beperkingen:

1. Grofmazige globale prompts: Een enkele veiligheidsprompt die op alle risicovolle inputs wordt toegepast, vestigt een sterke veiligheidsgrens maar resulteert vaak in generieke, oninformatieve weigeringen die geen risicospecifieke begeleiding of veilige alternatieven bieden.
2. Fijnmazige gerouteerde prompts: Methoden die categorie-specifieke prompts toewijzen aan verschillende risico's kunnen meer gerichte en behulpzame veilige reacties genereren. Ze hebben echter moeite om een stabiele veiligheidsgrens te handhaven, vooral wanneer inputs niet precies overeenkomen met vooraf gedefinieerde risicocategorieën of wanneer de expert-matching onvoldoende is.

De kernuitdaging is hoe men de fijnmazige constructiviteit van reacties kan balanceren met cross-categorie veiligheid door middel van input-adaptieve prompt-compositie, zonder het volledige backbone-model te bevriezen.

Methodologie: HiRoute

De auteurs stellen HiRoute (Hierarchical Routed Prompt Tuning) voor, een input-adaptief framework dat categorie-agnostische veiligheidscontrole scheidt van categorie-specifieke responsbegeleiding. Het framework werkt op een bevroren instruction-tuned LLM en bestaat uit drie hoofdonderdelen:

1. Hiërarchische Risico-routing

Een lichtgewicht router wordt getraind op representaties geëxtraheerd uit de bevroren LLM om twee taken uit te voeren:

• Grofmazige Detectie: Bepaalt of een input schadelijk of onschuldig is ($p_{safe}$ vs. $p_{risk}$).
• Fijnmazige Risicovoorspelling: Voorspelt multi-label risicoscores voor specifieke categorieën (bijv. cybercriminaliteit, privacy schendingen, economische criminaliteit, geweld).
  De router geeft een waarschijnlijkheidsverdeling die wordt gebruikt om fijnmazige prompt-experts dynamisch te wegen.

2. Hiërarchische Prompt-compositie

HiRoute leert twee soorten continue prompt-embeddings:

• Gedeelde grofmazige prompt ($P_c$): Biedt een categorie-agnostische veiligheidsgrens die van toepassing is op alle risicovolle inputs.
• Fijnmazige Prompt Experts ($\{P_k\}$): Een set van $K$ experts, die elk overeenkomen met een specifieke risicocategorie.
  Voor een gegeven input $x$ wordt de fijnmazige prompt $P_f(x)$ geconstrueerd als een gewogen som van de experts op basis van de door de router voorspelde risicoscores:
  $$P_f(x) = \sum_{k=1}^{K} \alpha_k(x) P_k$$
  waarbij $\alpha_k(x)$ genormaliseerde gewichten zijn. De uiteindelijke input-afhankelijke prompt is $P(x) = [P_f(x); P_c]$.

3. Tweestaps Ontkoppelde Optimalisatie

De training verloopt in twee afzonderlijke fasen om te voorkomen dat het generatie-doel de risicobeslissingsgrens verandert:

• Fase I (Router Leren): De router wordt getraind met Cross-Entropy (CE) voor grofmazige classificatie en Binary Cross-Entropy (BCE) voor fijnmazige multi-label classificatie. Het backbone-model blijft bevroren.
• Fase II (Prompt Optimalisatie): Het backbone-model en de getrainde router zijn bevroren. Alleen de prompt-parameters ($P_c$ en $\{P_k\}$) worden geoptimaliseerd met behulp van Direct Preference Optimization (DPO) op voorkeurstriplets (veilig/behulpzaam vs. onveilig/lage kwaliteit). Om interferentie te verminderen, gebruiken de auteurs alternerende gradiënt-updates: wanneer $P_c$ wordt bijgewerkt, worden de gradiënten voor $P_f(x)$ gemaskeerd (stop-gradient), en vice versa.

4. Input-Adaptieve Inferentie

Tijdens de inferentie berekent de router eerst de veiligheidswaarschijnlijkheid $p_{safe}(x)$.

• Onschuldige Inputs: Als $p_{safe}(x) > \tau_s$ (een veiligheidsdrempel), omzeilt het systeem alle veiligheidsprompts en genereert het direct vanuit de bevroren backbone, waardoor de algemene bruikbaarheid behouden blijft.
• Risicovolle Inputs: Als $p_{safe}(x) \leq \tau_s$, activeert het systeem de veiligheidstak en construeert het de hiërarchische prompt $P(x)$ om de generatie te conditioneren.

Belangrijkste Bijdragen

1. Empirische Validatie van Afruilen: De auteurs laten zien dat gedeelde grofmazige prompts een sterkere veiligheid bieden maar een lagere behulpzaamheid, terwijl gerouteerde fijnmazige prompt-mengsels de behulpzaamheid verbeteren maar minder veiligheidswinst opleveren.
2. HiRoute Framework: Ze stellen een nieuw architectuur voor die een gedeelde grofmazige prompt (voor een stabiele veiligheidsgrens) combineert met dynamisch gerouteerde fijnmazige experts (voor risicospecifieke begeleiding). Dit ontwerp voorkomt onnodige interventie op onschuldige inputs via een veiligheidshek (safety gate).
3. Uitgebreide Evaluatie: De methode wordt gevalideerd over drie instruction-tuned modellen (Mistral-7B, Vicuna-7B, Zephyr-7B) en meerdere benchmarks, waarbij wordt aangetoond dat het een hoge veiligheid en behulpzaamheid behoudt terwijl de algemene bruikbaarheid intact blijft.

Experimentele Resultaten

Experimenten werden uitgevoerd op drie modellen tegenover baselines inclusief Base modellen, RPO (Robust Prompt Optimization), DRO (Dynamic Risk Optimization) en ACD (Adversarial Contrastive Decoding).

• Veiligheid en Behulpzaamheid: HiRoute behaalde de hoogste gemiddelde veiligheidspercentages over alle modellen (bijv. 93,2% op Mistral, 97,7% op Vicuna, 94,8% op Zephyr) en bereikte tegelijkertijd de hoogste of bijna hoogste scores voor de behulpzaamheid van veilige reacties. Het presteerde 1,3–2,9 procentpunt beter dan de sterkste baselines op het gebied van veiligheid.
• Algemene Bruikbaarheid: In tegen tegenstelling tot methoden die algemene capaciteiten verslechteren, behield HiRoute de prestaties op GSM8K (wiskundig redeneren), MT-Bench (instructieopvolging) en TruthfulQA. Op Mistral verhoogde het bijvoorbeeld de veiligheid van 53,2% naar 93,2% terwijl 48,0% GSM8K-nauwkeurigheid behouden bleef.
• Over-refusal: HiRoute verminderde over-refusal op onschuldige inputs aanzienlijk. Op Mistral behaalde het een over-refusal rate van 2,5%, een enorme verbetering ten opz wzglęken met RPO (40,5%) en DRO (18,0%).
• Ablatie-studies:
  • Trainingsstrategie: Gezamenlijke alternerende training (het bijwerken van prompts in een consistente samengestelde context) presteerde aanzienlijk beter dan sequentiële training ("Coarse na fine"), waarbij de gemiddelde veiligheid steeg van 74,8% naar 93,2%.
  • Prompt Lengte: Een 15/5 split tussen grove en fijnmazige prompt lengtes (totaal 20) bood de beste balans tussen veiligheid en behulpzaamheid.
  • Drempelwaarde: Het verhogen van de veiligheidsgating-drempel $\tau_s$ versterkte de veiligheidscontrole met minimale impact op de algemene capaciteit.

Betekenis en Claims

Het artikel claimt dat HiRoute effectief de spanning tussen veiligheid en behulpzaamheid in LLM-afstemming aanpakt. Door risico-identificatie te ontkoppelen van responsgeneratie en een hiërarchische, input-adaptieve promptstructuur te gebruiken, bereikt de methode het volgende:

• Vestigt een stabiele, categorie-agnostische veiligheidsgrens zonder te vertroukken op agressieve, uniforme weigeringen.
• Biedt constructieve, risicospecifieke begeleiding en veilige alternatieven voor schadelijke verzoeken.
• Minimaliseert over-refusal door onschuldige inputs de mogelijkheid te geven de veiligheidsprompt volledig te omzeilen.
• Bereikt deze doelen met parameter-efficiëntie, waarbij het backbone-model bevroren blijft en alleen een klein aantal prompt-parameters en een lichtgewicht router wordt geoptimaliseerd.

De auteurs concluderen dat input-afhankelijke hiërarchische prompting een robuuste oplossing biedt voor veiligheidsafstemming die veiligheid, responskwaliteit, parameter-efficiëntie en algemene bruikbaarheid in balans brengt. Ze erkennen beperkingen met betrekking tot de nauwkeurigheid van de router bij out-of-distribution samengestelde risico's en de huidige focus op single-turn tekstinteracties, wat wijst op toekomstig werk in open-set risicorecognitie en multimodale instellingen.