Accelerating Suffix Jailbreak attacks with Prefix-Shared KV-cache

Probleemstelling

Suffix-jailbreak-aanvallen zijn een systematische methode om Large Language Models (LLM's) te testen (red-teaming) door kwaadaardige instructies te omzeilen via gegenereerde achtervoegsels (suffixen). Huidige methoden, zoals heuristische zoekalgoritmen (bijv. GCG, BEAST) of LLM-gebaseerde generators, lijden onder een prohibitief hoge rekenkosten.

De kern van het probleem ligt in de iteratieve aard van deze aanvallen:

1. Er moeten duizenden kandidaat-suffixen worden geëvalueerd tegen een doel-LLM.
2. Elk van deze kandidaten bestaat uit een identieke, lange kwaadaardige instructie (prefix) gevolgd door een variërende suffix.
3. Standaard inferentie-technieken (inclusing standaard KV-cache implementaties) behandelen elke kandidaat als een volledig nieuwe invoer. Dit betekent dat de Key-Value (KV) cache voor de lange prefix voor elke kandidaat opnieuw wordt berekend of fysiek wordt gedupliceerd in het geheugen.
4. Dit leidt tot enorme redundantie in zowel rekentijd als geheugengebruik, wat de batchgrootte beperkt en de evaluatie van kwetsbaarheden op grotere schaal onhaalbaar maakt. Bestaande inferentie-systemen zoals vLLM of SGLang zijn niet optimaal voor dit specifieke scenario omdat ze geen backpropagation (gradiënten) ondersteunen, wat essentieel is voor witte-doos-aanvallen.

Methodologie: Prefix-Shared KV Cache (PSKV)

De auteurs introduceren PSKV, een plug-and-play inferentie-optimalisatietechniek die specifiek is ontworpen voor de structuur van suffix-jailbreak-aanvallen.

Kernprincipes:

• Gedeelde Prefix Cache: PSKV maakt gebruik van de observatie dat alle kandidaat-prompten dezelfde prefix (de schadelijke instructie) delen. In plaats van de KV-cache voor deze prefix $N$ keer te dupliceren, berekent PSKV deze eenmalig en bewaart slechts één kopie.
• Lazy Layer-wise Expansion: Tijdens de inferentie van de kandidaat-suffixen haalt PSKV de KV-vectoren van de prefix dynamisch op uit de cache en dupliceert deze alleen op het moment dat een specifieke attention-laag ze nodig heeft. Dit gebeurt laag-voor-laag, wat de piekgeheugengebruik minimaliseert in vergelijking met het gelijktijdig vasthouden van volledige duplicaten.
• Suffix-Centric Alignment: Voor batch-verwerking van meerdere verschillende instructies (multi-instruction attacks) introduceert PSKV een uitlijningsstrategie. Hierbij worden instructies links-opgevuld (left-padded) en doelantwoorden rechts-opgevuld (right-padded). Hierdoor beginnen alle gegenereerde suffixen op hetzelfde indexpunt in de tensor, wat efficiënte vectorisatie en parallelle verwerking mogelijk maakt zonder de gedeelde prefix-cache te hoeven dupliceren voor elke instructie.

Complexiteitsanalyse:

• Rekencomplexiteit: PSKV behoudt dezelfde rekencomplexiteit als een standaard KV-cache (beide vermijden herberekening van de prefix), maar elimineert de overhead van het dupliceren van de cache voor de batch.
• Geheugencomplexiteit: Waar standaard KV-cache een geheugencomplexiteit heeft van $O(N_{cand} \cdot N_{prefix})$ (waarbij $N_{cand}$ het aantal kandidaten is), reduceert PSKV dit tot $O(B \cdot N_{prefix})$ (waarbij $B$ het aantal unieke instructies is). Dit verwijdert de afhankelijkheid van het geheugen van de zoekbreedte.

Belangrijkste Bijdragen

1. Identificatie van inefficiëntie: Het paper identificeert dat algemene inferentie-caching inefficiënt is voor de specifieke "één-prefix-veel-suffixen" structuur van jailbreak-optimalisatie.
2. PSKV Framework: Een lichtgewicht, workload-bewust framework dat virtuele cache-broadcasting en tensor-uitlijning implementeert om redundante berekeningen te elimineren.
3. Uitgebreide Evaluatie: Een grondige evaluatie over zes veelgebruikte suffix-aanvallen (zoals GCG, AutoDAN, BEAST) en vijf populaire open-source LLM's (Vicuna, Llama-2/3, Mistral, Qwen).

Resultaten

De experimentele resultaten tonen aanzienlijke verbeteringen aan:

• Snelheid: PSKV reduceert de inferentietijd met ongeveer 40% in vergelijking met een baseline zonder cache, en is aanzienlijk sneller dan standaard KV-cache implementaties die vaak vastlopen door geheugenproblemen (OOM).
• Geheugenefficiëntie: De piekgeheugengebruik wordt met 50% verlaagd ten opzichte van standaard KV-cache implementaties.
• Scalabiliteit: PSKV maakt het mogelijk om zoekbreedtes (batch sizes) te vergroten zonder dat het geheugenexploderend toeneemt. Bijvoorbeeld, bij een 4x vergroting van de zoekbreedte bleef het totale piekgeheugen vrijwel constant.
• Aanvalssucces (ASR): De Attack Success Rate blijft onveranderd. PSKV is een "drop-in" optimalisatie die de onderliggende aanvalsalgoritmen niet wijzigt, waardoor de effectiviteit behouden blijft.
• Compatibiliteit: In tegenstelling tot systemen zoals vLLM of SGLang, ondersteunt PSKV volledige gradiëntberekening, waardoor het werkt met zowel gradiënt-gebaseerde als gradiënt-vrije aanvallen.

Betekenis en Impact

Dit werk is van cruciaal belang voor de beveiliging van LLM's. Door de reken- en geheugenkosten van het evalueren van kwetsbaarheden drastisch te verlagen, stelt PSKV onderzoekers in staat om:

1. Grootschalige evaluaties uit te voeren op grotere modellen en in bredere zoekruimtes die eerder onbereikbaar waren.
2. Snellere iteraties te hebben bij het testen van defensieve maatregelen (red-teaming).
3. Een brug te slaan tussen hoge-prestatie inferentie en flexibele, gradiënt-gebaseerde adversariele optimalisatie.

Kortom, PSKV maakt het systematisch testen van LLM's op veiligheidskwetsbaarheden veel praktischer en schaalbaarder, wat essentieel is voor het ontwikkelen van robuustere AI-systemen.