Stochasticity in Tokenisation Improves Robustness

Titel: Stochasticiteit in Tokenisatie Verbetert Robuustheid

Auteurs: Sophie Steger, Rui Li, Sofiane Ennadir, Anya Sims, Arno Solin, Franz Pernkopf, Martin Trapp.

---

1. Het Probleem

Grootte Taalmodellen (LLMs) zijn kwetsbaar voor aanval via manipulatie van de tokenisatie van de invoer. Moderne LLMs gebruiken doorgaans een deterministische tokenisatie (bijv. Byte Pair Encoding - BPE), waarbij een tekstreeks altijd op exact dezelfde manier wordt omgezet in een reeks tokens (de "canonieke tokenisatie").

• Kwetsbaarheid: Onderzoek toont aan dat kleine wijzigingen in de tokenisatie (bijvoorbeeld het splitsen van een woord in andere, maar geldige, subwoorden) de prestaties van het model drastisch kunnen verstoren. Dit maakt modellen broos tegenover adversariële aanvallen, waarbij een aanvaller een tekst manipuleert zodat deze voor de mens hetzelfde blijft, maar voor het model een andere tokenisatie krijgt die de voorspelling verandert.
• Onderzoeksvraag: Kan het trainen van modellen met stochastische tokenisatie (waarbij tokens willekeurig op verschillende manieren worden gesplitst) de robuustheid tegen zowel willekeurige verstoringen als adversariële aanvallen verbeteren?

2. Methodologie

De auteurs onderzoeken dit probleem door verschillende trainingsregimes te analyseren: pre-training, supervised fine-tuning (SFT) en in-context learning (ICL). Ze vergelijken verschillende methoden voor het genereren van stochastische tokenisaties.

A. Bestaande Methode: STOCHASTOK

Ze beginnen met STOCHASTOK (Sims et al., 2026), een methode die tokens willekeurig splitst in subtokens volgens de merge-regels van het vocabulaire.

• Beperkingen: De auteurs ontdekken dat STOCHASTOK twee fundamentele problemen heeft:
  1. Bias: De verdeling is niet uniform; bepaalde splitsingen zijn waarschijnlijker dan andere.
  2. Onvolledige Support: Het kan niet alle mogelijke geldige tokenisaties genereren (bijvoorbeeld splitsingen die niet via een reeks van 2-way merges bereikbaar zijn).

B. Nieuwe Methodes: Uniforme Sampling

Om deze beperkingen op te lossen, introduceren ze twee nieuwe methoden:

1. STOCHASTOK-UNI: Een uitbreiding die de bias verwijdert. Het genereert eerst een splitsingsaantal per token (gebaseerd op de STOCHASTOK-verdeling) en sample vervolgens uniform binnen die specifieke splitsingsaantallen.
2. UNIFORM-K: Deze methode samplet uniform over alle mogelijke tokenisaties met een specifieke "edit distance" (aantal splitsingen) over de gehele invoerreeks. Dit biedt de grootste flexibiliteit en kan splitsingen over de grenzen van canonieke tokens heen maken.

C. Experimenteel Opzet

• Modellen: Tiny-LLM (vanaf scratch), Llama-1b, Llama-8b, GPT2-xl, Qwen-0.6b.
• Datasets: LANGUAGE GAME (subwoord-taken zoals letters tellen), CUTE, en diverse meerkeuze-benchmarks (ARC-E, COPA, etc.).
• Aanvalsscenario's:
  • Willekeurige perturbaties: Toevalsgewijs gekozen niet-canonieke tokenisaties.
  • Adversariële aanvallen: Een greedy search-algoritme (gebaseerd op Geh et al., 2025) dat probeert de tokenisatie te vinden die de voorspelling van het model het meest verstoort.

3. Belangrijkste Bijdragen

1. Empirisch Bewijs: Ze tonen aan dat training met stochastische tokenisatie de robuustheid tegen zowel willekeurige als adversariële tokenisatie-aanvallen aanzienlijk verbetert, zonder de nauwkeurigheid op schone (canonieke) data te verlagen.
2. Analyse van Sampling Bias: Ze identificeren dat de bestaande STOCHASTOK-methode een voorkeur heeft voor bepaalde tokenisaties en een beperkt bereik heeft. Ze bewijzen dat het oplossen van deze bias en het vergroten van de "support" (het bereik van mogelijke tokenisaties) cruciaal is voor maximale robuustheid.
3. Nieuwe Algoritmen: Introductie van STOCHASTOK-UNI en UNIFORM-K als superieure alternatieven voor STOCHASTOK.
4. Theoretische Analyse: Ze leveren een theoretische onderbouwing voor de verbeterde robuustheid. Ze tonen aan dat stochastische tokenisatie de "Lipschitz-constante" van het model verlaagt, wat betekent dat het model minder gevoelig is voor kleine veranderingen in de invoerrepresentatie.

4. Resultaten

• Robuustheid tegen Willekeurige Perturbaties:
  • Een canoniek getraind Llama-1b model verliest 29,8% nauwkeurigheid wanneer het getest wordt met uniform gesamplede niet-canonieke tokenisaties.
  • Modellen die zijn gefine-tuned met UNIFORM-K of STOCHASTOK-UNI behouden hun nauwkeurigheid (verlies < 1%) en zijn zeer resistent tegen perturbaties.
• Robuustheid tegen Adversariële Aanvallen:
  • Canoniek getrainde modellen zakken onder een adversariële aanval van ~94% naar 6,1% nauwkeurigheid.
  • Modellen getraind met UNIFORM-K zakken slechts naar 69,6%.
  • STOCHASTOK (zonder uniformiteit) biedt enige verbetering, maar presteert significant slechter dan de uniforme methoden, wat bevestigt dat uniformiteit en volledige support essentieel zijn.
• In-Context Learning (ICL):
  • Het gebruik van stochastische tokenisatie in de contextvoorbeelden (zonder het model te finetunen) biedt ook enige robuustheid, maar de winst is kleiner dan bij full fine-tuning.
• Kosten:
  • De training vereist stochastische tokenisatie, maar de inference-kosten (tijd en rekenkracht tijdens het gebruik) nemen niet toe. Het model kan tijdens het gebruik gewoon de standaard canonieke tokenisatie verwerken.

5. Betekenis en Conclusie

Dit paper levert een belangrijke bijdrage aan de veiligheid en betrouwbaarheid van LLMs.

• Praktische Implicatie: Het suggereert dat standaard fine-tuning van LLMs niet langer alleen met deterministische tokenisatie moet gebeuren. Het gebruik van uniforme stochastische tokenisatie tijdens training is een effectieve en goedkope (geen extra inferentie-kosten) methode om modellen veel robuuster te maken tegen manipulatie.
• Theoretische Inzicht: Het paper verbindt de kwetsbaarheid van tokenisatie direct aan de Lipschitz-continuïteit van het model. Door de trainingsverdeling te spreiden over meerdere tokenisaties, wordt de representatieruimte "geglad", waardoor het moeilijker wordt voor een aanvaller om een punt te vinden dat de voorspelling verandert.
• Toekomst: Hoewel de focus nu ligt op vraag-antwoord taken, biedt dit een blauwdruk voor het verbeteren van de veiligheid van generatieve modellen in het algemeen.

Samenvattend: Stochastische tokenisatie tijdens training is geen "nice-to-have", maar een noodzakelijke stap om LLMs te beschermen tegen de inherente kwetsbaarheid van deterministische tokenisatieschema's.