[Text Sampled Out]. (I) From Text Generation to Search Sampling

【搜出来的文本】⋅（一）从文本生成到搜索采样 苏剑林 · 2021-01-07 · Machine Learning text generation sampling discrete optimization MCMC

Recently, I've fallen into a new rabbit hole: doing certain text generation tasks based on the idea of discrete optimization. Simply put, the idea is to write down, in quantitative terms, the objective for the text we want to generate, construct a distribution from it, and then either search for the maximum point of this distribution or sample from it. This process typically requires no training on labeled data. Since language is discrete, continuous optimization methods like gradient descent are not applicable, and since this distribution usually has no easily sampleable form, direct sampling isn't feasible either. This calls for some specially designed sampling algorithms, such as Rejection Sampling, MCMC (Markov Chain Monte Carlo), MH Sampling (Metropolis-Hastings Sampling), Gibbs Sampling, and so on.

Some readers might find this a bit familiar — doesn't it feel like going back to those headache-inducing years of learning LDA (Latent Dirichlet Allocation)? Indeed, the sampling algorithms mentioned above are also essential background for understanding the LDA model. In this post, we'll review these various sampling algorithms, which will appear in the many text generation applications to be introduced later on. more

Clarifying the Objective

Very often, we need to generate a target text $\boldsymbol{x}$ based on some specific information $\boldsymbol{c}$, which in mathematical terms is a conditional language model $p(\boldsymbol{x}|\boldsymbol{c})$. However, we can't obtain enough paired corpora $(\boldsymbol{x},\boldsymbol{c})$ to directly supervise the training of a conditional language model; instead we can only train an unconditional language model $p(\boldsymbol{x})$. That said, we can artificially design a metric to quantitatively describe the relationship between $\boldsymbol{x}$ and $\boldsymbol{c}$. Given this, how to perform conditional text generation using the relationship between the unconditional language model $p(\boldsymbol{x})$ and $\boldsymbol{x},\boldsymbol{c}$ becomes our object of study. We can call this "Constrained Text Generation."

For example, take sentence generation from keywords: here $\boldsymbol{c}$ would be the set of keywords, and we can define an indicator function:

\begin{equation}\chi(\boldsymbol{x}, \boldsymbol{c})=\left\{\begin{aligned}&1,\,\,\text{if}\boldsymbol{x}\text{contains keyword set}\boldsymbol{c} \\ &0,\,\,\text{if}\boldsymbol{x}\text{does not contain keyword set}\boldsymbol{c}\end{aligned}\right. \end{equation}

and then define

\begin{equation}\rho(\boldsymbol{x}, \boldsymbol{c}) = p(\boldsymbol{x})\chi(\boldsymbol{x}, \boldsymbol{c})\end{equation}

$p(\boldsymbol{x})$ ensures the fluency of the generated sentence, while $\chi(\boldsymbol{x}, \boldsymbol{c})$ ensures that the generated sentence contains the required keywords. The problem then reduces to either a maximization operation $\mathop{\text{argmax}}\limits_{\boldsymbol{x}} \rho(\boldsymbol{x}, \boldsymbol{c})$ or a sampling operation $\boldsymbol{x}\sim \rho(\boldsymbol{x}, \boldsymbol{c})$. Of course, $\rho(\boldsymbol{x}, \boldsymbol{c})$ here isn't yet a probability distribution — it only becomes a true probability distribution after normalization:

\begin{equation}\frac{\rho(\boldsymbol{x}, \boldsymbol{c})}{\sum\limits_{\boldsymbol{x}}\rho(\boldsymbol{x}, \boldsymbol{c})} = \frac{p(\boldsymbol{x})\chi(\boldsymbol{x}, \boldsymbol{c})}{\sum\limits_{\boldsymbol{x}}p(\boldsymbol{x})\chi(\boldsymbol{x}, \boldsymbol{c})}\end{equation}

But the denominator is usually intractable to compute explicitly. That is to say, all we know about the distribution to be sampled is that it is proportional to some function $\rho(\boldsymbol{x}, \boldsymbol{c})$, without knowing the exact distributional expression.

There's no shortage of similar examples — text summarization, for instance. What is text summarization? It's essentially expressing, with fewer words $\boldsymbol{x}$, as much of the same meaning as the original text $\boldsymbol{c}$ as possible. In this case we can define:

\begin{equation}\rho(\boldsymbol{x}, \boldsymbol{c}) = p(\boldsymbol{x})\cdot \text{sim}(\boldsymbol{x}, \boldsymbol{c})\cdot \chi(\boldsymbol{x}, \boldsymbol{c})\end{equation}

Here $\text{sim}(\boldsymbol{x}, \boldsymbol{c})$ is some text similarity function, and $\chi(\boldsymbol{x}, \boldsymbol{c})$ is a length indicator function, i.e. it equals 1 if the length of $\boldsymbol{x}$ falls within some range (possibly dependent on $\boldsymbol{c}$), and 0 otherwise. Here again we obtain an unnormalized probability distribution $\rho(\boldsymbol{x}, \boldsymbol{c})$, which we need to either maximize or sample from. Clearly, this objective means that we want a piece of text that is as semantically similar as possible to the original, subject to a certain length constraint — and isn't that exactly the point of a summary? So the core starting point of this whole approach is: we need to pin down, in quantitative terms, exactly what we want to generate, and only then proceed to the next step.

Analyzing the Difficulty

So, setting aside the background above, the problem we now face is: we have a distribution $p(\boldsymbol{x})$, of which we only know $p(\boldsymbol{x})\propto \rho(\boldsymbol{x})$, i.e.

\begin{equation}p(\boldsymbol{x}) = \frac{\rho(\boldsymbol{x})}{\sum\limits_{\boldsymbol{x}} \rho(\boldsymbol{x})}\end{equation}

and we cannot explicitly compute the denominator. Throughout this series, $\boldsymbol{x}$ represents text, i.e. a sequence of discrete elements, but the conclusions below apply equally to the case where $\boldsymbol{x}$ is a continuous vector. Now we want to search for the argmax position $\mathop{\text{argmax}}\limits_{\boldsymbol{x}} p(\boldsymbol{x})$ or perform sampling $\boldsymbol{x}\sim p(\boldsymbol{x})$. As we'll see later, finding the maximum can actually be viewed as a special case of sampling, so we'll mainly focus on sampling methods.

As mentioned earlier, the reason we need specially designed algorithms to perform sampling is that direct sampling from $p(\boldsymbol{x})$ is difficult, and we need to understand exactly where this difficulty lies in order to truly grasp the key points of the sampling algorithms introduced later. Where's the difficulty? If the space of candidate values for $\boldsymbol{x}$ isn't too large — even a million candidates — we could compute $p(\boldsymbol{x})$ for each one and then do ordinary categorical sampling. However, in general the candidate space of $\boldsymbol{x}$ is far larger than a million; if $\boldsymbol{x}$ has 10 components, each with 10,000 choices (corresponding to the vocabulary size), then the total number of permutations is $10^{40}$ — there's no way to precompute the probability of every permutation and sample according to those probabilities.

So what do we do? As the saying goes, "without accumulating small steps, one cannot travel a thousand miles" — we simply have to proceed one step at a time. That is, if I can't directly do a "choose 1 out of $10^{40}$" in one shot, can I instead do ten rounds of "choose 1 out of $10^4$"? This corresponds to so-called "autoregressive generation":

\begin{equation}p(\boldsymbol{x})=p(x_1) p(x_2|x_1) p(x_3|x_1, x_2) \cdots p(x_n|x_1,\cdots,x_{n-1}) = \prod_{t=1}^n p(x_t|\boldsymbol{x}_{< t})\end{equation}

This way we can first sample a $x_1$ from $p(x_1)$, then sample a $x_2$ from $p(x_2|x_1)$, and so on recursively. However, autoregressive generation only applies to unconditional language models or supervised Seq2Seq models. If we want to add constraints to the generation process of an unconditional language model, as in the examples given earlier, the resulting model is no longer autoregressive, and so we can no longer sample via this kind of recursion.

This is precisely why we need the various sampling algorithms introduced below — they too embody the "one step at a time" philosophy, but apply to a broader class of distributions.

Importance Sampling

In earlier posts such as "A Unified View of Sampling and Optimization: Differentiable and Non-Differentiable Optimization" and "How to Split a Validation Set That's Closer to the Test Set?", we introduced the concept of "importance sampling": if we want to estimate the expectation $\mathbb{E}_{\boldsymbol{x}\sim p(\boldsymbol{x})}[f(\boldsymbol{x})]$, but $p(\boldsymbol{x})$ is not an easy distribution to sample from, we can find a distribution $q(\boldsymbol{x})$ that's close to $p(\boldsymbol{x})$ and easy to sample from, and then, via the following transformation

\begin{equation} \mathbb{E}_{\boldsymbol{x}\sim p(\boldsymbol{x})}[f(\boldsymbol{x})] = \sum_{\boldsymbol{x}} p(\boldsymbol{x}) f(\boldsymbol{x}) = \sum_{\boldsymbol{x}} q(\boldsymbol{x})\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}f(\boldsymbol{x}) = \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}\left[\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}f(\boldsymbol{x})\right] \end{equation}

turn the problem into computing the expectation of $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}f(\boldsymbol{x})$ under samples drawn from $q(\boldsymbol{x})$, i.e. weighting each sample by $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$ — hence the name "Importance Sampling." Even if we only know $p(\boldsymbol{x})\propto \rho(\boldsymbol{x})$, importance sampling can still be carried out, because

\begin{equation}1 = \sum_{\boldsymbol{x}} p(\boldsymbol{x}) = \sum_{\boldsymbol{x}} q(\boldsymbol{x})\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})} = \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}\left[\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}\right]\end{equation}

and therefore

\begin{equation} \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}\left[\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}f(\boldsymbol{x})\right] = \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}\left[\frac{p(\boldsymbol{x}) / q(\boldsymbol{x})}{\mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}[p(\boldsymbol{x}) / q(\boldsymbol{x})]}f(\boldsymbol{x})\right] \end{equation}

We can see here that the expression above depends only on the relative values of $p(\boldsymbol{x})$, not on its absolute value, so we can replace $p(\boldsymbol{x})$ with anything proportional to it, $\rho(\boldsymbol{x})$, which ultimately simplifies to:

\begin{equation} \mathbb{E}_{\boldsymbol{x}\sim p(\boldsymbol{x})}[f(\boldsymbol{x})] \approx \frac{\sum\limits_{i=1}^N \rho(\boldsymbol{x}_i) / q(\boldsymbol{x}_i) \cdot f(\boldsymbol{x}_i)}{\sum\limits_{i=1}^N \rho(\boldsymbol{x}_i) / q(\boldsymbol{x}_i)},\quad \boldsymbol{x}_1,\cdots,\boldsymbol{x}_N\sim q(\boldsymbol{x})\end{equation}

Rejection Sampling

The importance sampling in the previous section converts an expectation under a complicated distribution into an expectation under a simple distribution, but that's not actually our real goal — what we want to achieve is drawing samples from the distribution $p(\boldsymbol{x})$ itself, not just estimating some expectation of it. The idea is still similar to importance sampling: we introduce an easy-to-sample distribution $q(\boldsymbol{x})$, and then randomly discard certain samples so that the remaining samples follow the distribution $p(\boldsymbol{x})$.

Specifically, suppose we have a function $\alpha(\boldsymbol{x})\in [0, 1]$, and we sample according to the following procedure, called "Rejection Sampling":

Rejection Sampling: Draw a sample $\boldsymbol{x}$ from $q(\boldsymbol{x})$, draw a random number $\varepsilon$ from $U[0,1]$. If $\varepsilon \leq \alpha(\boldsymbol{x})$, accept the sample; otherwise reject it and repeat the procedure.

So what is the true probability distribution of the samples $\boldsymbol{x}$ obtained this way? It's not hard to work out: since the probability that a sample $\boldsymbol{x}$ is kept is $\alpha(\boldsymbol{x})$, its relative probability is $q(\boldsymbol{x})\alpha(\boldsymbol{x})$, and we just need to renormalize it:

\begin{equation}\frac{q(\boldsymbol{x})\alpha(\boldsymbol{x})}{\sum\limits_{\boldsymbol{x}} q(\boldsymbol{x})\alpha(\boldsymbol{x})}\end{equation}

and we obtain the true probability distribution corresponding to rejection sampling. From this form we can also see that, in theory, multiplying the acceptance rate by any number between 0 and 1 leaves the distribution corresponding to rejection sampling unchanged.

This process suggests that rejection sampling allows us to sample from a distribution proportional to $q(\boldsymbol{x})\alpha(\boldsymbol{x})$. So, based on $p(\boldsymbol{x})=q(\boldsymbol{x})\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$, we can take $\alpha(\boldsymbol{x})=\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$ as the acceptance probability and perform rejection sampling starting from $q(\boldsymbol{x})$, and the result is equivalent to sampling from $p(\boldsymbol{x})$. Of course, it's not quite that simple: by the normalization property of probabilities, unless $q(\boldsymbol{x})$ is identically equal to $p(\boldsymbol{x})$, $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$ cannot always stay within $[0, 1]$. But that's fine — as long as $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$ has an upper bound, we can choose a sufficiently large constant $M$ such that $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})\cdot M}\in [0, 1]$, and then use $\alpha(\boldsymbol{x})=\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})\cdot M}$ as the acceptance probability. As noted earlier, multiplying by a constant doesn't affect the distribution corresponding to rejection sampling. In other words, this process likewise doesn't depend on knowing $p(\boldsymbol{x})$ exactly — we can replace $p(\boldsymbol{x})$ with anything proportional to it, $\rho(\boldsymbol{x})$.

Regarding the acceptance rate $\alpha(\boldsymbol{x})$: although in theory it only needs to satisfy $\alpha(\boldsymbol{x})\in[0, 1]$, in practice it's still better to keep it $\max\limits_{\boldsymbol{x}}\alpha(\boldsymbol{x}) = 1$, because too small an acceptance rate leads to excessive rejection (almost every draw gets rejected), making sampling too inefficient and the cost of producing a reasonable sample too high. Similarly, although in theory the only requirements on $q(\boldsymbol{x})$ are that it be easy to sample from and that $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$ have an upper bound, in practice the closer $q(\boldsymbol{x})$ is to $p(\boldsymbol{x})$ the better, since otherwise the acceptance rate may still end up too low, making the sampling cost prohibitively high. So although rejection sampling appears to offer a way to sample from almost any distribution $p(\boldsymbol{x})$, in practical applications designing the approximating distribution $q(\boldsymbol{x})$ remains a considerable challenge.

Summary

Starting with this post, we've opened up a new line of exploration, attempting to tackle certain text generation tasks (constrained text generation) from the perspective of discrete optimization. The approach works by pinning down a quantitative evaluation objective, and then either maximizing this objective or sampling from it to obtain the output we want, without requiring supervised training of a new model on labeled data. The tools needed for this process are mainly sampling algorithms; this post introduced two very basic ones — importance sampling and rejection sampling. We'll continue to flesh out this series in future posts, so stay tuned.