💬 NLP

Stochasticity in Tokenisation Improves Robustness

This paper demonstrates that training large language models with uniformly sampled stochastic tokenisations significantly enhances their robustness against adversarial and random perturbations while preserving accuracy and inference efficiency.

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

Published 2026-04-20

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, well-read librarian (the AI model) who answers your questions. But there's a catch: this librarian only understands words if they are written in a very specific, rigid way.

In the world of Large Language Models (LLMs), this "writing style" is called tokenization. It's the process of chopping a sentence into small chunks (tokens) that the computer can read. Usually, this is done in one single, deterministic way. For example, the word "revolution" might always be chopped into one single chunk: revolution.

The Problem: The "Brittle" Librarian

The paper argues that if you train your librarian to only ever see the word "revolution" as one single chunk, they become brittle.

If someone sneaks into the library and changes the way the word is chopped—say, splitting it into re, vol, and ution—the librarian panics. They might think it's a completely different word or get confused, leading to a wrong answer. The paper calls this a "non-canonical tokenization." It's like if you told the librarian, "I'm asking about 'rev-ol-u-tion' today," and they suddenly forgot how to read "revolution."

The researchers found that these models are surprisingly fragile. A tiny, harmless change in how a word is split can cause the AI to fail completely, even if the meaning of the sentence hasn't changed at all.

The Solution: The "Flexible" Librarian

The paper proposes a simple but powerful fix: Stochastic Tokenization.

Instead of training the librarian to see "revolution" only as revolution, you train them to see it in many different ways randomly. Sometimes it's revolution, sometimes re-vol-ution, sometimes rev-o-lution.

Think of it like teaching a child to recognize a dog.

The Old Way (Deterministic): You only show the child a Golden Retriever. If they see a Poodle, they don't recognize it.
The New Way (Stochastic): You show the child a Golden Retriever, a Poodle, a Chihuahua, and a Great Dane. You teach them that "dog" can look many different ways.

By training the AI with this "randomness" (stochasticity), the model learns the essence of the word rather than just the specific way it was chopped up. It becomes robust.

The Experiments: What Happened?

The researchers tested this idea in three different "training camps":

Pre-training (The Foundation): They trained a small model from scratch using random chopping.
- Result: The model got better at understanding the tiny details of words, but it wasn't a miracle cure on its own.
Fine-tuning (The Specialization): They took a big, powerful model (like Llama) and gave it a "boot camp" where it had to solve puzzles using randomly chopped words.
- Result: Huge success! The model became incredibly tough. When they tried to trick it with weirdly chopped words later, it didn't even flinch. It kept its accuracy high, even when the input was messy.
In-Context Learning (The Quick Lesson): They tried to teach the model on the fly by showing it examples in the chat prompt.
- Result: It helped a little, but not as much as actually training the model.

The Secret Sauce: Uniformity

The paper also discovered that not all "randomness" is created equal.

The Flawed Method (STOCHASTOK): The previous method of random chopping was biased. It was like rolling a weighted die; some splits happened way more often than others.
The Better Method (Uniform Sampling): The researchers created new methods to ensure every possible way of chopping a word had an equal chance of being used.
- Analogy: Imagine a lottery. The old method had some numbers that were drawn way more often. The new method ensures every number has an equal shot. This "fairness" made the model even stronger against malicious attacks (adversarial attacks) designed to trick it.

The Big Takeaway

The paper concludes that randomness is a superpower for AI safety.

By intentionally training AI models to handle messy, inconsistent, and random ways of breaking down words, we make them:

Stronger: They don't break when someone tries to trick them with weird formatting.
Smarter: They understand the actual meaning of words, not just the specific code they were given.
Cheaper: You don't need more computer power to run them; you just need to change how you train them.

In short: If you want an AI that can't be easily fooled by a simple typo or a weird font, stop teaching it to be rigid. Teach it to be flexible, random, and adaptable. That's the key to a robust AI.

1. Problem Statement

Large Language Models (LLMs) rely on deterministic subword tokenization (e.g., Byte Pair Encoding) to convert text into sequences of tokens. This process is typically canonical, meaning a specific string always maps to the same token sequence.

Vulnerability: Recent studies indicate that LLMs are brittle to non-canonical tokenizations. Even if the underlying text remains unchanged, slight variations in how the text is split into tokens (e.g., revolution vs. re-vol-ution) can cause significant performance drops or allow adversarial attacks that bypass safety alignments without altering the semantic content.
Gap: While previous work (e.g., STOCHASTOK) introduced stochastic tokenization to improve subword understanding, it remains unclear whether training with stochastic tokenization effectively improves robustness against random perturbations and worst-case adversarial attacks. Furthermore, existing stochastic methods often suffer from sampling bias and incomplete support (failing to cover all valid token splits).

2. Methodology

The authors systematically investigate the impact of stochastic tokenization across three learning regimes: Pre-training, Supervised Fine-Tuning (SFT), and In-Context Learning (ICL).

A. Tokenization Schemes

The paper compares several tokenization strategies:

