Stochasticity in Tokenisation Improves Robustness

Abstract: The widespread adoption of LLMs has increased concerns about their robustness. Vulnerabilities in perturbations of tokenisation of the input indicate that models trained with a deterministic canonical tokenisation can be brittle to adversarial attacks. Recent studies suggest that stochastic tokenisation can deliver internal representations that are less sensitive to perturbations. In this paper, we analyse how stochastic tokenisations affect robustness to adversarial attacks and random perturbations. We systematically study this over a range of learning regimes (pre-training, supervised fine-tuning, and in-context learning), data sets, and model architectures. We show that pre-training and fine-tuning with uniformly sampled stochastic tokenisations improve robustness to random and adversarial perturbations. Evaluating on uniformly sampled non-canonical tokenisations reduces the accuracy of a canonically trained Llama-1b model by 29.8%. We find that training with stochastic tokenisation preserves accuracy without increasing inference cost.

• The paper demonstrates that incorporating stochastic tokenisation during fine-tuning nearly eliminates performance drops under non-canonical perturbations.
• It introduces StochasTok-uni and Uniform-k schemes that use uniform sampling over tokenisations to reduce vulnerabilities against adversarial attacks.
• Empirical and theoretical analyses show that randomising tokenisation improves subword-level representations without compromising clean accuracy.

Stochastic Tokenisation as a Mechanism for Improved LLM Robustness

Introduction

The paper "Stochasticity in Tokenisation Improves Robustness" (2604.16037) presents a comprehensive analysis of the impact of stochastic tokenisation on the robustness of LLMs to both random and adversarial input perturbations at the tokenisation level. The study systematically dissects the vulnerabilities of LLMs exposed to non-canonical tokenisations and explores whether introducing controlled stochasticity at various stages of LLM training and usage can enhance their resilience. The work contributes detailed empirical evaluations across pre-training, supervised fine-tuning, and in-context learning, providing both theoretical and practical insights into model behavior under alternative tokenisation regimes.

Background and Motivation

Modern LLMs rely almost universally on deterministic subword tokenisation schemes, such as Byte Pair Encoding (BPE) or Unigram LM, yielding a canonical segmentation of text into tokens. However, these encodings are inherently many-to-one mappings, with lexical and structural ambiguities permitting multiple valid token sequences that reconstruct the same input string. Canonically trained models are typically only exposed to single, fixed decompositions during training and inference.

Recent work demonstrates that even minor deviations from canonical tokenisation, whether random or adversarially crafted, can substantially degrade LLM performance—revealing a gap in subword-level generalisation and a concrete axis of vulnerability for prompt-based attacks [geh2025adversarial, zheng2025broken]. Previous approaches like BPE-dropout and StochasTok have primarily aimed to enrich the internal token-level representations, but a dedicated analysis of their effect on downstream robustness has been lacking.

Methodologies and Stochastic Tokenisation Schemes

The paper introduces and rigorously analyses several stochastic tokenisation protocols, building on but generalising prior art.

• StochasTok: Recursively splits canonical tokens into allowable subtokens according to the tokeniser's merge rules, introducing variability in input segmentations via controlled expansions.
• StochasTok-uni: Recognises the sampling bias in original StochasTok (non-uniform sampling, incomplete support) and introduces a uniform sampling scheme over tokenisations conditioned on a fixed per-token split budget. This ensures every admissible tokenisation with a given number of splits has equal probability.
• Uniform-k: Further generalises to marginally uniform sampling over all tokenisations within $k$ token-level edit distance of the canonical, leveraging Multi-valued Decision Diagrams (MDDs) for tractable enumeration and sampling.

Both StochasTok-uni and Uniform-k provide broader support in the space of non-canonical segmentations, crucial for robust model generalisation.

Experimental Analysis

Pre-training Regimes

Empirical analysis with small-scale LLMs demonstrates that stochastic tokenisation during pre-training enhances subword task accuracy and robustness to random tokenisation perturbations (as measured by normalised split distance). However, the robustness induced by pre-training with stochasticity is moderate, and clean (canonical) accuracy sees minimal degradation. Importantly, these benefits primarily manifest when aligned with corresponding fine-tuning strategies.

