Stochasticity in Tokenisation Improves Robustness

TL;DR

Stochastic tokenisation enhances the robustness of large language models against adversarial and random input perturbations without increasing inference costs.

Contribution

This study systematically demonstrates that training with stochastic tokenisation improves LLM robustness across various regimes and architectures.

Findings

Stochastic tokenisation reduces vulnerability to adversarial attacks.

Pre-training with stochastic tokenisation improves model robustness.

Training with stochastic tokenisation maintains accuracy without extra inference cost.

Abstract

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.