FlashRT: Towards Computationally and Memory Efficient Red-Teaming for Prompt Injection and Knowledge Corruption

1. Probleemstelling

Long-context Large Language Models (LLMs) (bijv. Gemini-3.1-Pro, Qwen-3.5) worden steeds vaker gebruikt in Retrieval-Augmented Generation (RAG), autonome agenten en AI-assistenten. Ze zijn echter kwetsbaar voor prompt injection en knowledge corruption aanvallen, waarbij adversarial tekst subtiel wordt ingebed binnen een lange context om de output van het model te manipuleren.

Hoewel optimatie-gebaseerde red-teaming methoden (zoals GCG en zijn variant nanoGCG) effectiever zijn in het vinden van kwetsbaarheden dan heuristische methoden, staan ze voor ernstige schaalbaarheidsproblemen in lange-context scenario's:

• Hoge GPU-geheugenvraag: De backward pass die nodig is voor gradientberekening verbruikt 2–4× meer geheugen dan de forward pass vanwege de noodzaak om intermediaire activaties en gradients op te slaan. Voor lange contexten (bijv. 32K tokens) kan dit het beschikbare GPU-geheugen overschrijden (bijv. vereisend >260 GB).
• Hoge rekenkosten: Het evalueren van duizenden kandidaat-adversarial teksten vereist talloze forward passes. Hoewel standaard KV-Caching helpt, vereist het nog steeds het herberekenen van hidden states voor alle tokens die volgen op de adversarial tekst, wat leidt tot enorme tijdsverliezen (bijv. uren per aanval).

Deze resource-beperkingen voorkomen een systematische beveiligingsevaluatie van lange-context LLMs en de beoordeling van verdedigingsstrategieën op grote schaal.

2. Methodologie: FlashRT

De auteurs stellen FlashRT voor, een framework ontworpen om zowel reken- als geheugenefficiëntie te optimaliseren voor white-box optimatie-gebaseerde aanvallen. Het adresseert de twee belangrijkste knelpunten via twee nieuwe technieken:

A. Vermindering van Forward Pass Berekening: Selective Recomputation

Standaard KV-Caching hergebruikt key-value (KV) paren voor het gedeelde prefix, maar vereist het herberekenen van KV-paren voor de adversarial tekst en alle daaropvolgende tokens in de context om de exacte loss te berekenen. FlashRT introduceert Selective Recomputation:

• Concept: In plaats van KV-paren voor de volledige suffix (context na de adversarial tekst) te herberekenen, herberekent FlashRT alleen een subset van tokens die de kans op de doeloutput significant beïnvloeden.
• Invloedsscore: De methode berekent een "invloedsscore" voor segmenten van de context op basis van attention weights tussen context-tokens en de doel-output tokens. Tokens met hoge attention weights vanuit de doeloutput worden als kritiek beschouwd.
• Implementatie: Alleen de KV-paren voor de top $\beta$ fractie van tokens met hoge invloed (plus de adversarial tekst, gebruikersinstructie en doeloutput) worden herberekend. De rest wordt hergebruikt vanuit de gecachede beste kandidaat. Dit benadert de loss-berekening met hoge nauwkeurigheid terwijl het aantal operaties drastisch wordt verminderd.

B. Vermindering van Backward Pass Geheugen: Context Subsampling

De backward pass is geheugenintensief omdat het gradients opslaat voor alle parameters en activaties over de volledige context.

• Concept: FlashRT observeert dat exacte gradients niet strikt noodzakelijk zijn voor het genereren van kandidaat-tokens; een redelijk nauwkeurige benadering volstaat.
• Implementatie: De context wordt opgedeeld in segmenten. Voor elke gradientberekeningsstap samplet de methode willekeurig een fractie $\gamma$ van deze segmenten uit zowel de linker- als rechtercontext. De backward pass wordt alleen uitgevoerd op deze gesubsamplede context.
• Gradient Resampling: Om te voorkomen dat de optimalisatie stagneert in lokale minima door ruis in de gradient-schattingen, bewaakt FlashRT de loss-trend. Als de loss na een vastgesteld aantal iteraties ($\tau$) niet verbetert, voert het gradient resampling uit (het verversen van de subset van context-tokens die worden gebruikt).

Unified Framework

FlashRT integreert deze technieken in een geünificeerde optimalisatielus (Algoritme 2):

