Guiding Giants: Lightweight Controllers for Weighted Activation Steering in LLMs

Titel: Guiding Giants: Lightweight Controllers for Weighted Activation Steering in LLMs

Auteurs: Amr Hegazy, Mostafa Elhoushi, Amr Alanwar.

---

1. Het Probleem

Grote Taalmodellen (LLMs) zijn kwetsbaar voor misbruik, zoals het genereren van schadelijke, vooroordeel-gebaseerde of onjuiste inhoud. Bestaande methoden om deze risico's te mitigeren, hebben aanzienlijke nadelen:

• Fine-tuning: Vereist kostbare herscholing, kan leiden tot "catastrophic forgetting" (verlies van algemene vaardigheden) en is resource-intensief.
• Bestaande inferentie-tijd sturing (Steering): Methoden zoals Activation Addition of Inference-Time Intervention (ITI) passen vaak vaste aanpassingen toe over alle lagen van het model. Dit gebrek aan fijnmazige controle leidt tot twee problemen:
  1. Als de sturing te sterk is, weigert het model ook onschadelijke vragen (verlies van nut).
  2. Als de sturing te zwak is, faalt het om schadelijke verzoeken te blokkeren.
• Gebrek aan adaptiviteit: Bestaande technieken passen zich niet dynamisch aan op basis van de specifieke context van een prompt of de interne staat van het model.

Het doel is een methode te vinden die het gedrag van een LLM tijdens de inferentie (tijdens het genereren) veilig kan sturen zonder de modelparameters te wijzigen, terwijl de nuttigheid voor veilige vragen behouden blijft.

---

2. Methodologie: Weighted Activation Steering (WAS)

De auteurs introduceren Weighted Activation Steering (WAS), een lichtgewicht, trainbaar controller-netwerk dat dynamisch ingrijpt op de activaties van het LLM tijdens de inferentie.

Kerncomponenten:

1. Stuurvector ($d_{steer}$):

   • Een vooraf berekende vector die de "weigerrichting" (refusal direction) in de embedding-ruimte voorstelt.
   • Deze wordt berekend door het gemiddelde van embeddings van weiger-tokens (bijv. "I cannot", "I apologize") af te trekken van het gemiddelde van antwoord-tokens (bijv. "Sure", "Here").
   • Formule: $d_{steer} = \frac{\bar{e}_{refuse} - \bar{e}_{answer}}{||\bar{e}_{refuse} - \bar{e}_{answer}||_2}$.

2. Het Controller Netwerk ($f_c$):

   • Een lichtgewicht Multi-Layer Perceptron (MLP) die de interne activaties van het gefrozen LLM observeert.
   • Input: Geconcateneerde activaties van de laatste derde van de lagen ($L_{input}$) op het laatste token van de prompt.
   • Output: Een tuple bestaande uit:
     • Een scalar magnitude ($s$): Een globale schaalfactor die bepaalt hoe sterk er moet worden ingegrepen (hoog voor schadelijke prompts, laag voor veilige).
     • Een vector van laag-gewichten ($w$): Een reeks gewichten (tussen 0 en 1) voor elke laag van het model, die bepaalt hoe sterk de stuurvector in die specifieke laag wordt toegepast.

3. Toepassing van de Patch:

   • Tijdens de inferentie wordt de modifiatie op de verborgen staat ($h_{l,t}$) van laag $l$ toegepast als volgt:
     $$h'_{l,t} = h_{l,t} + (s \cdot w_l \cdot \alpha_{global} \cdot d_{steer})$$
   • Waarbij $\alpha_{global}$ een vaste hyperparameter is voor de totale intensiteit.

4. Trainingsproces:

   • Het controller-netwerk wordt getraind op een dataset van schadelijke (harmful) en veilige (benign) prompts.
   • Doel: Het netwerk moet een hoge $s$ voorspellen voor schadelijke prompts en een lage $s$ voor veilige prompts.
   • Verliesfunctie: Mean Squared Error (MSE) tussen de voorspelde $s$ en de doelen (1.0 voor schadelijk, 0.0 voor veilig).
   • Interessant: Hoewel alleen $s$ expliciet wordt gesuperviseerd, leert het gedeelde verborgen laag van het netwerk via backpropagation ook om gestructureerde, niet-uniforme gewichten ($w$) te genereren die specifiek zijn voor de inhoud van de prompt.

