A Brief Exploration of Random Tokenization: From Viterbi Decoding to Viterbi Sampling

随机分词浅探：从Viterbi Decoding到Viterbi Sampling 苏剑林 · 2023-09-16 · Machine Learning probability random tokenization new word discovery

After my previous post A Problem with Large-Vocabulary Language Models on Continuation Tasks, and a Countermeasure was published, readers quickly pointed out that introducing randomness into tokenization during training could address the same problem, and that there were already papers and implementations doing exactly this. After digging further, I found this technique is called Subword Regularization, originally applied in NMT (neural machine translation), and now also implemented in SentencePiece. It does seem to alleviate the aforementioned problem, and may even help improve the fault tolerance of language models — so I started thinking about adding it to BytePiece.

That raises the question: how do we turn deterministic tokenization into random tokenization? BytePiece is based on the Unigram model, which uses the Viterbi algorithm to find the highest-probability segmentation. Since there's already a probability distribution involved, can random sampling be derived naturally from it? This post discusses this question and shares my own solution.more

Framing the Problem

Currently, Unigram tokenization directly outputs the highest-probability segmentation, which is typically a deterministic output. Specifically, suppose $\boldsymbol{w}=(w_1,w_2,\cdots,w_k)$ represents a segmentation scheme with corresponding score $P(\boldsymbol{w})=p(w_1)p(w_2)\cdots p(w_k)$, and $\Omega(S)$ denotes the set of all possible segmentations of sentence $S$. Then tokenization can be described as

\begin{equation}\boldsymbol{w}^* = \mathop{\text{argmax}}_{\boldsymbol{w}\in \Omega(S)}P(\boldsymbol{w})\end{equation}

This can be computed in linear time via the Viterbi algorithm, which is why we call this process "Viterbi Decoding." It might seem, since the Unigram model naturally comes with probabilities attached, that turning it into probability-weighted sampling shouldn't be hard. But on closer inspection, this turns out to be a non-trivial problem with quite a few subtle difficulties to overcome.

My idea was to design a recursive sampling procedure analogous to autoregressive language models. The hardest part here is trying to preserve, as much as possible, the original ranking of candidate segmentations — or, failing that, at least ensuring that the maximum-probability result stays the same. That is, the maximum-probability path $\boldsymbol{w}^*$ found by Viterbi decoding should also be the most likely outcome produced by whatever recursive sampling algorithm we design. Since the set of all segmentations $\Omega(S)$ forms a directed acyclic graph (DAG), I initially thought random walking directly on this DAG might work. But on further thought, it's hard to design suitable transition probabilities that preserve the maximum-probability path (because different outgoing edges from the same node are not equally weighted — you can't simply sample according to edge frequency).

Existing Approaches

Since I couldn't come up with a new idea right away, I decided to go check the "reference answer" — namely, how the original Subword Regularization paper, Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates, tackles this.

However, this "standard answer" left me a bit bemused. It turns out Subword Regularization's approach is remarkably simple and direct: first search for the top $n$ segmentations $\boldsymbol{w}^*_1,\boldsymbol{w}^*_2,\cdots,\boldsymbol{w}^*_n$ with the largest $P(\boldsymbol{w})$ (the $n$-best segmentations), then construct the following distribution

\begin{equation}p_i = \frac{P(\boldsymbol{w}^*_i)^{\alpha}}{\sum\limits_{j=1}^n P(\boldsymbol{w}^*_j)^{\alpha}}\end{equation}

and sample from these $n$ candidates according to this probability, where $\alpha > 0$ is a hyperparameter. This algorithm is already integrated into SentencePiece, and readers can try it themselves (see here for usage).

The problem is that "simple and direct" doesn't mean "efficient." Although finding the top-$n$ segmentations is also a linear-time operation (interested readers can look up N-best Viterbi), it's clearly much more expensive than plain top-1 Viterbi Decoding (theoretically $n$ times the complexity). The direct consequence is that turning on random sampling makes tokenization considerably slower than deterministic tokenization — which is not the ideal sampling method I had in mind.

My Own Approach

So my thinking hit a wall again. At an impasse, I decided to go back and re-examine my reasoning: the goal is to find a random sampling algorithm with complexity similar to Viterbi Decoding. Given that, Viterbi Decoding itself should be a good place to start looking for a breakthrough. So I opened up the tokenization code again — the tokenization function at the time looked like this:

def _tokenize(self, bytes text):
    cdef int e, k, s
    cdef double v, score
    cdef list routes = [(0, None)] + [(-INFINITY, None) for _ in text]
    cdef list tokens = []
    for e, (k, v) in self._automaton.iter(text):
        s, e = e - k + 1, e + 1
        score = routes[s][0] + v
        if score > routes[e][0]:
            routes[e] = score, s
    while text:
        s = routes[e][1]
        tokens.append(text[s:e])
        text, e = text[:s], s
    return tokens[::-1]

