SlowGuess/NeurIPS_2023_OCR huggingface dataset:
One of the primary advantages of using gradient-based techniques for controlled generation is that it provides a means to sample from the conditional distribution Eq. (4) without having to calculate the normalization factor  $Z_{t}$ . For example in HMC algorithm (Algorithm 1), all we need to calculate regarding the potential energy  $U(\mathbf{V}) = -\log p(\mathbf{V}\mid t)$  is two terms: (1) the gradient of  $U$  with respect to  $\mathbf{V}$ :  $\nabla U$ , and (2) the difference between the potential energy of two points  $\Delta U$  for the Metropolis criterion. Fortunately, both terms are independent of  $Z_{t}$ , and we can sample from the conditional distribution without the need to compute  $Z_{t}$ .
MuCoLA. As defined in Eq. (4),  $p(\mathbf{V} \mid t)$  is  $\mathbb{R}^d$ -valued, so it is tempting to apply a gradient-based sampling technique to sample from it. And, indeed, Kumar et al. [21] have proposed such a scheme based on Langevin dynamics. They define the potential energy of the embeddings as
Algorithm 1 HMC
Input:  $\mathbf{x}^t$ : current sample,  $U$ : potential energy function,  $\varepsilon$ : step size,  $L$ : number of leapfrog steps  
Output: next sample:  $\mathbf{x}^{t+1}$
1:  $\mathbf{r} \sim \mathcal{N}(\mathbf{0}, \mathbb{I})$
2:  $H^t\gets U(\mathbf{x}^t) + K(\mathbf{r})$
3:  $\mathbf{x}^{t + 1}\gets \mathbf{x}^t$
4: for  $l = 1, \dots, L$ :
5:  $\mathbf{r} \gets \mathbf{r} - \frac{\varepsilon}{2} \nabla U(\mathbf{x}^{t + 1})$
6:  $\mathbf{x}^{t + 1}\gets \mathbf{x}^{t + 1} + \varepsilon \mathbf{r}$
7:  $\mathbf{r} \gets \mathbf{r} - \frac{\varepsilon}{2}\nabla U(\mathbf{x}^{t + 1})$
8:  $H^{t + 1}\gets U(\mathbf{x}^{t + 1}) + K(\mathbf{r})$
9:  $\Delta H\gets H^{t + 1} - H^t$
10: if  $\mathbf{s} \sim \mathcal{U}(0, 1) < e^{-\Delta H}$ :
11: return  $\mathbf{x}^{t + 1}$
12: else
13: return  $\mathbf{x}^{t + 1}\gets \mathbf{x}^t$
There exist sophisticated methods [15] to tune the mass matrix, but given that LM gradients are expensive to compute, we will not attempt such methods in this paper.
$U(\mathbf{V}) \stackrel{\mathrm{def}}{=} -\log p(\mathbf{V} \mid t)$  and apply Langevin dynamics out of the box; see Algorithm 3. However, since  $U(\mathbf{V})$  is zero for any vector other than vectors in  $\mathcal{B}$  (defined in §2.1), they modify the sampling process and project the proposed sample to the base embeddings set at each step of the sampling process (line 3 of the algorithm). The added projection step, however, is neither volume-preserving nor time-reversible. Hence, this sampling procedure does not sample from the intended distribution and is not a valid MCMC algorithm.
COLD and other schemes. Alternative approaches have been proposed in the literature to reformulate a language model as potential energy over the logit [COLD; 38] or simplex space [14, 20]. However, these formulations are not suitable for principled gradient-based sampling. COLD only employs a heuristic energy function to select among the candidate generations obtained via top- $k$  sampling, and the simplex-based approach requires an extra constraint to ensure the sample stays on the simplex.
6 Structured Voronoi Sampling
Given our structured Voronoi measure  $p_{\mathrm{V}}$ , one can apply HMC to sample from it. In this section, we take one step further and propose a variation of HMC that is more suitable to sample from  $p_{\mathrm{V}}$ . Importantly,  $p_{\mathrm{V}}$  contains discontinuities whereas the generic leapfrog integrator does not account for such sudden jumps in the potential function. In other words, even if the leapfrog integrator itself is volume preserving and time-reversible, the sudden jumps in potential can lead to large deviations in the Hamiltonian value, causing a low acceptance rate. We therefore would like to find an alternative to leapfrog in such situations.
A Physical Analogy. In classical mechanics, discontinuity with smooth boundary in the potential function occurs naturally, e.g., in collision dynamics or a slab magnetic field, and is referred to as a potential barrier (or interface). Upon encountering a potential barrier, a particle will either be reflected from or transmitted through the barrier surface, depending on whether it has enough kinetic energy to overcome the potential jump (e.g., [8, §4.6.2]). Such behavior is similar to reflection-refraction phenomenon in optics. The key insight here is that, in both cases, the Hamiltonian is conserved.
Reflection and Refraction. To give a precise mechanical description of this behavior, suppose a particle encounters a potential barrier at position  $\mathbf{x}$  with momentum  $\mathbf{r}$ . We can decompose momentum as  $\mathbf{r} = \mathbf{r}_{\perp} + \mathbf{r}_{\parallel}$  where  $\mathbf{r}_{\perp}$  is normal to the barrier and  $\mathbf{r}_{\parallel}$  parallel to it. Let  $\Delta U$  be the signed potential energy difference between the two sides of the barrier. If  $\| \mathbf{r}_{\perp} \|_2^2 > 2\Delta U$ , then the particle has enough kinetic energy to overcome the barrier, and its momentum's normal component will instantaneously become  $\mathbf{r}_{\perp}' = \sqrt{\|\mathbf{r}_{\perp}\|_2^2 - 2\Delta U} \cdot \frac{\mathbf{r}_{\perp}}{\|\mathbf{r}_{\perp}\|_2^2}$  after being transmitted through the barrier (refraction). Otherwise, if  $\| \mathbf{r}_{\perp} \|_2^2 \leq 2\Delta U$ , the particle will be reflected from the barrier and the normal component will instantaneously become  $\mathbf{r}_{\perp}' = -\mathbf{r}_{\perp}$ . We show in App. E.1 that Hamiltonian is conserved in either case. The reflect-refract process is summarized in Algorithm 5.
6.1 A Sampling Algorithm
Noting that the discontinuity surfaces of  $p_{\mathrm{V}}$  are all piecewise smooth, we can build on the above and develop a sampling algorithm for  $p_{\mathrm{V}}$  to handle discontinuity in a principled and effective way. In fact, we only need to make one change to the generic HMC, which is updates  $(\mathbf{x}, \mathbf{r})$  according to the mechanical description given above. Concretely, we need to replace step 2 in the HMC outline in §5.1: When a single step of leapfrog encounters no discontinuity, we may advance to the next point as in HMC; however, when there is discontinuity, if a full leapfrog step is taken, we need to proceed by repeatedly computing where the discontinuity is encountered, taking a smaller step up to the discontinuity and refracting-reflecting based on the potential energy difference. This process is continued until we have exhausted the step size. Since refracting-reflecting conserves Hamiltonian (App. E.1), this process yields a better acceptance rate in the presence of a discontinuity. See Algorithm 4 for the details of this sampling procedure, and App. F.1 for how to efficiently find discontinuities. We will supply proof of correctness in App. E.2.
A note on calculating the base measure. To adjust the momentum, one needs to compute  $\Delta U$ , which implies computing the difference between two base measures, as defined in Eq. (12). However, computing such an integral in a high-dimensional space is not practical. Therefore, make an
Figure 2: Left: JS divergence between the reference probability and empirical probability distribution. Voronoi Sampling clearly outperforms others in low temperatures. Right: reference probability distribution annealed with 3 temperatures: 0.25 (peaked), 1, and 2 (close to uniform).
Figure 3: Perplexity of 100 samples taken with different gradient-based algorithms, compared to 1000 samples taken with the ancestral sampling (in green). While LANGEVIN, SVS, and MUCOLA are comparably close to the ancestral samples' distribution, SVS models the tail of the distribution better.
assumption that the base measures of Voronoi cells are equal, and thus do not have an effect on  $\Delta U$ . However, such an assumption might not hold. See App. A for limitations.
We empirically assess the performance of Voronoi sampling in a series of experiments. In each experiment, we perform a grid search to find the best set of hyperparameters, these experimental details can be found in App. H. Our open-source implementation will be available upon publication.
7.1 Toy Example
We first apply our Voronoi sampling $^{8}$  method on the toy example discussed earlier in Example 1, where the reference probability distribution is tractable and known  $p(\mathbf{x})$ . The potential energy is then set to  $U(\mathbf{x}) = -\log p_{\mathrm{V}}(\mathbf{x})$ . Importantly, the toy experiment is intentionally designed such that the base measure of all of the Voronoi cells is equal, therefore, we can safely ignore calculating the base measure and arrive at exact sampling methods.
We compare Voronoi sampling to MUCoLA and HMC. To make a fair comparison, we add the Metropolis criterion<sup>9</sup> to the MUCoLA algorithm and only do one leapfrog step in all algorithms. Furthermore, to see the effect of the reference distribution on the performance of the sampling algorithm, we anneal this distribution with 6 temperatures, where the lower temperatures lead to peaky distributions and the higher temperatures to uniform-like distributions.
We take 200 samples after 500 burn-in iterations, and compare the empirical distributions of samples to the reference by measuring the Jensen-Shannon (JS) divergence between the two. As results in Fig. 2 show, the JS divergence between the reference distribution and the empirical distribution of Voronoi samples is the smallest. The difference between the methods is more pronounced at lower temperatures, as the change in potential energy is greater, resulting in more errors in the leapfrog integrator. In App. I we provide more empirical support that Voronoi sampling converges faster to the reference distribution, especially when we increase the dimensionality of the sampling space.
7.2 Sampling from Language Models
Next, we apply our method to sample from a language model. The underlying LM is a finetuned GPT-2 $^{10}$  on E2E dataset [34]; see App. G for dataset statistics. As opposed to the previous
Note that in the case of the toy experiment, we are not sampling a sequence of text, but rather a single embedding. To highlight this point, we use Voronoi sampling instead of structured Voronoi sampling to refer to our algorithm in this section.
We accept transitioning from  $\mathbf{x}^t$  to  $\mathbf{x}^{t + 1}$  with probability  $e^{H^t - H^{t + 1}}$ .
<sup>10</sup>We use gpt2 checkpoint from the Huggingface library [47].
experiment, the reference distribution of the LM is not tractable. Therefore, we use the empirical distribution of ancestral samples as an unbiased estimate of the reference distribution. Note that ancestral sampling incrementally draws samples from a language model, where at each step of the generation a token  $w_{n}$  is sampled with the probability given by the LM:  $p(w_{n} \mid w_{< n})$ . Therefore, the process can give unbiased estimates of the reference distribution.
We follow §4 to define a structured Voronoi measure  $p_{\mathrm{V}}$  using the LM probability. We then implement SVs, where the potential energy is set to  $-\log p_{\mathrm{V}}$ . To empirically measure the benefit of reflection-refraction step in SVs, we compare it to applying Langevin dynamics directly to  $p_{\mathrm{V}}$ . We implement MUCOLA as a baseline, which writes the potential energy using an embedding-augmented LM, i.e., Eq. (3).
We show the distribution of samples' perplexity in Fig. 3. The green trace is the empirical distribution of 1000 ancestral samples. While all the sampling methods result in distributions comparably close to ancestral samples' distribution, we observe that SVS manages to model the tail of the distribution better. On the other hand, MUCOLA and LANGEVIN tend to take samples from the mode of the distribution more often.
7.3 Controlled Generation
Finally, we apply our structured Voronoi sampling to 2 controlled generation task. The goal of the first task is to generate restaurant reviews for a target food type  $t$ , e.g., Italian, Fast food, Japanese, etc. The goal of the second task is to control the sentiment of the generations to enforce a positive sentiment. We train classifiers to predict the target  $t$  (food type or positive sentiment) from the input sequence  $p(t \mid w)$ .<sup>11</sup> We implement two baselines:
FUDGE. Yang and Klein [48] offer a heuristic approach to sample from the conditional distribution. They incrementally sample tokens under the language model. At each sampling step, they adjust the probabilities given by the LM, by feeding each candidate prefix to a classifier and obtain the probability of that prefix following the control target.
MUcOLA. Kumar et al. [21] treat  $p(\mathbf{V}\mid t)$  , Eq. (4), as a distribution in  $\mathbb{R}^d$  and apply Langevin dynamics directly to sample a sequence of embeddings  $\mathbf{V}$  . The potential energy is defined as  $-\log p(\mathbf{V}\mid t)$  . When rewriting this potential energy with Bayes' rule, it has been shown empirically, that adding a hyperparameter  $\gamma$  is helpful to keep the balance between the classifier and LM. Therefore, the final potential energy is defined as:
$$
U (\mathbf {V}) \stackrel {\text {d e f}} {=} - \log p (\mathbf {V}) - \gamma \log p (t \mid \mathbf {V}). \tag {17}
$$
As mentioned earlier,  $p(\mathbf{V} \mid t)$  only places a positive probability on a countable set. We, therefore, use Def. 3 to define structured Voronoi measures and set the potential energy to
$$
U (\mathbf {x}) \stackrel {\text {d e f}} {=} - \log p _ {\mathrm {V}} (\mathbf {x}) - \gamma \log p _ {\mathrm {V}} (t \mid \mathbf {x}). \tag {18}
$$
We then apply Langevin dynamics and SVS to sample according to this potential energy.
Evaluation. We sample 120 sentences of length  $20^{12}$  and evaluate the generations on three metrics:
- Success: is defined as the percentage of generations that adhere to the control target. To determine whether a generation conforms to the specified target we use an evaluator classifier.
- Fluency: is measured by the mean and standard deviation of perplexity under the language model.
- Diversity: is measured by the mean number of distinct  $n$ -grams ( $n = 1,2,3$ ) in a set of samples, normalized by the length of the sequence.
As results in Table 1 show, $^{13}$  FUDGE tends to achieve the highest diversity, however, it fails to follow the control target. MUCOLA either generates fluent results without paying enough attention to the control, or sacrifices fluency in favor of following the control target; thus, the high variance in success rates. Both LANGEVIN and SVS result in a high success rate and maintain fluency and diversity, and SVS is effective in maintaining a balance between various metrics and producing fluent sentences that adhere to control targets.
<sup>11</sup>See App. H for more experimental details about classifiers.
12We sample 20 sentences per control target.
<sup>13</sup>Please refer to Fig. 7 to see a visualization of the topic control results, and to Table 6 for results per target type.
Topic Control	Sentiment Control
Success(↑)	PPL(↓)	Dist-1(↑)	Dist-2(↑)	Dist-3(↑)	Success(↑)	PPL(↓)	Dist-1(↑)	Dist-2(↑)	Dist-3(↑)
GPT2	0.12 ± 0.10	5.10 ± 2.06	0.40	0.56	0.67	0.55 ± 0.49	21.33 ± 27.17	0.40	0.60	0.71
FUDGE	0.30 ± 0.12	5.59 ± 0.60	0.39	0.55	0.65	0.57 ± 0.49	24.27 ± 12.46	0.40	0.60	0.70
MUCOLA	0.58 ± 0.23	33.09 ± 36.32	0.26	0.40	0.51	0.66 ± 0.47	85.74 ± 152.18	0.28	0.42	0.53
LANGEVIN	0.91 ± 0.12	14.26 ± 2.55	0.24	0.39	0.51	0.82 ± 0.38	26.76 ± 14.42	0.16	0.30	0.41
SVS	0.92 ± 0.05	13.9 ± 2.04	0.22	0.37	0.49	0.84 ± 0.36	32.73 ± 16.70	0.14	0.28	0.41
Table 1: Evaluation of different sampling methods on controlled generation, using three criteria: success in following the control target (measured by the evaluator classifier), fluency (measured by perplexity), and diversity.
Controlled Generation. Numerous approaches have been proposed to enforce controls during the text generation process [19, 24, 41, inter alia]. For example, weighted decoding [10, 17] scores each candidate token with a weighted sum of its score under the language model and its adherence to control targets, subsequently selecting candidates with the highest scores. FUDGE method adopts a similar scoring function, resembling a Bayesian formulation for  $p(\boldsymbol{w} \mid t)$ . After making simplifying assumptions and factorizing  $p(\boldsymbol{w} \mid t)$ , FUDGE samples tokens autoregressively based on their scores. More recently, a line of research attempts to directly sample from  $p(\boldsymbol{w} \mid t)$  by reformulating it as an energy-based model and sampling from it using efficient gradient-based sampling algorithms. As discussed in §5.2 COLD reformulates  $p(\boldsymbol{w} \mid t)$  as an energy-based model on the logit space and uses that to select samples with high energy from a number of candidate generations. MuCoLA offers a sampling algorithm motivated by Langevin Dynamics that operates in the embedding space.
Gradient-based Sampling. Our work is closely related to the line of research that makes use of gradient information to sample from complex distributions [7, 31, 46]. Gradient-based samplers [15, 32] are shown to be highly effective when sampling from continuous distributions [3, 4, 37]. However, it is a difficult problem to adapt gradient-based samplers to discrete settings [36, 50]. More recently, several papers proposed promising gradient-based MCMC for discrete distribution that are reversible chains [11, 39, 49]. Our work instead formulates an irreversible Markov chain based on HMC. We leave it to future work to explore the utility of these recent papers on text generation.
In this work, we propose structured Voronoi sampling, a principled gradient-based sampling method for text generation. To formulate the energy function used in SVs, we define structured Voronoi measures on the embedding space and show how such measures can encode language models. In a controlled generation task, SVs outperformed other sampling methods in following the control target while producing comparably fluent and diverse samples.
It has been repeatedly shown that LMs can generate harmful, toxic, or non-factual content [9, 35, 42]. In fact, an application of the controlled generation scheme discussed in this paper could be to mitigate such issues. However, the same method could be used to generate misinformation, or toxic content intentionally.
We thank Tim Vieira, Tiago Pimentel, and Clara Meister for their feedback on this work. We also thank the anonymous reviewers for their constructive feedback during the review process. Afra Amini is supported by ETH AI Center doctoral fellowship.
[47] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumont, Clement Delangue, Anthony Moi, Pierrick Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online. Association for Computational Linguistics.
[48] Kevin Yang and Dan Klein. 2021. FUDGE: Controlled text generation with future discriminators. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3511-3535, Online. Association for Computational Linguistics.
[49] Ruqi Zhang, Xingchao Liu, and Qiang Liu. 2022. A Langevin-like sampler for discrete distributions. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 26375-26396. PMLR.
[50] Yichuan Zhang, Zoubin Ghahramani, Amos J. Storkey, and Charles Sutton. 2012. Continuous relaxations for discrete hamiltonian monte carlo. In Advances in Neural Information Processing Systems, volume 25.
Approximating the base measure. In this work, we go one step closer to implementing a principled approach for text generation. However, in our text generation experiments, we follow the previous work in using uncorrected Langevin dynamics. Moreover, when implementing SVs for text generation, we assume all Voronoi cells to have an equal base measure, which might not hold.
Efficiency. As reported in App. H, all gradient-based samplers are considerably slower than ancestral sampling and heuristics such as FUDGE. Further efforts are needed to make gradient-based methods faster.
Text Quality. As shown in previous work, sequences with high probability under the LM can be repetitive, dull, or degenerate [16]. Therefore, as with any other sampling method, svS might sample degenerate sentences, and the quality of the samples depends on the LM probability distribution. This is an active area of research with various proposals from changing the loss function during training [45], to modifying the decoding objective [28].
B A Note on the Choice of the Gaussian Measure
We thank the reviewers and the meta reviewer for their careful assessment of this work and their thoughtful feedback. As a response to an important point raised by the meta reviewer, here we explain the reasoning behind using a Gaussian measure in Eq. (12). The gradient  $\mathbf{g}_m$  is only defined at Voronoi centers and the goal is to follow a direction similar to  $\mathbf{g}_m$  when we are in the vicinity of the Voronoi center, i.e., in the Voronoi cell. We note that this direction needs to be adjusted depending on the position in the Voronoi cell  $\mathbf{x}$ . A natural choice for this would be to use a truncated Gaussian that is centered at  $\mathbf{g}_m$ , which implies that the gradient in  $\mathbf{x}$  should be  $\mathbf{g}_m - \mathbf{x}$ , to follow a similar direction  $\mathbf{g}_m$  at the center of the Voronoi.
Proposition 1. Let  $\pmb{p} = [p_1, \dots, p_M]$  be an embedding-augmented distribution with embeddings  $\{\mathbf{v}_m\}_{m=1}^M \subset \mathbb{R}^d$ , and let  $p_V$  be the corresponding Voronoi measure Eq. (7). Then,  $p_V(C_m) = p_m$  where  $C_m$  is defined as in Eq. (6). See App. C.1 for proof.
$$
\begin{array}{l} p _ {\mathrm {V}} \left(C _ {m}\right) = \int_ {C _ {m}} p _ {\mathrm {V}} (\mathbf {x}) \mathrm {d} \lambda (19a) \\ = \int_ {C _ {m}} \frac {p _ {m}}{\mu \left(C _ {m}\right)} \frac {\mathrm {d} \mu}{\mathrm {d} \lambda} (\mathbf {x}) \mathrm {d} \lambda (19b) \\ = \frac {p _ {m}}{\mu \left(C _ {m}\right)} \int_ {C _ {m}} \frac {\mathrm {d} \mu}{\mathrm {d} \lambda} (\mathbf {x}) \mathrm {d} \lambda (19c) \\ = \frac {p _ {m}}{\mu \left(C _ {m}\right)} \mu \left(C _ {m}\right) = p _ {m} (19d) \\ \end{array}
$$
Proof. Let  $\pmb{m} = [m_1, \dots, m_N] \in [M]^N$ . We have
$$
\begin{array}{l} \mu \left(C _ {\boldsymbol {m}}\right) = \mu \left(\prod_ {n = 1} ^ {N} C _ {m _ {n}}\right) (20a) \\ = \prod_ {n = 1} ^ {N} \mu \left(C _ {m _ {n}}\right) (20b) \\ \end{array}
$$
Proof. Choose  $E \in \mathcal{B}(\mathbb{R}^d)$  such that  $\lambda(E) = 0$ .<sup>14</sup> Note that  $E = (E \setminus \mathcal{K}) \cup \bigcup_{m=1}^{M} (E \cap C_m)$  where  $\mu(E \setminus \mathcal{K}) = 0$  by Eq. (13). Since the Gaussian measure over  $\mathbb{R}^d$  itself is absolutely continuous with respect to  $\lambda$ , we can also conclude that  $\mu(E \cap C_m) = 0$  for any  $m$ . Hence,  $\mu(E) \leq \mu(E \setminus \mathcal{K}) + \sum_{m} \mu(E \cap C_m) = 0$  which means  $\mu(E) = 0$ . So  $\mu \ll \lambda$ .
Proposition 4. The gradient of the log of the Voronoi measure  $p_{\mathbb{V}}$  is given by
Proof. We have three cases. Suppose  $\mathbf{x} \in \operatorname{int}(C_m)$  for some  $m \in [M]$ . Then, by Observation 1, we have that  $p_{\mathrm{V}}(\mathbf{x})$ , and, thus,  $\log p_{\mathrm{V}}(\mathbf{x})$  is differentiable. Direct computation reveals: