Pruning Unsafe Tickets: A Resource-Efficient Framework for Safer and More Robust LLMs

Probleemstelling

Hoewel grote taalmodellen (LLMs) en vision-language modellen (VLMs) zoals Mistral en LLaVA indrukwekkende prestaties leveren, vertonen ze vaak nog steeds onveilig gedrag dat is overgeërfd van het vooraf trainen (pre-training). Bestaande methoden voor uitlijning (alignment), zoals Supervised Fine-Tuning (SFT) en Reinforcement Learning from Human Feedback (RLHF), richten zich voornamelijk op het genereren van gewenste antwoorden, maar verwijderen niet expliciet de onderliggende neurale subnetwerken die schadelijke output veroorzaken.

Dit leidt tot twee belangrijke problemen:

1. Fragiliteit: Onveilig gedrag kan opnieuw optreden bij adversariële prompts (bijv. jailbreaks) of distributieveranderingen.
2. Resource-intensiteit: Bestaande methoden vereisen vaak enorme datasets, veel GPU-geheugen en lange trainingstijden. Lichtere alternatieven, zoals prompt engineering of interne modelinterventies, zijn vaak niet schaalbaar of intrinsiek onvoldoende om onveiligheid uit het model zelf te verwijderen.

Methodologie

De auteurs introduceren een resource-efficiënt pruning-framework dat direct parameters identificeert en verwijdert die geassocieerd zijn met onveilig gedrag, terwijl de nuttige functionaliteit van het model behouden blijft. Het framework baseert zich op de Lottery Ticket Hypothesis (LTH), waarbij wordt verondersteld dat modellen "onveilige tickets" (dunne subnetwerken die schadelijk gedrag veroorzaken) bevatten die verwijderd kunnen worden om "veilige tickets" bloot te leggen.

Het proces bestaat uit vier fasen:

1. Gedragsprofilering (Behavior Profiling):

   • Er wordt een model-specifiek dataset gegenereerd door het model te bevragen met diverse prompts.
   • Antwoorden worden automatisch gelabeld als "veilig" of "onveilig" met behulp van een externe classifier (bijv. LLaMA-Guard).
   • Om redundantie te minimaliseren, worden antwoorden geclusterd en representatieve paren geselecteerd. Een vaste responslengte wordt gehandhaafd om kwaliteitsverlies te voorkomen.

2. Attributie-analyse (Attribution Analysis):

   • In plaats van dure backpropagation, gebruiken de auteurs een gradient-free attributiemethode gebaseerd op Wanda.
   • Ze passen een binair masker toe om alleen de activaties tijdens de respons (niet de prompt) te tellen. Dit isoleert de parameters die specifiek bijdragen aan het genereren van de output.
   • De score voor een parameter wordt berekend als het product van de gewichtsgrootte en de L2-norm van de bijbehorende activaties.

3. Component Scoring:

   • Er wordt een contrastieve belangrijkheidsscore berekend voor elke modelcomponent (bijv. attention-heads, MLP-lagen).
   • De score $I$ is de verhouding tussen de attributiescores voor onveilige data ($S'_u$) en veilige data ($S'_s$): $I = S'_u / (S'_s + \epsilon)$.
   • Na normalisatie (z-score) worden de parameters met de hoogste scores geselecteerd, aangezien deze disproportioneel actief zijn tijdens onveilige generaties.

4. Iteratief Pruning:

   • In plaats van alles in één keer te verwijderen, wordt iteratief gepreemd. Twee strategieën worden onderzocht:
     • Greedy Pruning: Verwijdert bij elke iteratie de component met de hoogste score. Dit is snel en efficiënt.
     • Beam Search Pruning: Onderzoekt meerdere pruning-paden tegelijkertijd om een optimale balans te vinden tussen veiligheid en bruikbaarheid, gebaseerd op een doelwitfunctie die veilige data behoudt en onveilige data straft.

Belangrijkste Bijdragen

• Resource-efficiënt Framework: Een post-hoc uitlijningsmethode die geen fijnafstemming (fine-tuning) of prompt engineering vereist, en werkt op zowel LMs als VLMs, inclusief gekwantiseerde varianten.
• Nieuw Perspectief (LTH): Het paper koppelt pruning aan de Lottery Ticket Hypothesis, bewijzend dat het verwijderen van "onveilige tickets" leidt tot een subnetwerk dat veilig is zonder prestatieverlies.
• Gradient-vrij: De methode vereist geen backpropagation, wat het toepasbaar maakt op modellen waarvoor geen volledige training mogelijk is.
• Schaalbaarheid: Werkt effectief op verschillende architecturen (Mistral, LLaMA, Gemma, Qwen) en modali teiten (tekst en beeld).

Resultaten

De methode is geëvalueerd op modellen zoals Mistral-7B, Mistral-Nemo, Mistral-Small en LLaVA-v1.6.

• Veiligheid:
  • Bij Mistral-7B daalde het percentage onveilige antwoorden van 22,8% naar 1,17% (met Beam Search).
  • Bij LLaVA-v1.6-Mistral-7B daalde dit naar 2%.
  • De methode is effectief tegen jailbreak-aanvallen (GCG, AutoDAN, PAIR, TAP), waarbij de aanvalssuccespercentages aanzienlijk dalen zonder de inferentie-efficiëntie te verliezen.
• Bruikbaarheid (Utility):
  • In tegenstelling tot methoden zoals Circuit Breakers (CB) die vaak leiden tot "over-refusal" (het weigeren van veilige vragen), behoudt het pruning-framework een hoge bruikbaarheidsscore (Utility ~7.0-8.0).
  • De over-refusal rate blijft redelijk (rond de 46% voor Beam vs. 100% voor CB), wat betekent dat het model nog steeds nuttige antwoorden geeft op veilige prompts.
• Efficiëntie:
  • Het pruningproces is zeer snel: het verwijderen van parameters op Mistral-7B duurt slechts 455 seconden (Greedy) met een piekgeheugengebruik van 18 GB.
  • Dit is aanzienlijk sneller en minder geheugenintensief dan fine-tuning methoden zoals DPO of CB.
  • Er is geen extra inferentie-tijd of token-overhead tijdens het gebruik van het gepreemde model.

Significantie

Dit paper biedt een lichtgewicht, schaalbare oplossing voor het veiligstellen van AI-modellen in resource-beperkte omgevingen. Door onveilige subnetwerken fysiek te verwijderen in plaats van ze alleen te maskeren of te overstemmen, creëert het framework een intrinsiek robuuster model.

De bevindingen suggereren dat veiligheid en bruikbaarheid niet noodzakelijk in conflict hoeven te zijn; door de juiste "onveilige tickets" te verwijderen, kan een model zowel veiliger als efficiënter worden. Dit maakt de techniek ideaal voor de implementatie van veilige AI in productieomgevingen waar rekenkracht en geheugen beperkt zijn, en waar snelle, effectieve uitlijning na het vooraf trainen vereist is.