Closing the Distribution Gap in Adversarial Training for LLMs

Probleemstelling: De Distributie-Gap

Hoewel Adversarial Training (AT) een veelbelovende methode is om de robuustheid van Large Language Models (LLMs) tegen kwaadaardige aanvallen te verbeteren, vertonen huidige modellen nog steeds kwetsbaarheden voor eenvoudige, "in-distribution" exploits. Voorbeelden hiervan zijn het herschrijven van prompts in de verleden tijd of het vertalen ervan naar andere talen.

De auteurs stellen dat deze kwetsbaarheid voortkomt uit een fundamentele beperking in bestaande AT-algoritmen:

1. Empirisch vs. Populatie-Risico: Bestaande methoden minimaliseren het empirische robuuste risico op een vaste, eindige dataset ($D$). Dit is echter een slechte benadering van het ware populatie-robuuste risico ($R_{pop}$) over de volledige verdeling van natuurlijke taal ($q$).
2. De Gap: Er ontstaat een "distributie-gap" waarbij modellen wel bestand zijn tegen specifieke, geoptimaliseerde aanvallen (model-specifieke perturbaties), maar falen bij data-specifieke generalisatiefouten. Deze fouten treden op omdat de trainingsdata de volledige steun (support) van de verdeling van prompts die tot schadelijke antwoorden leiden, niet dekt.

Methodologie: Distributional Adversarial Training (DAT)

Om deze kloof te dichten, introduceren de auteurs Distributional Adversarial Training (DAT). Dit kader combineert generatieve sampling met continue adversarial training om de populatie-risico beter te benaderen.

Kerncomponenten:

1. Generatieve Surrogaatmodellen (Diffusion LLMs):

   • In plaats van te vertrouwen op een statische dataset, gebruiken de auteurs voorgeïmplementeerde Diffusion LLMs.
   • In tegenstelling tot autoregressieve modellen die alleen de conditionele verdeling $q(y|x)$ leren, modelleren diffusion modellen de gezamenlijke verdeling $q(x, y)$.
   • Dit maakt het mogelijk om het inverse probleem op te lossen: het genereren van diverse, hoog-likelihood prompts ($x$) die specifiek gericht zijn op een schadelijke respons ($y$). Dit wordt gedaan via "inpainting"-achtige conditionering.
   • Deze gegenereerde prompts zijn data-specifiek (ze komen voor in de natuurlijke verdeling) en niet alleen model-specifiek (zoals geoptimaliseerde suffixen van GCG).

2. Het DAT-Trainingscircuit:

   • Buitenste lus (Monte Carlo Sampling): Er worden diverse prompts gegenereerd uit de diffusion-verdeling, geconditioneerd op schadelijke antwoorden. Dit dekt de "distributie-gap" door de trainingsdata uit te breiden naar gebieden die in de oorspronkelijke dataset ontbraken.
   • Binnenste lus (Continue Adversarial Training - CAT): Op deze gegenereerde samples wordt continue adversarial training toegepast. Dit minimaliseert het worst-case verlies door continue perturbaties ($\delta$) toe te passen op token-embeddings.
   • Doelfunctie: De totale loss ($L_{DAT}$) combineert het adversarial verlies met een KL-divergentie-term op een "retain set" (veilige data) om te voorkomen dat de modelkwaliteit voor nuttige taken verslechtert.

Theoretische Rechtvaardiging:
De auteurs bewijzen een stelling (Theorem 3.1) die aantoont dat als de diffusion-surrogaat de conditionele verdeling $q(x|y)$ nauwkeurig benadert (binnen een foutmarge $\epsilon$), de kloof tussen het empirische risico en het ware populatie-risico lineair wordt beperkt.

Belangrijkste Bijdragen

1. Formalisatie van de Robuustheidsgap: Het paper onderscheidt duidelijk tussen model-specifieke perturbaties en data-specifieke generalisatiefouten, en positioneert het falen van huidige defenses als een probleem van distributie-benadering.
2. Nieuw Trainingskader (DAT): Een innovatieve methode die generatieve sampling (via Diffusion LLMs) combineert met continue adversarial optimalisatie om zowel de approximatiefout als de optimalisatiefout aan te pakken.
3. Verbeterde Robuustheid: Experimenteel bewijs dat DAT modellen oplevert die aanzienlijk robuuster zijn tegen een breed scala aan state-of-the-art aanvallen, inclusief die welke traditionele AT-methoden omzeilen.

Resultaten

De experimenten zijn uitgevoerd op modellen zoals Llama3-8B en Qwen2.5-14B en vergeleken met baselines zoals CAT, LAT, MixAT-GCG en Circuit Breakers.

• Worst-Case Robuustheid:

  • Bestaande methoden (zoals CAT en LAT) falen vaak bij "Inpainting"-aanvallen (waarbij prompts worden herschreven) en "Best-of-N" (BoN) aanvallen. Bijvoorbeeld, LAT had een aanvalsuccesratio (ASR) van 94% tegen Inpainting op Llama3-8B.
  • DAT presteert aanzienlijk beter: de ASR voor Inpainting daalt naar 32% (Llama3-8B) en 14% (Qwen2.5-14B).
  • De Best-of-All (BoA) ASR (de ergste case over alle aanvallen) wordt voor Llama3-8B verlaagd naar 36% (tegenover 77% voor de volgende beste methode, MixAT). Voor Qwen2.5-14B is dit 18% (tegenover 75-93% voor anderen).

• Nuttigheid (Utility) vs. Robuustheid:

  • DAT behoudt een uitstekende balans tussen veiligheid en nut. De modellen presteren vergelijkbaar met andere AT-baselines op benchmarks zoals XSTest (hulpvaardigheid) en MMLU/ARC (algemene capaciteiten).
  • De Pareto-grens (trade-off tussen robuustheid en nut) wordt door DAT naar een gunstiger positie verschoven vergeleken met alle andere methoden.

• Ablatie-studies:

  • Alleen het gebruik van diffusion-sampling ("Diffusion-only") verbetert de weerstand tegen Inpainting, maar faalt tegen andere aanvallen. Dit bevestigt dat zowel de distributie-benadering (buitenste lus) als de continue adversarial optimalisatie (binnenste lus) noodzakelijk zijn.
  • Het aantal gegenereerde samples per instructie heeft een direct lineair effect op de robuustheid: meer unieke samples leiden tot lagere ASR, zonder afbreuk te doen aan de nuttigheid.
  • Data-specifiekheid is cruciaal: Het trainen op "slechte" (laag-likelihood) samples die niet effectief zijn, levert weinig robuustheid op. Alleen samples die de hoge-likelihood manifold van natuurlijke aanvallen dekken, zijn effectief.

Betekenis en Impact

Dit werk is significant voor de AI-veiligheid omdat het een fundamentele beperking van bestaande defensiemethoden voor LLMs adresseert. Door te erkennen dat het simpelweg "meer data" of "sterkere perturbaties" op een vaste dataset niet voldoende is, biedt DAT een theoretisch onderbouwd pad naar echte generalisatie.

De methode toont aan dat het gebruik van generatieve modellen (Diffusion LLMs) als surrogaat voor de data-verdeling een krachtige manier is om de "distributie-gap" te dichten. Dit resulteert in modellen die minder vatbaar zijn voor real-world jailbreaks, zoals vertalingen of herschrijvingen, wat essentieel is voor de betrouwbare implementatie van LLMs in kritieke toepassingen.