Supervised Fine-tuning

Fine-tuning canonically pre-trained models with stochastic tokenisation confers pronounced robustness benefits, notably with no trade-off in clean accuracy. For instance, evaluating on randomly sampled non-canonical tokenisations reduces the accuracy of a LoRA-fine-tuned Llama-1b model by nearly 30% after canonical fine-tuning, but stochastic fine-tuning with StochasTok-uni or Uniform-k almost completely eliminates this drop (residual performance drops are ≤1%). In the adversarial setting, canonical fine-tuned models can suffer 80–90% absolute drops in task accuracy, while stochastic fine-tuned models retain up to 70% aboslute accuracy.

In-Context Learning

Stochastic tokenisation employed in context examples during in-context learning (ICL) yields partial robustness to non-canonical query perturbations, but with smaller effect sizes than full-parameter fine-tuning. The practical implication is that while ICL can compensate for some tokenisation sensitivity, it is insufficient as a standalone mechanism for hardening LLMs.

Detailed Insights on Sampling Bias and Support

The experiments reveal that the full support and uniformity of sampling over alternative tokenisations critically mediate adversarial robustness. StochasTok-uni and Uniform-k, by expanding the distributional coverage of tokenisation variations, prevent overfitting to highly localised support in the segmentation space, which adversaries otherwise exploit.

Comparison to Baselines

StochasTok-uni and Uniform-k outperform both canonical fine-tuning and earlier stochastic schemes (e.g., BPE-dropout), especially under maximal perturbations and adversarially optimised tokenisation attacks.

Theoretical Analysis

A formal robustness analysis quantifies adversarial risk under tokenisation perturbations. Under assumptions of 1-Lipschitz activation functions and bounded embeddings, the adversarial risk (maximum output shift under tokenisation-space perturbations within an edit-distance budget $k$) is upper bounded as

$$\gamma = \sqrt{2} \lVert W \rVert \left(\frac{d}{d-1}\right)^2 C_1 C_2 \sqrt{k}$$

where $W$ is the embedding matrix and $C_1, C_2$ depend on (parameter norms of) the attention and MLP blocks [ennadir2025pool]. This characterisation demonstrates that vulnerabilities to tokenisation adversaries grow with the local Lipschitz constant associated with the initial embedding step, underscoring the importance of smoothing this space via randomisation in input segmentations.

Empirically, the local representation distances between canonical and non-canonical segmentations are substantially reduced after stochastic fine-tuning, suggesting that the norm of the difference in hidden representations (local Lipschitz constant) is also reduced, further corroborating the theoretical findings.

Implications and Future Directions

These results have several immediate implications for both the design and deployment of LLMs:

• Fine-tuning Protocols: For downstream tasks, performing stochastic tokenisation during fine-tuning significantly mitigates both random and targeted attacks, without incurring cost at inference. Uniform tokenisation schemes should become standard for robust alignment.
• Tokenizer Design: Tokenisers and tokenisation schemes with broader non-canonical support and unbiased sampling should be incorporated into LLM development toolkits to harden models to input-level attacks.
• Benchmarking Robustness: Canonical tokenisation performance may overstate true model capability; evaluation should incorporate random and adversarial tokenisation setups as a robustness check.
• Theoretical Guarantees: Developing explicit, quantifiable robustness guarantees based on model and tokenisation properties can inform model selection and fine-tuning strategies under threat models targeting the tokenisation stage.

Potential future developments include extension of these methods to generation tasks (rather than multiple-choice QA), exploration of stochastic decoding at inference, and further optimisation of sampling tractability for large vocabularies and long sequences.

Conclusion

This work provides compelling evidence, both empirical and theoretical, that stochastic tokenisation—particularly when inducing uniform coverage over alternative segmentations—fundamentally improves LLM robustness to both random and adversarial input perturbations in the tokenisation space. Fine-tuning with these schemes closes the gap in performance between canonical and non-canonical inputs while obviating additional computational cost at inference. These findings set a new methodological standard for robust LLM deployment and suggest revisiting canonical design decisions in LLM training and evaluation pipelines (2604.16037).