After reading through it a few times, I finally had a flash of insight: the key line in Viterbi Decoding is if score > routes[e][0]:, which keeps track of the optimal segmentation up to the current position, where score is the score (log probability) of the new candidate segmentation and routes[e][0] is the best score seen so far — if the new candidate is better, it overwrites the old one. This reminded me of the acceptance-rate design in MCMC algorithms. What if we introduced randomness right at this step? Wouldn't that randomize the tokenization result?

Let $r\in \{1, 0\}$ denote accepting/rejecting the new candidate. Since this step is just a binary choice, turning it into a probabilistic one is quite simple:

\begin{equation} r_i = \left\{\begin{aligned}&\,1\,, \,\, s_i > s_{i-1} \\ &\,0\,, \,\, \text{else}\end{aligned}\right.\qquad\longrightarrow\qquad r_i = \left\{\begin{aligned}&\,1\,, \,\, \varepsilon < \sigma(\alpha(s_i - s_{i-1})) \\ &\,0\,, \,\, \text{else}\end{aligned}\right. \end{equation}

Here $\varepsilon\sim U[0,1]$ is a uniform random number, $\alpha > 0$ is a hyperparameter, $\sigma(t)=1/(1+e^{-t})$ is the sigmoid function, and $s_i,s_{i-1}$ are the scores (log probabilities) of the new and old candidates respectively. It's easy to see that the deterministic sampling rule on the left corresponds to the case $\alpha\to\infty$ in the probabilistic version.

This gives us a very natural and lightweight random sampling algorithm built directly on top of Viterbi decoding, which I'll call "Viterbi Sampling." Implementing it only requires replacing the criterion if score > routes[e][0]: with a version that incorporates a random number. Because of the monotonicity of the sigmoid function, when $s_i > s_{i-1}$, the new candidate is naturally assigned a higher probability — so it's clear that the original maximum-probability segmentation is still the most likely outcome under Viterbi Sampling. Moreover, the larger $s_i - s_{i-1}$ is, the larger $\sigma(\alpha(s_i - s_{i-1}))$ becomes, which means segmentations with higher original scores are also more likely to be sampled — to some extent preserving the ranking of segmentations (though I haven't proven this is strictly order-preserving in general, from a practical standpoint approximate order-preservation is good enough).

A Quick Test

Starting from version 0.4.0, Viterbi Sampling is built into BytePiece's tokenization function — just pass an alpha parameter greater than 0 to tokenizer.tokenize or tokenizer.encode, and the result will be randomized:

import bytepiece
assert bytepiece.__version__ >= '0.4.0'

tokenizer = bytepiece.Tokenizer('bytepiece_160k.model')
text = '今天天气不错'
print(tokenizer.tokenize(text))  # alpha默认值为-1，alpha≤0 都代表确定性分词
for i in range(5):
    print(tokenizer.tokenize(text, alpha=0.1))

# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6', b'\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8', b'\x8d', b'\xe9\x94', b'\x99']
# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9', b'\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
# [b'\xe4\xbb', b'\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9', b'\xe6\xb0\x94\xe4\xb8\x8d', b'\xe9\x94', b'\x99']

Let's compare the speed of SentencePiece's Subword Regularization against BytePiece's Viterbi Sampling (with $\alpha=0.1$ set for both random tokenization runs):

$$\begin{array}{c|cc} \hline & \text{deterministic tokenization} & \text{random tokenization} & \\ \hline \text{SP-BPE} & \text{1.36M bytes/sec} & \text{1.25M bytes/sec} \\ \text{SP-Unigram} & \text{5.65M bytes/sec} & \text{1.28M bytes/sec} \\ \text{BytePiece} & \text{1.95M bytes/sec} & \text{1.36M bytes/sec}\\ \hline \end{array}$$

As you can see, once Subword Regularization is turned on (the "SP-Unigram" row), tokenization speed drops to less than a quarter of its original value, showing that Subword Regularization's sampling algorithm is quite inefficient. By contrast, the Viterbi Sampling proposed in this post only drops by about 30%, which is clearly much more efficient — the slowdown here comes from generating random numbers and computing the sigmoid function, and if these two parts were further optimized, speed could be improved even more. As for the BPE model, its random tokenization variant is called BPE Dropout, a method specific to BPE models — interested readers can look into it on their own; I won't cover it here.

Summary

This post explored strategies for turning deterministic Unigram tokenization into random tokenization. While a method called "Subword Regularization" already exists for this purpose, it's relatively inefficient. To address this, I proposed a more efficient sampling algorithm, Viterbi Sampling, which only requires a simple modification to deterministic Viterbi Decoding, thereby largely preserving the original efficiency. Experiments show that the new algorithm samples noticeably faster than Subword Regularization. The corresponding implementation is now built into the latest version of BytePiece.