---

3. Belangrijkste Bijdragen

• Adaptieve, Laag-specifieke Sturing: In tegenstelling tot eerdere methoden die een vaste sturing toepassen, leert WAS om de sterkte en de locatie (welke lagen) van de ingreep dynamisch aan te passen op basis van de input.
• Lichtgewicht Architectuur: De controller is een klein MLP-netwerk dat geen herscholing van het basis-LLM vereist, wat het zeer efficiënt maakt voor productieomgevingen.
• Discriminatie tussen Schadelijk en Veilig: De methode is specifiek ontworpen om schadelijke verzoeken te weigeren zonder de prestaties op veilige, nuttige taken te degraderen.
• Implementatie: De auteurs tonen aan dat dit efficiënt kan worden geïmplementeerd met PyTorch forward hooks zonder de broncode van het basismodel aan te passen.

---

4. Resultaten

De methode is geëvalueerd op modellen zoals Llama-3.1-8B, Llama-3.2-1B en Mistral-7B.

• Veiligheid (Refusal Rates):

  • WAS presteert aanzienlijk beter dan bestaande methoden (CAST, ACT, ITI) op veiligheidsbenchmarks zoals ToxicChat, Jailbreak Prompts en AdvBench.
  • Op ToxicChat bereikte WAS een weigerpercentage van 93,0% (Llama-3.1-8B), terwijl concurrenten vaak onder de 50% bleven.
  • Op AdvBench werd een score van 98,8% behaald.

• Behoud van Algemene Vaardigheden:

  • Op benchmarks voor algemene nuttigheid (AlpacaEval) bleef het winpercentage tegen het basismodel statistisch ononderscheidbaar van 50% (bijv. 49,82% voor Llama-3.1-8B).
  • Prestaties op redeneringstaken (GSM8K) en kennis (MMLU, HellaSwag) bleven nagenoeg ongewijzigd, wat aantoont dat de "nuttigheid" van het model niet is aangetast.

• Efficiëntie:

  • WAS voegt een minimale latentie toe aan de inferentie (token-per-tijd), wat aanzienlijk sneller is dan methoden zoals ITI, terwijl het wel betere veiligheidsresultaten levert.

• Ablatie Studies:

  • Het gebruik van laag-specifieke gewichten ($w$) bleek cruciaal. Een variant die alleen de scalar $s$ gebruikte met uniforme gewichten, resulteerde in een significante daling van de weigerpercentages (bijv. van 93% naar 71,8% op ToxicChat).
  • De controller leerde dat bepaalde lagen (vooral midden- en late lagen) kritischer zijn voor het detecteren van schadelijke inhoud dan andere.

---

5. Betekenis en Conclusie

Dit paper introduceert een nieuwe paradigma voor LLM-veiligheid: inference-time control via een trainbare, adaptieve controller.

• Praktische Toepassing: WAS biedt een kosteneffectief alternatief voor fine-tuning. Het kan worden ingezet als een "safety layer" voor reeds getrainde modellen om ze direct veiliger te maken zonder de onderliggende architectuur te wijzigen.
• Fijnmazige Controle: Het paper bewijst dat het dynamisch aanpassen van de stuurkracht per laag en per prompt essentieel is om de balans tussen veiligheid en nuttigheid te vinden.
• Beperkingen: De methode is afhankelijk van de kwaliteit van de vooraf berekende stuurvector en kan kwetsbaar zijn voor geavanceerde adversariale aanvallen die specifiek zijn ontworpen om de controller te misleiden (bijv. prompts die schadelijk zijn maar activatiepatronen hebben die lijken op veilige prompts).

Kortom, WAS biedt een robuuste, efficiënte en aanpasbare oplossing om LLM's te "sturen" naar veiliger gedrag, waarbij de modelcapaciteiten voor legitieme taken intact blijven.