1. Initialiseer de beste adversarial tekst ($T_{best}$).
2. Bereken gradients met Context Subsampling (laag geheugen).
3. Genereer kandidaat-teksten.
4. Evalueer kandidaat-loss met Selective Recomputation (laag rekenvermogen) met gecachede KV-paren.
5. Update $T_{best}$ en de KV-cache als een betere kandidaat wordt gevonden.
6. Herstelt periodiek gradients als de optimalisatie vastloopt.

3. Belangrijkste Bijdragen

1. Eerste Efficiënt Framework voor Long-Context Red-Teaming: FlashRT is het eerste framework dat specifiek de geheugen- en rekenknelpunten van optimatie-gebaseerde aanvallen in lange-context LLMs adresseert.
2. Nieuwe Algorithmische Technieken: Het introduceert Selective Recomputation (benadering van loss via attention-gebaseerde token-selectie) en Context Subsampling (benadering van gradients via willekeurige segment-sampling), die orthogonaal zijn aan bestaande systeemniveau-optimalisaties zoals FlashAttention.
3. Generaliseerbaarheid: Hoewel gefocust op white-box aanvallen, blijkt het framework compatibel te zijn en in staat om black-box methoden te versnellen (bijv. TAP, AutoDAN) wanneer de aanvaller toegang heeft tot de interne log-probabilities van het model.
4. Open Source: De code is vrijgegeven om systematische beveiligingsevaluatie door de onderzoeksgemeenschap te faciliteren.

4. Experimentele Resultaten

De auteurs evalueerden FlashRT tegen state-of-the-art baselines (nanoGCG, nanoGCG-OPT, Heuristische aanvallen) op datasets zoals NarrativeQA, MuSiQue, GovReport, NQ en HotpotQA met verschillende LLMs (Llama-3.1, Qwen-2.5, Meta-SecAlign).

• Snelheidswinst: FlashRT bereikt een 2×–7× snelheidswinst in totale runtime.
  • Voorbeeld: Runtime gereduceerd van ~2.736 seconden (nanoGCG) naar ~1.039 seconden (FlashRT) op NarrativeQA met Llama-3.1-8B.
• Geheugenreductie: FlashRT reduceert het piek GPU-geheugengebruik met 2×–4×.
  • Voorbeeld: Geheugen gereduceerd van 164,8 GB (nanoGCG) naar 53,7 GB (FlashRT) voor een 32K token context.
  • Betekenis: Dit maakt aanvallen op grote modellen (bijv. Llama-3.1-70B) mogelijk die eerder onuitvoerbaar waren vanwege geheugenbeperkingen.
• Aanvalsuccesratio (ASR): FlashRT behoudt of verbetert de aanvalseffectiviteit.
  • Het bereikt een vergelijkbare of hogere ASR dan nanoGCG (bijv. 98% vs 98% op NarrativeQA).
  • In sommige gevallen (bijv. redenerende modellen zoals DeepSeek-R1) presteert het aanzienlijk beter dan baselines door het interne denkproces van het model effectief te sturen.
• Verdedigingsevaluatie: FlashRT slaagde erin robuuste modellen zoals Meta-SecAlign (8B en 70B) te red-teamen, wat aantoont dat zelfs uitgelijnde modellen kwetsbaar zijn voor geoptimaliseerde aanvallen wanneer ze efficiënt worden geëvalueerd.

5. Betekenis

• Systematisch Beveiligingsonderzoek mogelijk maken: Door de resource-barrière te verwijderen, stelt FlashRT academische onderzoekers en modelproviders in staat om systematisch de beveiliging van lange-context LLMs en hun verdedigingen op grote schaal te evalueren.
• Schaalbaarheid: Het maakt het haalbaar om modellen met 70B+ parameters en enorme contextvensters te testen, wat cruciaal is naarmate de industrie zich richt op langere contextmogelijkheden.
• Verbetering van Verdediging: Het vermogen om efficiënt sterke adversarial voorbeelden te genereren faciliteert adversarial training, wat helpt bij het bouwen van robuustere LLMs tegen prompt injection en knowledge corruption.
• Brede Toepasbaarheid: De technieken (selectieve herberekening en gradient-benadering) zijn generiek en kunnen worden toegepast op andere optimatie-gebaseerde taken buiten red-teaming.

Concluderend overbrugt FlashRT de kloof tussen de theoretische noodzaak van rigoureuze beveiligingstests voor lange-context LLMs en de praktische beperkingen van huidige hardware, en biedt het een schaalbaar, efficiënt en effectief hulpmiddel voor de AI-veiligheidsgemeenschap.