Random Tokenization Revisited: From Viterbi Sampling to a Perfect Sampling Algorithm
In the post A Brief Exploration of Random Tokenization: From Viterbi Decoding to Viterbi Sampling, I proposed a random tokenization algorithm called "Viterbi Sampling," which is just a small modification on top of Viterbi Decoding (which finds the optimal solution). It retains the simplicity and speed of the Viterbi algorithm, and is noticeably more efficient than the existing Subword Regularization approach. However, a reader on Zhihu, @鶴舞, pointed out that the sampling algorithm as it stood might "dilute" the probability of certain segmentation schemes through repeated pairwise choices, with the direct consequence that the segmentation with the highest score would not necessarily appear with the highest probability.
After thinking it over carefully, I found that this issue does indeed exist — at the time, in my rush to get a new sampling algorithm out, I was admittedly a bit sloppy about the details. So in this post I'll further refine the Viterbi Sampling algorithm, and prove that the refined version is equivalent in effect to Subword Regularization.
Problem Analysis
First, let's look at the original comment:
In subword regularization, sampling can be guaranteed to follow the probabilities (with a temperature hyperparameter). In the proposed method, for each endpoint e, the route computed first gets "challenged" multiple times in 1v1 comparisons — won't the resulting probability distribution end up quite different from that of the existing algorithm?
For example, suppose "watching" has three possible segmentations: watch+ing, wat+ching, and w+atching, each with probability one-third. Under the proposed scheme, wouldn't the sampling probabilities end up being one-quarter for the first two and one-half for the third?
Actually the comment already explains it quite clearly, but let me elaborate a bit further in case some readers are still unsure. Suppose there are three segmentation schemes, each with the same score; naturally we'd want each scheme to have an equal probability of $1/3$ of being sampled. However, Viterbi Sampling turns the "choose one out of many" sampling process into a sequence of pairwise "choose one out of two" steps:
\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}
In this process, the first two segmentation schemes are compared pairwise first, each with probability $\frac{1/3}{1/3+1/3}=1/2$; once one of them is chosen, it's compared pairwise against the third scheme. Since the probability is computed based on each candidate's own score, the probability at this stage is still $1/2$ for each. So over the full sampling process, the first two schemes end up with probability $1/4$ each, while the last scheme ends up with probability $1/2$ — the later a scheme enters the comparison, the more it benefits, while the earlier ones get their probability diluted more severely. Unfortunately, given the order in which BytePiece's Aho-Corasick automaton returns matches, longer tokens (which typically have higher scores) tend to appear earlier. So in the original Viterbi Sampling, the higher-scoring schemes were actually more likely to have their probability diluted.
The Fix
As it turns out, the fix is quite simple: every time we make a pairwise choice, we also cache the cumulative probability. Then, starting from the second step, each new candidate entering the pairwise comparison is not compared against the score of the previously chosen candidate, but against the cumulative probability score. This is essentially the well-known "reservoir sampling" algorithm.
Using the earlier example: first the two segmentation schemes come in, and one is chosen with probability $\frac{1/3}{1/3+1/3}=1/2$; their combined cumulative probability is then $2/3$. Next, the currently selected candidate is compared against the new scheme, and the probability of the new scheme being selected should be $\frac{1/3}{2/3+1/3}=1/3$ — that is, it's compared against the cumulative probability, not against the probability of the currently selected candidate alone. With this scheme, over the whole sampling process each segmentation ends up with probability $1/3$.
For Viterbi Sampling, at each endpoint there are multiple candidate segmentations, and we need to sample one out of many, with the selection probability constructed from each candidate's score as $p_i = e^{\alpha s_i}/Z$, where $Z$ is the normalizing factor. Because we process this recursively, we don't know in advance how many candidates there are in the "many," and we can't compute $Z$ directly — but that's fine, because knowing $e^{\alpha s_i}$ is enough. This is because computing the conditional sampling probability at each step doesn't actually require the full $Z$, but rather the recursive quantity $Z_i$:
$$\begin{array}{c|c|c} \hline \text{Viterbi Decoding} & \text{old version Viterbi Sampling} & \text{new version Viterbi Sampling} \\ \hline r_i = \left\{\begin{aligned}&\,1\,, \,\, s_i > s_{i-1} \\ &\,0\,, \,\, \text{else}\end{aligned}\right. & r_i = \left\{\begin{aligned}&\,1\,, \,\, \varepsilon < \sigma(\alpha(s_i - s_{i-1})) \\ &\,0\,, \,\, \text{else}\end{aligned}\right. & \begin{aligned}Z_i =&\, Z_{i - 1} + e^{\alpha s_i} \\[1pt] r_i =&\, \left\{\begin{aligned}&\,1\,, \,\, \varepsilon < e^{\alpha s_i} / Z_i \\ &\,0\,, \,\, \text{else}\end{aligned}\right.\end{aligned} \\ \hline \end{array}$$
In practice, directly caching $Z_i$ runs a high risk of numerical overflow due to exponential blow-up, so we generally cache its logarithm $Z^{\log}_i$ instead, and use the $\text{logsumexp}$ function to avoid overflow:
\begin{equation} \begin{aligned}&\,Z^{\log}_i = \text{logsumexp}(Z^{\log}_{i-1}, \alpha s_i) \\ &\qquad e^{\alpha s_i} / Z_i \to e^{\alpha s_i - Z^{\log}_i} \end{aligned},\qquad \text{logsumexp}(x,y) = \left\{\begin{aligned}&\,x + \log(1+e^{y-x}),\,\, x \geq y \\ &\,y + \log(1 + e^{x-y}),\,\,x < y \end{aligned}\right. \end{equation}
The corresponding implementation is already built into bytepiece>=0.5.0.
Perfect Sampling
Overall, the flaw in the old version of Viterbi Sampling stemmed from having moved too quickly at the time. So now let's carefully supply the mathematical proof for the new version of Viterbi Sampling. Interestingly, it turns out that the updated Viterbi Sampling and Subword Regularization are both "perfect sampling" algorithms in the same sense.
As discussed previously, Subword Regularization takes a fairly "brute-force" approach: it directly finds the top $k$ highest-scoring segmentation schemes, and then computes the selection probability via $p_i = e^{\alpha s_i}/Z$, where $s_i$ is the score of the $i$-th scheme. Aside from its high computational complexity, there's nothing wrong with this approach. When $k$ is unrestricted (i.e., we enumerate all segmentation schemes), we obtain a random sample over the full set of segmentations, where the probability of sampling each scheme is proportional to $e^{\alpha s_i}$ — a monotonically increasing function of the score $s_i$, meaning the ranking of sampling probabilities matches the ranking of scores exactly. I call any sampling scheme satisfying both of these conditions "perfect sampling."
Decoding
To prove that the new version of Viterbi Sampling is also "perfect sampling," let's first revisit Viterbi Decoding. Suppose we have a byte string $c_1,c_2,\cdots,c_l$ of length $l$, and let $S^*(c_1,c_2,\cdots,c_l)$ denote the score of the optimal segmentation. Assuming we know that a split must occur between $c_k,c_{k+1}$, then necessarily
\begin{equation}S^*(c_1,c_2,\cdots,c_l) = S^*(c_1,c_2,\cdots,c_k) + S^*(c_{k+1},c_{k+2},\cdots,c_l)\end{equation}
That is, the sub-segmentation of the optimal segmentation restricted to a substring must itself be the optimal segmentation of that corresponding sub-byte-string — this is the fundamental basis of dynamic programming. Of course, in reality we cannot know in advance where a split will occur, so we must enumerate:
\begin{equation}S^*(c_1,c_2,\cdots,c_l) = \max\left\{\begin{aligned} &\,\color{green}{s\left(\overline{c_1,\cdots,c_l}\right)} \\ \color{red}{S^*(c_1)} \,+&\, \color{green}{s\left(\overline{c_2,\cdots,c_l}\right)} \\ \color{red}{S^*(c_1,c_2)} \,+&\, \color{green}{s\left(\overline{c_3,\cdots,c_l}\right)} \\ \vdots \\ \color{red}{S^*(c_1,\cdots,c_{l-2})} \,+&\, \color{green}{s\left(\overline{c_{l-1},c_l}\right)} \\ \color{red}{S^*(c_1,\cdots,c_{l-1})} \,+&\, \color{green}{s\left(\overline{c_l}\right)} \end{aligned}\right\}\label{eq:core}\end{equation}
where $s\left(\overline{c_1,\cdots,c_l}\right)$ denotes the score of the byte string $c_1, \cdots,c_l$ when treated as a single token (if it is not a token in the vocabulary, this is set to $-\infty$). In this way, computing $S^*(c_1,c_2,\cdots,c_l)$ reduces to computing $S^*(c_1),S^*(c_1,c_2),\cdots,S^*(c_1,\cdots,c_{l-1})$, and so on — computing $S^*(c_1,c_2,\cdots,c_{l-1})$ in turn reduces to computing $S^*(c_1),S^*(c_1,c_2),\cdots,S^*(c_1,\cdots,c_{l-2})$, and so forth. In other words, the result of $S^*$ can be reused. So the whole procedure boils down to one sentence:
As we scan to each position, we record the optimal segmentation up to that position and its score.
Of course, naively applying the recursion in $\eqref{eq:core}$ would in theory have complexity $\mathcal{O}(l^2)$, but in practice not every sub-byte-string will be a token in the vocabulary. So we can use a trie, an Aho-Corasick automaton, or similar structures to pre-scan, based on the vocabulary, all possible tokens that could occur. The complexity is then proportional to the number of candidate tokens found, which scales linearly in $l$. If we want a concrete estimate: assuming the maximum token length in the vocabulary is $m$, then the number of tokens found by scanning a byte string of length $l\geq m$ is at most
\begin{equation}l + (l - 1) + \cdots + (l - m + 1) = lm - \frac{1}{2}m(m-1) = \mathcal{O}(lm)\end{equation}
Sampling
With the Decoding section as groundwork, understanding Sampling becomes relatively easier. The key again lies in equation $\eqref{eq:core}$. Let $Z(c_1,c_2,\cdots,c_l)$ denote the normalizing factor over all segmentation schemes of the byte string $c_1,c_2,\cdots,c_l$ (for perfect sampling); then we have
\begin{equation}Z(c_1,c_2,\cdots,c_l) = \sum\left\{\begin{aligned} &\,\color{green}{e^{\alpha\cdot s\left(\overline{c_1,\cdots,c_l}\right)}} \\ \color{red}{Z(c_1)} &\, \color{green}{e^{\alpha\cdot s\left(\overline{c_2,\cdots,c_l}\right)}} \\ \color{red}{Z(c_1,c_2)} &\, \color{green}{e^{\alpha\cdot s\left(\overline{c_3,\cdots,c_l}\right)}} \\ \vdots \\ \color{red}{Z(c_1,\cdots,c_{l-2})} &\, \color{green}{e^{\alpha\cdot s\left(\overline{c_{l-1},c_l}\right)}} \\ \color{red}{Z(c_1,\cdots,c_{l-1})} &\, \color{green}{e^{\alpha\cdot s\left(\overline{c_l}\right)}} \end{aligned}\right\}\label{eq:core-2} \end{equation}
This identity also tells us how to sample from all segmentations of $c_1,c_2,\cdots,c_l$ with weights proportional to $e^{\alpha s}$: randomly pick one from all segmentations of $c_1,\cdots,c_{l-1}$ and append token $\overline{c_l}$, randomly pick one from all segmentations of $c_1,\cdots,c_{l-2}$ and append token $\overline{c_{l-1},c_l}$, randomly pick one from all segmentations of $c_1,\cdots,c_{l-3}$ and append token $\overline{c_{l-2},c_{l-1},c_l}$, and so on. Having obtained these $l$ sampling results, we then pick one of them with weights $Z(c_1,\cdots,c_{l-1}) e^{\alpha\cdot s\left(\overline{c_l}\right)}$, $Z(c_1,\cdots,c_{l-2}) e^{\alpha\cdot s\left(\overline{c_{l-1},c_l}\right)}$, $Z(c_1,\cdots,c_{l-3}) e^{\alpha\cdot s\left(\overline{c_{l-2},c_{l-1},c_l}\right)}$, ....
Following the same logic as in the Decoding case, computing $Z(c_1,\cdots,c_{l-1})$ can reuse the result of $Z(c_1),Z(c_1,c_2),\cdots,Z(c_1,\cdots,c_{l-2})$, computing $Z(c_1,\cdots,c_{l-2})$ can reuse the result of $Z(c_1),Z(c_1,c_2),\cdots,Z(c_1,\cdots,c_{l-3})$, and so on — and likewise the sampling results themselves can be reused. So similarly, the whole Sampling algorithm can also be summarized in one sentence:
As we scan to each position, we sample among all segmentation schemes ending at that position with weights $e^{\alpha s}$, and record the sampling result along with the cumulative weight $Z$.
Taking logarithms on both sides, equation $\eqref{eq:core-2}$ can be equivalently rewritten as
\begin{equation}Z^{\log}(c_1,c_2,\cdots,c_l) = \text{logsumexp}\left\{\begin{aligned} &\,\color{green}{\alpha\cdot s\left(\overline{c_1,\cdots,c_l}\right)} \\ \color{red}{Z^{\log}(c_1)} \,+&\, \color{green}{\alpha\cdot s\left(\overline{c_2,\cdots,c_l}\right)} \\ \color{red}{Z^{\log}(c_1,c_2)} \,+&\, \color{green}{\alpha\cdot s\left(\overline{c_3,\cdots,c_l}\right)} \\ \vdots \\ \color{red}{Z^{\log}(c_1,\cdots,c_{l-2})} \,+&\, \color{green}{\alpha\cdot s\left(\overline{c_{l-1},c_l}\right)} \\ \color{red}{Z^{\log}(c_1,\cdots,c_{l-1})} \,+&\, \color{green}{\alpha\cdot s\left(\overline{c_l}\right)} \end{aligned}\right\} \end{equation}
The difference from Viterbi Decoding's equation $\eqref{eq:core}$ is that $Z^{\log}$ replaces $S^*$, and $\text{logsumexp}$ replaces $\max$, and $\text{logsumexp}$ is precisely a smooth approximation of $\max$, so as $\alpha\to\infty$ it degenerates back to Viterbi Decoding. On the other hand, in actual computation, multiple segmentation schemes sharing the same endpoint arrive one at a time rather than all at once, which is why we need to convert the single-step "choose one out of many" into multiple steps of "choose one out of two" — this is exactly what was discussed in the "The Fix" section. With this, we have proven (or rather, re-derived starting from Viterbi Decoding) that the modified Viterbi Sampling is in fact a perfect sampling algorithm, just like Subword Regularization.
Summary
This post refines the previously proposed random tokenization algorithm Viterbi Sampling, and mathematically proves that it is a "perfect sampling" algorithm equivalent in effect to Subword Regularization, while being noticeably more efficient in practice than Subword Regularization.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.