CANON: Standard deterministic tokenization.
STOCHASTOK: A method that randomly splits tokens into subtokens based on merge rules. The authors identify that this induces a biased distribution (preferring splits with fewer alternatives) and has incomplete support (cannot generate all valid non-canonical splits).
STOCHASTOK-UNI (Proposed): A two-stage method that first samples the number of splits per token (following the STOCHASTOK distribution) and then uniformly samples the specific splits for each token. This removes bias while maintaining the split-count distribution.
UNIFORM-K (Proposed): Samples tokenizations uniformly from the set of all valid non-canonical tokenizations at a specific edit distance $k$ (total split count) across the entire sequence. This uses a Multi-valued Decision Diagram (MDD) to handle the exponential search space efficiently.
UNIFORM: Full uniform sampling over all possible valid tokenizations (computationally expensive).

B. Experimental Setup

Models: Tiny-LLM (50M params, trained from scratch), Llama-1b, Llama-8b, GPT2-xl, Qwen-0.6b, and Gemma-1b.
Datasets:
- Language Games & CUTE: Tasks requiring fine-grained subword understanding (e.g., counting letters, detecting substrings).
- General Benchmarks: ARC-E, COPA, CSQA, HellaSwag, PIQA, SocialIQA.
Evaluation Metrics:
- Robustness to Random Perturbations: Testing on non-canonical tokenizations with varying levels of "splits" (edit distance).
- Adversarial Robustness: Using a greedy search algorithm (Geh et al., 2025) to find tokenizations that maximize the margin of error (force the model to predict the wrong answer).
Theoretical Analysis: The authors derive an upper bound on adversarial risk ( $\gamma$ ) for a classifier, showing that stochastic tokenization reduces the Lipschitz constant of the embedding space, thereby smoothing the decision boundary.

3. Key Contributions

Empirical Evidence of Robustness: Demonstrated that training with stochastic tokenization significantly improves robustness against both random and adversarial tokenization attacks across pre-training, fine-tuning, and ICL.
Identification of Sampling Bias: Showed that the original STOCHASTOK method is biased and has incomplete support, making it vulnerable to adversarial attacks compared to uniform sampling methods.
Novel Sampling Algorithms: Introduced STOCHASTOK-UNI and UNIFORM-K, which provide uniform sampling distributions and full support over non-canonical tokenizations, leading to superior robustness.
Theoretical Justification: Provided a theoretical bound proving that stochastic tokenization reduces the model's sensitivity to input perturbations by smoothing the embedding space.
Practical Guide: Established that fine-tuning with stochastic tokenization is a low-cost, high-reward strategy to harden LLMs against tokenization-based attacks without increasing inference costs.

4. Key Results

A. Robustness to Random Perturbations

Fine-Tuning: Fine-tuning a canonically pre-trained Llama-1b model with UNIFORM-K or STOCHASTOK-UNI drastically reduces accuracy loss when tested on non-canonical inputs.
- Canonical Fine-Tuning: Accuracy drops by 29.8% when evaluated on uniformly sampled non-canonical tokenizations.
- Stochastic Fine-Tuning (UNIFORM-K): Accuracy drops by only 0.6% (or even improves slightly) under the same conditions.
Pre-Training: Pre-training with stochastic tokenization improves performance on subword tasks and provides mild robustness, but fine-tuning is the critical step for maximizing robustness.

B. Adversarial Robustness

Catastrophic Failure of Canonical Models: When attacked with adversarial tokenizations, a canonically fine-tuned Llama-1b model's accuracy plummeted from 94.0% to 6.1%.
Resilience of Stochastic Models: Models fine-tuned with UNIFORM-K maintained an accuracy of 69.6% under the same adversarial conditions.
Uniformity Matters: Methods with uniform sampling (UNIFORM-K, STOCHASTOK-UNI) consistently outperformed the biased STOCHASTOK method in adversarial settings.

C. In-Context Learning (ICL)

Using stochastic tokenization in the few-shot context (examples) improves robustness to stochastic queries, though the gains are smaller than full fine-tuning due to the limited number of examples ( $K=10$ ) compared to the thousands of training samples used in SFT.

D. Theoretical Findings

The derived bound $\gamma \propto \sqrt{k}$ (where $k$ is the perturbation budget) suggests that stochastic tokenization effectively reduces the local Lipschitz constant of the model.
Empirical measurements of hidden state distances confirmed that stochastic training results in smaller distances between embeddings of canonical and non-canonical tokenizations, indicating a smoother embedding space.

5. Significance and Implications

Security: The paper highlights a critical vulnerability in current LLMs: they are easily "jailbroken" or misled by manipulating token boundaries rather than content. Stochastic tokenization offers a robust defense mechanism.
Training Paradigm Shift: The authors argue that fine-tuning should generally be performed with stochastic tokenization instead of canonical tokenization. This improves robustness without increasing inference costs (since inference can still be deterministic or use the canonical path if desired, though the paper suggests testing with perturbations is beneficial).
Subword Understanding: The method forces models to learn the internal structure of words (character-level composition) more effectively, leading to better performance on tasks requiring fine-grained linguistic reasoning.
Future Directions: The work suggests that future tokenization schemes should prioritize uniform sampling and full support to maximize model robustness, moving away from the standard deterministic canonical approach.

In conclusion, the paper establishes that stochasticity in tokenization is not just a regularization technique for subword understanding, but a fundamental defense mechanism that significantly enhances the robustness of LLMs against adversarial attacks and input perturbations.