Stochasticity in Tokenisation Improves Robustness
For LLM developers, stochastic tokenisation offers a simple, cost-free method to enhance robustness against tokenisation-based attacks.
Stochastic tokenisation during pre-training and fine-tuning improves LLM robustness to adversarial and random perturbations, preserving accuracy without increasing inference cost. Canonically trained Llama-1b accuracy drops 29.8% on non-canonical tokenisations.
The widespread adoption of large language models (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.