Title: Structured Voronoi Sampling
Venue: NeurIPS 2023
Accepted: yes
Keywords: Natural Language Processing

Text Generation

Controlled Generation

MCMC

HMC

Langevin Dynamics
Abstract: Gradient-based sampling algorithms have demonstrated their effectiveness in text generation, especially in the context of controlled text generation. However, there exists a lack of theoretically grounded and principled approaches for this task. In this paper, we take an important step toward building a principled approach for sampling from language models with gradient-based methods. We use discrete distributions given by language models to define densities and develop an algorithm based on Hamiltonian Monte Carlo to sample from them. We name our gradient-based technique Structured Voronoi Sampling (SVS). In an experimental setup where the reference distribution is known, we show that the empirical distribution of SVS samples is closer to the reference distribution compared to alternative sampling schemes. Furthermore, in a controlled generation task, SVS is able to generate fluent and diverse samples while following the control targets significantly better than other methods.

Structured Voronoi Sampling

Abstract

Gradient-based sampling algorithms have demonstrated their effectiveness in text generation, especially in the context of controlled text generation. However, there exists a lack of theoretically grounded and principled approaches for this task. In this paper, we take an important step toward building a principled approach for sampling from language models with gradient-based methods. We use discrete distributions given by language models to define densities and develop an algorithm based on Hamiltonian Monte Carlo to sample from them. We name our gradient-based technique Structured Voronoi Sampling (SVS). In an experimental setup where the reference distribution is known, we show that the empirical distribution of SVS samples is closer to the reference distribution compared to alternative sampling schemes. Furthermore, in a controlled generation task, SVS is able to generate fluent and diverse samples while following the control targets significantly better than other methods.

https://github.com/AfraAmini/svs

1 Introduction

Gradient-based sampling algorithms such as Hamiltonian Monte Carlo [HMC; 32] and Langevin dynamics [46] are widely used in Bayesian inference due to their efficiency in drawing samples from high-dimensional space [4]. Such algorithms construct a Markov Chain that has the desired distribution as its stationary distribution and use the gradient information of this distribution to efficiently navigate the state space. Additionally, gradient-based sampling schemes have recently been deployed in computer vision to generate high-quality images from state-of-the-art models [13, 43], and are a popular choice for tasks such as image synthesis [5] and image-to-image translation [40].

In natural language processing, there have been several attempts to apply gradient-based sampling techniques to sampling text from neural language models [21, 23, 38]. The motivation behind this approach to text generation is to sample from energy-based probabilistic models, where the normalization factor is not tractable. One such case is in controlled text generation, where energy functions are usually defined as a linear combination of LM probabilities and the probability of satisfying a set of predefined constraints [21, 23, 38]. In contrast to computer vision, however, applying gradient-based sampling schemes to text generation is nuanced as text, in contrast to images, is discrete.

Upon closer inspection, none of the proposed algorithms actually defines a valid Markov Chain Monte Carlo [MCMC; 12, 29] scheme that will draw samples from the model in the limit. For instance, Qin et al. [38] relax the language model from a distribution over strings to a distribution over logits. While the relaxation does transform the language model into a continuous distribution, it introduces bias. Kumar et al. [MUCOLA; 21] take a different approach. They derive a constrained gradient-based

sampler where the constraint is enforced through a projection. However, the projection invalidates the MCMC procedures, leading to an algorithm without guarantees. 1

We derive a principled Hamiltonian Monte Carlo scheme for generating text from language models, which we call a structured Voronoi sampler. Our scheme consists of two steps. First, we give a recipe for encoding discrete distributions as densities over R d ; we term the resulting encoding Voronoi measures. 2 Second, we derive a refractive Hamiltonian Monte Carlo algorithm [30] for sampling from an arbitrary Voronoi measure. In our theoretical analysis, we show that, despite the presence of discontinuities, we are able to give proof that our sampler satisfies the detailed balance condition and, thus, is a correct MCMC scheme.

To empirically evaluate the performance of structured Voronoi sampling, we begin by applying it to a toy example where the exact reference probability distribution is known. We compare the empirical distribution of drawn samples with the reference distribution and show Voronoi sampler's distribution is closer to the reference distribution than MUCOLA or unconstrained HMC. Furthermore, we use our sampling scheme for controlled generation, where the goal is to use GPT-2 to generate restaurant reviews for a target type of food, e.g., Italian, Fast food, or Japanese, and separately to generate text with a positive sentiment. We find that structured Voronoi sampling outperforms FUDGE [48], MUCOLA, and Langevin Dynamics algorithms in terms of adhering to the control target. Additionally, the samples generated by structured Voronoi sampling are comparably fluent and diverse to those produced by the other methods.

2 Language Models

Let Σ be an alphabet, a finite, non-empty set. By Σ ∗ def = ⋃ ∞ n =0 Σ n , we denote the Kleene closure of Σ . 3 A probability distribution over Σ ∗ is called a language model ( LM ). The elements of Σ may be characters, subword pieces, or words; the choice lies with the modeler. Language models can be factored autoregressively by means of the chain rule of probability, i.e., for any string w = w 1 · · · w N ∈ Σ ∗ , we can write

where EOS ̸∈ Σ is a distinguished end-of-sequence token and w <n is the prefix of length ( n -1) of the string w . We define Σ def = Σ ∪ { EOS } , and require the conditional distributions p ( · | w <n ) to be defined over Σ . While all language models can be factored autoregressively, not all conditionals p ( · | w <n ) can be assembled into a language model. In some cases, probability mass must be placed on infinite sequences [6]. In this work, we assume working with tight language models, i.e., that they indeed define valid probability distributions over Σ ∗ .

2.1 Language Modeling with Embeddings

Most neural language models make use of embeddings. Moreover, in most language models, the weights are shared between the language model head and the embedding layer. Such an embedding-based language model is defined as follows

where v w ∈ R d is the embedding of w and enc : Σ ∗ → R d is a real-valued encoding of the context. Notably, the context embedding enc ( w <n ) ∈ R d is obtained by inputting the context w <n into the language model, converting it to embeddings v w <n , passing it through neural network layers, and extracting the encoding from the model at position n -1 .

Interestingly, one can lift an embedding-based language model to be a distribution over the set of base embeddings : B = { v w : w ∈ Σ } . Specifically, we can associate a sequence of embeddings V = [ v 1 , . . . , v N ] to any string w = w 1 · · · w N . Substituting each w with the corresponding v w , we can rewrite Eq. (1) as

This transformation encodes a language model as a distribution over real-valued vectors. However, p ( V ) only places a positive probability on a countable set, and is zero everywhere else.

2.2 Controlled Language Modeling with Embeddings

In a controlled generation task, we are interested in a subset of strings w that have a target property t . Therefore, we want to model and sample from a conditional distribution p ( w | t ) , e.g., sample a sentence given a topic. Following Bayes' rule, one can write p ( w | t ) ∝ p ( t | w ) p ( w ) . Previous papers model p ( t | w ) with an embedding-augmented classifier. Such a classifier receives embeddings V associated with w and predicts the probability of the target t . Notably, if the classifier and the LM share the same base embeddings B , the controlled LM can also be lifted as a distribution over the base embeddings

where Z t = ∑ V p ( t | V ) p ( V ) is an intractable normalization factor, that sums over the embeddings of all possible strings. 4 Identically to p ( V ) , p ( V | t ) only places a positive probability on a countable set.

3 Voronoi Measures

In this section, we demonstrate how to encode an embedding-based language model as a density that places positive probability on a set with a measure greater than zero. Such encoding allows us to derive a principled gradient-based sampling approach to generate samples in §6. We start with some definitions.

Definition 1. An embedding-augmented probability distribution over the first M positive integers [ M ] is an array p = [ p 1 , . . . , p M ] such that p m ≥ 0 and ∑ M m =1 p m = 1 where we assume that there is a real-valued embedding { v m } M m =1 ⊂ R d associated with each m ∈ [ M ] .

Embedding-augmented distributions can be viewed as densities over R d using the following simple encoding

Eq. (5), however, yields a density that is 0 almost everywhere (with respect to the standard Lebesgue measure) and its gradient with respect to p m is also zero almost everywhere. Thus, Eq. (5) is not amenable to gradient-based sampling, and to derive a meaningful gradient-based sampling we require a more nuanced encoding.

To provide such an encoding, we introduce the Voronoi measure . Given an embedding-augmented distribution p = [ p 1 , . . . , p M ] with embeddings { v m } M m =1 , and a compact set K ⊂ R d that covers the embeddings, i.e., { v m } M m =1 ⊂ K , we define the Voronoi cell for the m th item with respect to the compact set K as follows

̸

Now, using the definition of a Voronoi cell C m given in Eq. (6), we can define a density that is not zero almost everywhere as follows. The strategy is to spread out the probability mass p m over the entirety of the set C m . To do so, we assume access to a set of base measures { µ m } M m =1 that give us a reference for how to judge the probability mass in each C m . We make this encoding formal in the following definition.

Definition 2. Let p = [ p 1 , . . . , p M ] be an embedding-augmented distribution with embeddings { v m } M m =1 ⊂ R d , and let K be a compact set such that { v m } M m =1 ⊂ K . Furthermore, let { µ m } M m =1 be a set of base measures over R d that are absolutely continuous with respect to the standard Lebesgue measure λ over R d , i.e., µ m ≪ λ . Define the ( K , µ ) -Voronoi measure as follows

where we define projection

with ties broken arbitrarily. 5

In the following proposition, we make precise the sense in which a Voronoi measure encodes the original embedding-augmented distribution.

Proposition 1. Let p = [ p 1 , . . . , p M ] be an embedding-augmented distribution with embeddings { v m } M m =1 ⊂ 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.

Example 1. Suppose p = [ p 1 , . . . , p 4 ] is a categorical distribution, and there are 4 embeddings in R 2 associated with each p i , namely: v 1 = [1 , 1] , v 2 = [ -1 , 1] , v 3 = [ -1 , -1] , v 4 = [1 , -1] . Given the K = [ -2 , 2] × [ -2 , 2] and the embedding-augmented probability distribution p , Eq. (7) defines a Voronoi measure over this space, where the Voronoi

0.11

0.02

0.31

0.56

=

p

(

v

1

)

p

(

v

2

) =

p

(

v

3

) =

p

(

v

4

) =

∥

r

⊥

∥ 2

< 2Δ

U

∥

r

⊥∥ 2

> 2Δ

U

refract

reflect

0

10

20

30

40

20

40

60

80

100

1.

0.8

0.6

0.4

0.2

C

2

C

3

C

4

C

1

2

0

-2

2

0

/u1D4A6

40

30

20

10

0

cells are visualized in Fig. 1. We will discuss in §6 how Voronoi sampling navigates this space.

3.1 Structured Voronoi Measures

To encode language models as densities more naturally, we introduce a generalization of the Voronoi measure, which we term a structured Voronoi measure. Now, rather than a distribution over M elements, we assume to have a sequence of length N . Each token in the sequence takes value in [ M ] . Let m = [ m 1 , . . . , m N ] ∈ [ M ] N . We define a structured Voronoi cell as C m = ∏ N n =1 C m n , where ∏ denotes the Cartesian product and we define the individual Voronoi cell as

̸

Proposition 2. Let µ be a measure on R d . Then, we have the product measure space as µ ( C m ) = ∏ N n =1 µ ( C m n ) . See App. C.2 for proof.

Definition 3. Let p be an embedding-augmented distribution over [ M ] N . For m ∈ [ M ] N , we denote m 's probability as p m , and m 's embedding as V m ∈ R N × d . Let K be a compact set that covers the embeddings V m and let µ ≪ λ be a base measure absolutely continuous with respect to the Lebesgue measure λ . We define the ( K , µ ) -structured Voronoi measure as follows

where we define the structured projection

4 Application to Text Generation

Def. 3 gives us the flexibility to use any probability distribution over sequences of embeddings to define a structured Voronoi measure. For example, one can substitute p m ⋆ ( x ) with an embedding-augmented LM, i.e., Eq. (3), to encode a language model as a structured Voronoi measure. Another example is to encode a controlled LM as a structured Voronoi measure by substituting p m ⋆ ( x ) with Eq. (4).

Base Measure. We have explained how to encode our desired distribution as a structured Voronoi measure. However, in order to actually implement a gradient-based sampler, we need to specify the base probability measures { µ m } M m =1 . Given an embedding-augmented probability p ( V ) it is natural to follow the gradient of log p ( V ) with respect to the word embedding, i.e., g m = ∇ V log p ( V ) . Thus, if we want to follow a direction similar to g m , one natural choice for µ is a Gaussian measure centered at the gradient g m restricted to Voronoi cell C m , 6 which we define:

The normalizer ensures the measure is a probability measure. Furthermore, we have that µ 's Radon-Nikodym derivative with respect to the Lebesgue measure λ is given by

Eq. (13) should be recognizable as the standard Gaussian density, albeit one that is truncated to the Voronoi cell C m [27].

Proposition 3. Eq. (13) is absolutely continuous with respect to the Lebesgue measure λ . See App. C.3 for proof.

Proposition 4. The gradient of the log of the Voronoi measure p V is given by

where the first two blocks in the case statement apply if there exists some m such that x ∈ int( C m ) or x ∈ ∂C m . See App. C.4 for proof.

5 Gradient-Based Sampling

In this section, we first review gradient-based sampling methods, namely HMC and Langevin dynamics. Then we discuss how they have been used in previous papers on text generation. This will set the stage for our algorithm in §6, which is based on HMC.

5.1 Hamiltonian Monte Carlo

The goal of HMC, originally proposed by Duane et al. [7], is to design a better proposal distribution in the standard Metropolis-Hastings MCMC by taking advantage of the gradient information in a principled way. Concretely, to sample from a given distribution p ( x ) where x ∈ R d , HMC treats x as the coordinates of the particles in some fictitious physical system. It then introduces an auxiliary momentum variable r ∈ R d associated with each coordinate and defines a Hamiltonian function H ( x , r ) . Here, the Hamiltonian H has the intuitive physical interpretation of the total energy of some conservative system, and, in classical mechanics, decomposes into the potential energy U ( x ) and the kinetic energy K ( r ) , i.e., H ( x , r ) = U ( x ) + K ( r ) . This formulation is convenient in part because if we define the joint distribution p ( x , r ) ∝ e -H ( x , r ) as in energy-based models, then

which means we can treat x and r as independent variables. Naturally, we can let U ( x ) = -log p ( x ) so that x has the marginal of the target distribution. It is also common practice to set K ( r ) = r ⊤ M -1 r / 2 so that the momentum variable has a Gaussian distribution. Here, M ∈ R d × d is called the mass matrix and commonly set to identity. 7

The Hamiltonian H determines the equations of motion in a physical system, given by the Hamiltonian equations, which is also known as Hamiltonian dynamics,

We are now ready to give a high-level description of how HMC generates a single sample (see Algorithm 1): First sample a momentum r from a Gaussian (line 1), then evolve the system using the Hamiltonian equations Eq. (16) for some predetermined amount of time (lines 4 to 7), and finally accept the new state with the MetropolisHastings acceptance criterion (lines 8 to 13). Note that the Hamiltonian equations often admit no closed-form solution in practice, and hence one needs to use numerical integrators for approximation. In particular, the leapfrog integrator, corresponding to lines 5 to 7, is almost always used in HMC.

Ingeniously designed by Duane et al. [7], the efficiency of this elegant procedure is a result of several favorable properties of Hamiltonian mechanics-namely volume preservation, reversibility, and conservation. Upon exact simulation of the Hamiltonian dynamics, these properties will lead to the acceptance probability of one, i.e., every proposed move will be accepted. In App. D, we will give intuition about these properties. Since the method developed later in §6 will subsume HMC as a special case, we delay the proof of correctness until then.

Langevin Dynamics. Langevin dynamics is a simplification of HMC, where only one leapfrog step is performed. Furthermore, what is usually referred to as Langevin dynamics is in fact the uncorrected Langevin dynamics, where the acceptance criterion is ignored, i.e., every proposed sample is accepted; see Algorithm 2. While uncorrected Langevin dynamics is guaranteed to converge when the energy function is smooth [31, §5.3], there is no such guarantee with the presence of a discontinuity in the energy function. Nevertheless, due to its simplicity, Langevin dynamics is the only gradient-based sampling method that has been applied for text generation applications.

5.2 Applications to Controlled Generation

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 ( V ) = -log p ( V | t ) is two terms: (1) the gradient of U with respect to V : ∇ U , and (2) the difference between the potential energy of two points ∆ 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 ( V | t ) is R d -valued, so it is tempting to apply a gradientbased 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: x t : current sample, U : potential energy function, ε : step size, L : number of leapfrog steps Output: next sample: x t +1

)

U ( V ) def = -log p ( V | t ) and apply Langevin dynamics out of the box; see Algorithm 3. However, since U ( V ) is zero for any vector other than vectors in 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 topk 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 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 V . Importantly, p 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 x with momentum r . We can decompose momentum as r = r ⊥ + r ∥ where r ⊥ is normal to the barrier and r ∥ parallel to it. Let ∆ U be the signed potential energy difference between the two sides of the barrier. If ∥ r ⊥ ∥ 2 2 > 2∆ U , then the particle has enough kinetic energy to overcome the barrier, and its momentum's normal component will instantaneously become r ′ ⊥ = √ ∥ r ⊥ ∥ 2 2 -2∆ U · r ⊥ ∥ r ⊥ ∥ 2 2 after being transmitted through the barrier (refraction). Otherwise, if ∥ r ⊥ ∥ 2 2 ≤ 2∆ U , the particle will be reflected from the barrier and the normal component will instantaneously become r ′ ⊥ = -r ⊥ . 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 V are all piecewise smooth, we can build on the above and develop a sampling algorithm for p 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 ( x , 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 ∆ 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

-6GLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>'LYHUJHQFH

-6GLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>'LYHUJHQFH

5HIHUHQFHGLYPH<c=3,font=/AAAAAE+TimesNewRomanPSMT>3UREDELOLW\GLYPH<c=3,font=/AAAAAE+TimesNewRomanPSMT>'LVWULEXWLRQ

5HIHUHQFHGLYPH<c=3,font=/AAAAAG+TimesNewRomanPSMT>3UREDELOLW\GLYPH<c=3,font=/AAAAAG+TimesNewRomanPSMT>'LVWULEXWLRQ

GLYPH<c=21,font=/AAAAAL+TimesNewRomanPSMT>

5HIHUHQFHGLYPH<c=3,font=/AAAAAI+TimesNewRomanPSMT>3UREDELOLW\GLYPH<c=3,font=/AAAAAI+TimesNewRomanPSMT>'LVWULEXWLRQ

GLYPH<c=22,font=/AAAAAL+TimesNewRomanPSMT>

7HPSHUDWXUH

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=26,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAB+TimesNewRomanPSMT>

6DPSOLQJGLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>0HWKRG

0X&R/D

+0&GLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

9RURQRLGLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=26,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAB+TimesNewRomanPSMT>

6DPSOLQJGLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>0HWKRG

0X&R/D

+0&GLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

9RURQRLGLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

7HPSHUDWXUH

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=26,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAB+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAB+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAB+TimesNewRomanPSMT>

6DPSOLQJGLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>0HWKRG

0X&R/D

+0&GLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

9RURQRLGLYPH<c=3,font=/AAAAAB+TimesNewRomanPSMT>6DPSOLQJ

-6GLYPH<c=3,font=/AAAAAC+TimesNewRomanPSMT>'LYHUJHQFH

GLYPH<c=19,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAD+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAD+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAD+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=21,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=22,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAF+TimesNewRomanPSMT>GLYPH<c=23,font=/AAAAAF+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=22,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAH+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAH+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAH+TimesNewRomanPSMT>

1

2

3

1

2

3

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=21,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>

GLYPH<c=19,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=17,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=20,font=/AAAAAL+TimesNewRomanPSMT>GLYPH<c=24,font=/AAAAAL+TimesNewRomanPSMT>

5HIHUHQFHGLYPH<c=3,font=/AAAAAM+TimesNewRomanPSMT>3UREDELOLW\GLYPH<c=3,font=/AAAAAM+TimesNewRomanPSMT>'LVWULEXWLRQ

5HIHUHQFHGLYPH<c=3,font=/AAAAAM+TimesNewRomanPSMT>3UREDELOLW\GLYPH<c=3,font=/AAAAAM+TimesNewRomanPSMT>'LVWULEXWLRQ

7HPSHUDWXUH

2

4

6

8

10

12

14

SVS

Ancestral

MuCoLa

Perplexity

assumption that the base measures of Voronoi cells are equal, and thus do not have an effect on ∆ U . However, such an assumption might not hold. See App. A for limitations.

7 Experiments

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 ( x ) . The potential energy is then set to U ( x ) = -log p V ( 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 9 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

.

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 | w <n ) . Therefore, the process can give unbiased estimates of the reference distribution.

We follow §4 to define a structured Voronoi measure p V using the LM probability. We then implement SVS, where the potential energy is set to -log p V . To empirically measure the benefit of reflection-refraction step in SVS, we compare it to applying Langevin dynamics directly to p 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 | w ) . 11 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 ( V | t ) , Eq. (4), as a distribution in R d and apply Langevin dynamics directly to sample a sequence of embeddings V . The potential energy is defined as -log p ( V | t ) . When rewriting this potential energy with Bayes' rule, it has been shown empirically, that adding a hyperparameter γ is helpful to keep the balance between the classifier and LM. Therefore, the final potential energy is defined as:

As mentioned earlier, p ( V | 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

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.

8 Related Work

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 ( w | t ) . After making simplifying assumptions and factorizing p ( w | t ) , FUDGE samples tokens autoregressively based on their scores. More recently, a line of research attempts to directly sample from p ( w | 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 ( w | 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.

9 Conclusion

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.

Broader Impacts

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.

References

Luigi Ambrosio. 2008. Transport Equation and Cauchy Problem for Non-Smooth Vector Fields , pages 1-41. Springer Berlin Heidelberg, Berlin, Heidelberg.

V. I. Arnold. 1989. Mathematical Methods of Classical Mechanics . Springer, New York, NY.

Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman. 2018. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research .

Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. 2017. Stan: A probabilistic programming language. Journal of statistical software , 76(1).

Prafulla Dhariwal and Alexander Nichol. 2021. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems , volume 34, pages 8780-8794. Curran Associates, Inc.

Li Du, Lucas Torroba Hennigen, Tiago Pimentel, Clara Meister, Jason Eisner, and Ryan Cotterell. 2022. A measure-theoretic characterization of tight language models.

Simon Duane, A. D. Kennedy, Brian J. Pendleton, and Duncan Roweth. 1987. Hybrid Monte Carlo. Physics Letters B , 195(2):216-222.

J. Michael. Finn. 2010. Jones and Bartlett Publishers, Sudbury, MA. [link].

Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A. Smith. 2020. RealToxicityPrompts: Evaluating neural toxic degeneration in language models. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 3356-3369, Online. Association for Computational Linguistics.

Marjan Ghazvininejad, Xing Shi, Jay Priyadarshi, and Kevin Knight. 2017. Hafez: an interactive poetry generation system. In Proceedings of ACL 2017, System Demonstrations , pages 43-48, Vancouver, Canada. Association for Computational Linguistics.

Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. 2021. Oops I took a gradient: Scalable sampling for discrete distributions. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 3831-3841. PMLR.

W. K. Hastings. 1970. Monte Carlo sampling methods using Markov chains and their applications. Biometrika , 57(1):97-109.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. 2020. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , volume 33, pages 6840-6851. Curran Associates, Inc.

Cong Duy Vu Hoang, Gholamreza Haffari, and Trevor Cohn. 2017. Towards decoding as continuous optimisation in neural machine translation. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing , pages 146-156, Copenhagen, Denmark. Association for Computational Linguistics.

Matthew D. Hoffman and Andrew Gelman. 2014. The No-U-Turn sampler: Adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research , 15(47):15931623.

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020. The curious case of neural text degeneration. In Proceedings of the 8th International Conference on Learning Representations .

Ari Holtzman, Jan Buys, Maxwell Forbes, Antoine Bosselut, David Golub, and Yejin Choi. 2018. Learning to write with cooperative discriminators. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1638-1649, Melbourne, Australia. Association for Computational Linguistics.

Shi Jin and Xin Wen. 2005. Hamiltonian-preserving schemes for the Liouville equation with discontinuous potentials. Communications in Mathematical Sciences , 3(3):285 - 315.

Ben Krause, Akhilesh Deepak Gotmare, Bryan McCann, Nitish Shirish Keskar, Shafiq Joty, Richard Socher, and Nazneen Fatema Rajani. 2021. GeDi: Generative discriminator guided sequence generation. In Findings of the Association for Computational Linguistics: EMNLP 2021 , pages 4929-4952, Punta Cana, Dominican Republic. Association for Computational Linguistics.

Sachin Kumar, Eric Malmi, Aliaksei Severyn, and Yulia Tsvetkov. 2021. Controlled text generation as continuous optimization with multiple constraints. In Advances in Neural Information Processing Systems .

Sachin Kumar, Biswajit Paria, and Yulia Tsvetkov. 2022. Constrained sampling from language models via langevin dynamics in embedding spaces. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics.

Cornelius Lanczos. 1949. The Variational Principles of Mechanics . Mathematical Expositions, No. 4. University of Toronto Press, Toronto.

Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. 2022. Diffusion-LM improves controllable text generation. In Advances in Neural Information Processing Systems .

Alisa Liu, Maarten Sap, Ximing Lu, Swabha Swayamdipta, Chandra Bhagavatula, Noah A. Smith, and Yejin Choi. 2021. DExperts: Decoding-time controlled text generation with experts and anti-experts. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 6691-6706, Online. Association for Computational Linguistics.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2020. RoBERTa: A robustly optimized BERT pretraining approach.

Aaron Lou and Stefano Ermon. 2023. Reflected diffusion models. In Proceedings of the 40th International Conference on International Conference on Machine Learning .

Hassan Maatouk and Xavier Bay. 2016. A new rejection sampling method for truncated multivariate Gaussian random variables restricted to convex sets. In Monte Carlo and QuasiMonte Carlo Methods , pages 521-530, Cham. Springer International Publishing.

Clara Meister, Tiago Pimentel, Gian Wiher, and Ryan Cotterell. 2022. Locally typical sampling. Transactions of the Association for Computational Linguistics .

Nicholas Metropolis, Arianna W. Rosenbluth, Marshall N. Rosenbluth, Augusta H. Teller, and Edward Teller. 1953. Equation of state calculations by fast computing machines. The Journal of Chemical Physics , 21(6):1087-1092.

Hadi Mohasel Afshar and Justin Domke. 2015. Reflection, refraction, and Hamiltonian monte carlo. In Advances in Neural Information Processing Systems , volume 28. Curran Associates, Inc.

Radford M. Neal. 1993. Probabilistic inference using Markov chain Monte Carlo methods . Department of Computer Science, University of Toronto Toronto, ON, Canada.

Radford M. Neal. 2011. MCMC using Hamiltonian dynamics. In Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng, editors, Handbook of Markov Chain Monte Carlo , chapter 5. Chapman and Hall/CRC.

Akihiko Nishimura, David B. Dunson, and Jianfeng Lu. 2020. Discontinuous Hamiltonian Monte Carlo for discrete parameters and discontinuous likelihoods. Biometrika , 107(2):365380.

Jekaterina Novikova, Ondˇ rej Dušek, and Verena Rieser. 2017. The E2E dataset: New challenges for end-to-end generation. In Proceedings of the 18th Annual SIGdial Meeting on Discourse and Dialogue , pages 201-206, Saarbrücken, Germany. Association for Computational Linguistics.

Artidoro Pagnoni, Vidhisha Balachandran, and Yulia Tsvetkov. 2021. Understanding factuality in abstractive summarization with FRANK: A benchmark for factuality metrics. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 4812-4829, Online. Association for Computational Linguistics.

Ari Pakman and Liam Paninski. 2013. Auxiliary-variable exact hamiltonian monte carlo samplers for binary distributions. In Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc.

Du Phan, Neeraj Pradhan, and Martin Jankowiak. 2019. Composable effects for flexible and accelerated probabilistic programming in numpyro. arXiv preprint arXiv:1912.11554 .

Lianhui Qin, Sean Welleck, Daniel Khashabi, and Yejin Choi. 2022. COLD decoding: Energybased constrained text generation with Langevin dynamics. In Advances in Neural Information Processing Systems .

Benjamin Rhodes and Michael U. Gutmann. 2022. Enhanced gradient-based MCMC in discrete spaces. Transactions on Machine Learning Research .

Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. 2022. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings , New York, NY, USA. Association for Computing Machinery.

Timo Schick, Sahana Udupa, and Hinrich Schütze. 2021. Self-Diagnosis and Self-Debiasing: A Proposal for Reducing Corpus-Based Bias in NLP. Transactions of the Association for Computational Linguistics , 9:1408-1424.

Emily Sheng, Kai-Wei Chang, Prem Natarajan, and Nanyun Peng. 2021. Societal biases in language generation: Progress and challenges. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 4275-4293, Online. Association for Computational Linguistics.

Yang Song and Stefano Ermon. 2020. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems , volume 33, pages 12438-12448. Curran Associates, Inc.

Molei Tao and Shi Jin. 2022. Accurate and efficient simulations of hamiltonian mechanical systems with discontinuous potentials. Journal of Computational Physics , 450:110846.

Sean Welleck, Ilia Kulikov, Stephen Roller, Emily Dinan, Kyunghyun Cho, and Jason Weston. 2020. Neural text generation with unlikelihood training. In International Conference on Learning Representations .

Max Welling and Yee Whye Teh. 2011. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning , page 681-688, Madison, WI, USA. Omnipress.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric 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: Stateof-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.

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.

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.

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.

A Limitations

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 g m is only defined at Voronoi centers and the goal is to follow a direction similar to g m when we are in the vicinity of the Voronoi center, i.e., in the V oronoi cell. We note that this direction needs to be adjusted depending on the position in the Voronoi cell x . A natural choice for this would be to use a truncated Gaussian that is centered at g m , which implies that the gradient in x should be g m -x , to follow a similar direction g m at the center of the V oronoi.

C Proofs

C.1 Proof of Proposition 1

Proposition 1. Let p = [ p 1 , . . . , p M ] be an embedding-augmented distribution with embeddings { v m } M m =1 ⊂ 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.

Proof.

C.2 Proof of Proposition 2

Proposition 2. Let µ be a measure on R d . Then, we have the product measure space as µ ( C m ) = ∏ N n =1 µ ( C m n ) . See App. C.2 for proof.

■

Proof. Let m = [ m 1 , . . . , m N ] ∈ [ M ] N . We have

■

C.3 Proof of Proposition 3

Proposition 3. Eq. (13) is absolutely continuous with respect to the Lebesgue measure λ . See App. C.3 for proof.

Proof. Choose E ∈ B ( R d ) such that λ ( E ) = 0 . 14 Note that E = ( E \K ) ∪ ⋃ M m =1 ( E ∩ C m ) where µ ( E \ K ) = 0 by Eq. (13). Since the Gaussian measure over R d itself is absolutely continuous with respect to λ , we can also conclude that µ ( E ∩ C m ) = 0 for any m . Hence, µ ( E ) ≤ µ ( E \ K ) + ∑ m µ ( E ∩ C m ) = 0 which means µ ( E ) = 0 . So µ ≪ λ . ■

C.4 Proof of Proposition 4

Proposition 4. The gradient of the log of the Voronoi measure p V is given by

where the first two blocks in the case statement apply if there exists some m such that x ∈ int( C m ) or x ∈ ∂C m . See App. C.4 for proof.

Proof. We have three cases. Suppose x ∈ int( C m ) for some m ∈ [ M ] . Then, by Observation 1, we have that p V ( x ) , and, thus, log p V ( x ) is differentiable. Direct computation reveals:

Next, suppose x / ∈ ⋃ M m =1 C m . Then, the measure is zero, so the gradient is as well. Finally, we have that x ∈ ∂C m for some m . Then, by the next Observation (Observation 1), we have that p V is discontinuous, so the derivative is not defined. ■

Observation 1. A ( K , µ ) -Voronoi measure p V is differentiable with respect to p m on the set ∪ M m =1 int ( C m ) , and discontinuous on the set K \ ∪ M m =1 int ( C m ) , i.e., the union of the Voronoi cells' boundaries ∪ M m =1 ∂C m .

D Properties of Hamiltonian Dynamics

Preservation of Measure. First of all, the Hamiltonian equations are volume- or measurepreserving. Intuitively, this means that if we move a region in the phase space along the dynamics for an arbitrary amount of time, the volume of the region would stay unchanged. 15 Concretely, we can define a function

as moving every point in the phase space along the Hamiltonian equations for some time t . 16 Then, using g t as a change of variable would leave the underlying probability measure unchanged.

Time Reversibility. Next, the Hamiltonian equations are time reversible, meaning that if the equations can move ( x , r ) to ( x ′ , r ′ ) , then it would also take ( x ′ , -r ′ ) to ( x , -r ) . Due to our choice of distribution of r is always symmetric, these properties simplify the acceptance probability to be only the energy difference, i.e., we can ignore the conditional probabilities in the standard Metropolis-Hastings acceptance probability.

Conservation of Hamiltonian. Finally, we can quickly verify that the Hamiltonian H is conserved by the Hamiltonian dynamics: 17

This shows that the acceptance probability will be exactly 1 if the Hamiltonian equations are integrated exactly. In practice, numerical errors will lead to fluctuations in H . The Metropolis-Hastings acceptance probability ensures detailed balance.

E Details of Voronoi Sampling

E.1 Conservation of Hamiltonian in Refraction-Reflection

Proposition 5. A single step of refraction-reflection conserves the Hamiltonian.

Proof. Suppose an instantaneous refraction-reflection takes ( x , r ) to ( x ′ , r ′ ) . In other words, ( x , r ) is the particle's coordinate and momentum at the boundary prior to the refraction-reflection, and ( x ′ , r ′ ) is the particle's coordinate and momentum after the refraction-reflection, which includes instantaneous changes in its potential energy and momentum. Recall the refraction-reflection equation is

where r = r ⊥ + r ∥ is the normal decomposition of r and r ′ = r ′ ⊥ + r ′ ∥ . Then,

E.2 Proof of Correctness

In this section, we give a proof of correctness of our sampler by establishing its detailed balance. We will first prove several useful properties of the sampler in E.2.1 to E.2.3 and then combine them to prove the detailed balance in App. E.2.4.

E.2.1 Measure Preservation in Leapfrog

First of all, we note that the two leapfrog steps are measure-preserving, since they are composed of a series of shear transformations.

Proposition 6. The leapfrog steps are measure-preserving.

Proof. Recall that the two leapfrog steps used in HMC (algorithm 1) are ( x , r ) → ( r , r -ε 2 ∇ U ( x )) and ( x , r ) → ( x + ε r , r ) .

By the change-of-variable formula, to show a transformation is measure-preserving, it suffices to show that its Jacobian has determinant 1. The leapfrog step ( x , r ) → ( x , r -ε 2 ∇ U ( x )) as a function has the Jacobian

where ∇ 2 U ( x ) denotes the Hessian of U . This Jacobian (28) has determinant 1 and hence the leapfrog transformation ( x , r ) to ( x , r -ε 2 ∇ U ( x )) is measure preserving. Similarly, the other leapfrog step ( x , r ) → ( x + ε r , r ) has Jacobian

■

■

which also has determinant 1.

E.2.2 Measure Preservation in Refraction-Reflection

The computation for showing that refraction and reflection steps are measure preserving is slightly more involved than the leapfrog case. Our proof below largely follows the one found in Mohasel Afshar and Domke [30], where we simplified and corrected some details.

Proposition 7. The refraction-reflection step is measure-preserving.

Proof. We consider a single step of refraction-reflection as a function γ that takes ( x , r ) to ( x ′ , r ′ ) with total step size ε and with only one discontinuity. Note that, without loss of generality, we may only consider the discontinuity surface to be the hyperplane x 1 = 0 since it is equivalent to any tangent hyperplane of the discontinuity (hyper)surface through an affine transformation. Then, a normal vector to the hyperplane x 1 = 0 is e 1 = (1 , 0 , . . . , 0) and

since r ∥ is not affected by refraction-reflection. Therefore,

In this case, and

̸

Eq. (32) and Eq. (33) shows that the absolute determinant of the Jacobian of γ is the same as

We analyze all four quantities in Eq. (34) individually below. Again, without loss of generality, suppose x 1 ≤ 0 and x ′ 1 ≥ 0 . Let ∆ U ( x ) be the function of the signed potential difference as x 1 changes from negative to positive, defined only when x 1 = 0 . Let y = (0 , y 2 , . . . , y d ) be the point where the discontinuity is encountered. In other words, the particle begins at x , moved past y and arrived at x ′ .

Refraction. In the case of refraction, r 2 1 > 2∆ U ( y ) and

Let δ be the step size it takes to reach y . Then

and

Then,

and

Then,

where

Hence, Eq. (43) and Eq. (44) together imply that

So γ is measure-preserving in the case of refraction.

Reflection. In the case of reflection, r 2 ≤ 2∆ U ( y )

As in the case of refraction, let δ be the step size it takes to reach y . Then

We can then directly calculate

Hence,

So γ is measure preserving in the case of reflection as well.

E.2.3 Time Reversibility

Proposition 8. The refraction-reflection step is time-reversible.

Proof. A refraction-reflection step is time-reversible means that, if a single step of refractionreflection procedure takes ( x , r ) to ( x ′ , r ′ ) , then it would also take ( x ′ , -r ′ ) to ( x , -r ) . We can show this by considering each cases:

Reflection: When reflection happens, ∥ r ⊥ ∥ 2 2 ≤ 2∆ U , and hence ∥ -r ′ ⊥ ∥ 2 2 = ∥ -r ⊥ ∥ 2 2 = ∥ r ⊥ ∥ 2 2 ≤ 2∆ U , meaning that the reverse trajectory would be reflected as well;

Refraction: If ∆ U > 0 , then the reverse trajectory is crossing the boundary from the other side, seeing a sudden decrease in potential, and hence would be refracted and regain the magnitude of momentum lost in the forward step. If ∆ U < 0 , then ∥ -r ′ ⊥ ∥ 2 2 > 2∆ U and hence the reverse trajectory would also be refracted, ending in the original momentum with the sign reversed.

■

Proposition 8, combined with time reversibility of leapfrog step as in basic HMC, shows that Voronoi sampling is time reversible. Hence we only need to use the Hamiltonian difference in the MetropolisHastings acceptance probability.

E.2.4 Detailed Balance

Theorem 1. A step of Structured Voronoi Sampling (Algorithm 4) satisfies the detailed balance.

Proof Sketch. From measure preservation (App. E.2.1 and App. E.2.2), the change of variable as introduced by integrating Hamiltonian equations have Jacobian with absolute determinant 1 and hence can be omitted. From time reversibility (App. E.2.3), we only need to use the Hamiltonian difference in the Metropolis-Hastings acceptance probability. And, finally, by using a Metropolis-Hastings accepting step (lines 8 to 13) we ensure the detailed balance. ■

E.3 Related Work

Such an integrator scheme appears to be well-known in the computational physics community, dating back to as early as Jin and Wen [18], and quite possibly even earlier. Integration algorithms based on this idea continue to receive active research [44].

The first usage of such integrator with discrete stepin the context of MCMC appeared in Mohasel Afshar and Domke [30]. 18 Mohasel Afshar and Domke [30] primarily based their motivation on the optical reflection-refraction analogy. We note that, in fact, Hamiltonian mechanics originated from Hamilton's formulation of geometrical optics (or Hamiltonian optics) [22, §VIII.7]. This connection, important to modern physics, should be interpreted with care in our context since momentum isn't a natural quantity in optical rays.

■

Later, Nishimura et al. [33] provided a derivation of this integrator from the first principles of the calculus of variations [1]. Importantly, it should be noted that, due to the potential function being discontinuous, Eq. (16) loses its well-posedness as a system of differential equations, hence the necessity of justification based on non-smooth analysis.

F Algorithms

Algorithm 2 Langevin Dynamics

Input:

x t : current sample, U : potential energy function, ε : step size

Output:

next sample:

x

t

+1

4: return x t +1

Algorithm 3 MUCOLA

Input:

V t : current sample, U : potential energy function, ε : step size

Output:

next sample:

V

t

+1

4: return V t +1

Algorithm 4 Structured Voronoi Sampling

Input: x t : current embeddings, U : potential function, ε : step size Output: next sample: x t +1 1: r ∼ N ( 0 , ε I ) 2: H t ← U ( V m ⋆ ( x t ) ) + K ( r ) 3: r ← r -ε 2 ∇ U ( V m ⋆ ( x t ) ) 4: x t +1 ← FINDDISC ( x t , r ) 5: r ← r -ε 2 ∇ U ( V m ⋆ ( x t +1 ) ) 6: H t +1 ← U ( x t +1 ) + K ( r ) 7: ∆ H ← H t +1 -H t 8: if s ∼ U (0 , 1) < e -∆ H : 9: return x t +1 10: else 11: return x t +1 ← x t

Algorithm 5 REFRACTREFLECT

Input: r t : current momentum, b normal vector of the boundary, ∆ U : difference in potential energy Output: next momentum: r t +1

5: else

8: return r t +1

▷ Metropolis criterion

▷ refract

▷ reflect

F.1 Efficiently Finding Discontinuities

As explained in §6, the key insight in SVS is to adjust the momentum when facing discontinuity in the potential function. Therefore, we first need to find discontinuities efficiently and then reflect/refract on the boundary of the discontinuity. Remember that at each leapfrog step, we move from x t to x t +1 = x t + ε r . To find discontinuities along this trajectory, we divide each leapfrog step into fractions of a step and check for discontinuity. Concretely, we only advance the embedding by a fraction α , i.e. x ′ = x + αε r (line 4 in Algorithm 6). Then, we check for discontinuity by looking at the difference between the potential energies. If we don't observe any changes in the potential function, we take the next fraction of a step. Otherwise, we find the boundary and adjust the momentum (line 10 in Algorithm 6). Note that in SVS, finding the boundary is straightforward, and it is a hyperplane characterized by the normal vector, as the line goes through the two Voronoi centers on each side of the hyperplane.

Algorithm 6 Find Discontinuity

G Dataset Statistics

This dataset is made available under the CC BY-SA 4.0 license. Our use of this dataset is for i) fine-tuning GPT-2, and ii) training classifiers for topic control, which clearly matches the intended use of this dataset mentioned by the creators as end-to-end training and generating text with content selection. A summary of dataset statistics can be found in Table 2.

H Experimental Details

Sampling algorithms. Hyperparameters for each experiment are reported in Table 4. Following prior work [46], in algorithms based on Langevin dynamics, we apply an exponential decay to the step size by decreasing it to 0.05 after 500 steps. In all settings, we take 500 burn-in steps.

:

Food Classifiers. We train 3 classifiers. First, for experiments with FUDGE, we follow the experimental detail as in the original paper [48] and train a 3 -layered BILSTM classifier with 0 . 5 dropout. The hidden dimension is set to 300 , and the embedding layer is trained from scratch. Second, in experiments with MUCOLA, LANGEVIN, and SVS, to make the setup as close as to FUDGE, we train a 3 -layered BILSTM classifier with 0 . 5 dropout. However, this time the BILSTM is trained on top of frozen GPT-2 representations, thus sharing the embedding layer that is necessary for these methods to work. Finally, to evaluate the success of the methods in following the control targets,

Table 3: Inference time for different methods on the controlled generation experiments. All experiments are done on a single A100-40GB GPU.

we finetune a ROBERTA [25] base model. The accuracy of all the classifiers and the number of trained parameters are reported in Table 5. We train all the models on a single gtx_1080_ti GPU with approximately 2 hours of total computational budget.

Inference times. We report inference times based on sec/batch . The number of decoded sentences per batch depends on the GPU memory and the size of the model. As depicted in table 3, using Voronoi measures does not increase the inference time (compare MUCOLA and LANGEVIN). We observe that SVS inference time is longer, because of the extra momentum adjustment steps. However, one can reduce the inference time by increasing α .

I More Experiments on The Toy Model

To better understand the advantages of Voronoi sampling over HMC or MUCOLA, we further look at the JS divergence between the reference distribution and the distribution of samples when increasing the number of iterations. As depicted in Fig. 4, while the divergence decreases with more iterations across all sampling methods, Voronoi sampling converges to the reference distribution with fewer iterations. We then look at the distribution of sampled elements in Fig. 5. We observe that with 100 iterations, MUCOLA undersamples the element with the maximum probability while oversampling other elements. Finally, we extend the toy model from 4 squared cells in R 2 to 2 k hypercube cells in R k (Fig. 6). As the dimensionality increases, the divergence between the samples' distribution and the reference distribution also increases in all sampling methods. Importantly, Voronoi sampling consistently converges faster across different values for k .

0.5

1

1.5

2

0

0.05

0.1

0.15

Sampling Method

MuCoLa

HMC Sampling

Temperature

JS Divergence

(a)

500

iterations

0.5

1

1.5

2

0

0.01

0.02

0.03

0.04

0.05

0.06

0.07

Sampling Method

MuCoLa

HMC Sampling

Temperature

JS Divergence

(b)

1000

iterations

0.5

1

1.5

2

0

0.2

0.4

0.6

0.8

1

1.2

1.4

Sampling Method

MuCoLa

HMC Sampling

Temperature

JS Divergence

(c)

100

iterations

J Controlled Generation Results

0

1

2

3

0

0.2

0.4

0.6

0.8

1

MuCoLa

HMC Sampling

Reference

Item

Probability

(a)

500

iterations

0

1

2

3

0

0.2

0.4

0.6

0.8

1

MuCoLa

HMC Sampling

Reference

Item

Probability

(b)

1000

iterations

0

1

2

3

0

0.2

0.4

0.6

0.8

MuCoLa

HMC Sampling

Reference

Item

Probability

(c)

100

iterations

2

3

4

5

6

7

8

0

0.05

0.1

0.15

0.2

0.25

0.3

0.35

0.4

Sampling Method

MuCoLa

HMC Sampling

k

JS Divergence

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper pushes the frontier for gradient based sampling for auto regressive models. The paper lays clear theoretical issues to apply such techniques and address two issues: 1) the major contribution of the paper consists in proposing to construct voronoi like probability space over the output token embeddings to enable gradient based sampling process 2) the SVS algorithm also contains the novelty to deal with the non-continuity at the voronoi border by proposing a novel sampling algorithm.\n\nThe authors test the proposed algorithm at a toy setting, confirming their theoretical superiority transfers to this toy dataset over gradient sampling baselines. The authors also test their algorithm at generation and conditional generation tasks where the algorithm shows good performance in terms of success, perplexity and diversity.\n\nSTRENGTHS:\nThe paper has made good contributions in making the gradient based sampling more practical to tackle real problems. The theoretical contribution of constructing voronoi distribution based on embeddings to perform HMS sampling is a significant one; furthermore, the paper addresses some practical issues via a novel sampling algorithm. The theoretical layout is also quite clear, well highlighting the solved issues. \n\nThe paper confirms their theoretical findings empirically as well as testing on two concrete NLP problems where the paper shows better than some existing NLP popular sampling techniques FUDGE.\n\nWEAKNESSES:\nMy major problem for the paper is its presentation. The paper proposes a novel algorithm for sampling, however, all the algorithms have to be found in Appendix (minorly, these are not clearly indicated that it is in appendix for example for line 189). Note that the paper itself should be self contained and the appendix is really for interested readers to learn further details. \n\nAnother minor point for presentation is that the empirical part for this paper is relatively short and doesn't contain very detailed analysis. While the theoretical explanations is key for this paper, the similarity between eq(3) and eq(1) or eq(7) and eq(10) makes me think that the theoretical part can be shortened while maintaining the readability. \n\nFinally, sampling is certainly not the only way (and arguably not the SOTA) way to perform conditional generation, RLHF used in InstructGPT or related techniques such as DPG (Khalifa et al. 2021) should be mentioned or in the best case compared to better situate the current work on its empirical aspect.\n\nQUESTIONS:\nIn Table 1, all gradient sampling algorithms have a relatively high perplexity, which is reflected also in the examples in the appendix. Do the authors think it is an inherent limitation of this class of algorithms?\n\nLIMITATIONS:\nI appreciate the authors discuss the broader impact as well as transparently give insightful limitations. \n\nThe text quality issue is largely mitigated in modern large language models such as GPT-4 etc. This might inspire the authors might to rethink ways to improve on the quality aspect.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors have proposed a novel framework for gradient-based sampling from neural autoregressive LMs named Structured Voronoi Sampling. The core idea of the proposed approach is to map LM distribution to the embedding based version of it and then use newly proposed structured voronoi cells to perform sampling based on HMC. Authors carefully described each step along transformations and why it is theoretically sound. They discussed the conneciton of the proposed approach to related work such as COLD decoding and MuCoLa. They have done experiments in both toy task and controlled text generation to show the effectiveness and superiority of their approach.\n\nSTRENGTHS:\n# Originality\n\nthis work proposes a very coherent way of treating discrete neural based LMs using continuous densities. Opposed to prior work, their method uses less approximations and heuristics.\n\n# Significance \n\nResearch problems around decoding strategies including sampling based decoding is essential given the wide spread of large scale language models. This work performs an important step towards better understanding of the gradient-based sampling methods and how to apply it to LMs parameterizing discrete/categorical distributions.\n\nWEAKNESSES:\n# Experiments\n\n* Authors included samples from their methods and related work they have re-implemented, and these samples look pretty bad. Their proposed method repeats the same tokens right after each other making the sample look very unrealistic w.r.t. unknown data distribution e.g. \"It is located in *the the* city centre near The Portland Arms.The Eagle is an Italian restaurant\". This makes samples based results evaluations in the main text to be much less convincing. I tend to believe that the reason behind this is a poorly finetuned model: authors used GPT2 model which is quite out-dated while much better alternatives exist. Given that authors used modern GPUs (A100 40gb), they had all opportunities for choosing a stronger initial model.\n\n* Authors used only 1000 samples to analyze empirical distributions which is quite limited considering the vocabulary size and avg sentence length. Using much more samples and a better trained model could reveal more interesting observations.\n\n# Efficiency\n\nThe proposed approach is very slow even compared to Langevin dynamics. It would help to address this in the main text and not in the Appendix. Also it would help to put usual ancestral sampling time to see the gap between gradient based sampling and ancestral sampling in general.\n\nQUESTIONS:\nGeneral comment: I believe this is very strong theoretical work here which has rather weak experimental setting because of poor choice of the initial model (GPT-2) for fine-tuning and the task setting. I think this could become a much stronger submission upon stronger experiments are done,\n\nQuestions:\n\nsec 7.2 lines 293 - 299. Its unclear what do you mean by reference distribution here and how LM is related to it? IIUC you are saying that ancestral sampled from GPT-2 finetuned on E2E dataset resembles the unknown reference/data distribution? I think that is a very rough approximation. I think ancestral samples from GPT-2 gives unbiased estimates of sequence level distribution induced by GPT-2, thats it. The underlying connection between finetuned GPT-2 and reference distribution is unknown but likely not that strong given the samples you shown in the appendix.\n\nTable 1: in the text you claim SVS outperforms everything else given metrics, but is it? It goes very close to Langevin per success and PPL (~ same given std), and diversity is lower than FUDGE / MuCoLa. Moreover, I think diversity should be compared to values computed over the data distribution, could you report that as well?\n\nLIMITATIONS:\nAuthors included broader impact section and discussed how their framework could help to alleviate negative implications of large LMs as well as being a generator of intentionally toxic content.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
The authors present Structured Voronoi Sampling, which is a gradient-based sampling approach. To be specific, the authors map the discrete distribution by a language model and defines densities; the density is then used to sample, which the process is based on hamiltonian monte carlo. The novelty of this paper comes from the theoretical analysis, but the core weakness is the empirical result; not much performance gain is seen across the experiments.

STRENGTHS:
- The paper is well-written and structured. The flow of the paper is easy to follow, and the authors explain the introduced concept step-by-step.

- The paper proposes a novel gradient-based sampling method that caters for controlled generation tasks.

WEAKNESSES:
- Weak/wrong claims: i.e. Line 65-66. Not all language models share input and output embeddings. Authors should rephrase the sentence to avoid possible misunderstanding

- The empirical result is weak. There is not much of difference compared to Langevin in success and ppl score in Table 1. The same applies in Figure 3.

QUESTIONS:
Q1. Is there any particular reason why the authors did not test on popular controllable text generation task?

LIMITATIONS:
Please refer to summary and weakness section
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new gradient-based sampling approach called Structured Voronoi Sampling (SVS) for controlled text generation. The key idea is to extend the discrete point distribution over word embeddings given by language models into a continuous density that spreads out probability over their corresponding Voronoi cells. A Hamiltonian Monte Carlo scheme is then devised to efficiently sample from that density, handling discontinuities between cells through a volume-preserving refraction/reflection trick. Empirical results on a toy problem and a more realistic controlled text generation problem show SVS can better match target distributions and control constraints compared to baselines.\n\nSTRENGTHS:\nDeveloping provably sound sampling methods for text generation remains an important open problem for language models, especially for controlled generation. The proposed  method is well-motivated from first principles and provides formal guarantees unavailable with prior heuristic approaches, with rigorous mathematical derivations provided in appendices.\n\nA number of innovative steps are taken to get the core Voronoi idea to work, including (i) lifting the discrete token distribution to the embedding space, (ii) smoothing the measure-zero discrete distribution to a continuous density with the Voronoi transformation, and (iii) handling discontinuities in sampling with the refraction/reflection trick. These moves are independently interesting in their own right and a large segment of the NeurIPS community will likely find at least one of them to be novel and potentially applicable in future work.\n\nRelated methods like MuCoLa and Fudge are discussed and compared in the experiments. The toy domain is pedagogically useful for exposing the limitations of existing approaches, while the scaled-up controlled generation results are promising, showing competitive fidelity and control.\n\nThe paper is well-written and structured overall, and the writing systematically builds up the relevant concepts for a typical NeurIPS reader to understand the significance. The key generalizable ideas are appropriately emphasized.\n\nWEAKNESSES:\n(1) The most significant issue limiting the applicability of SVS as currently posed is the problem of calculating the base measure $\\mu$ in a high-dimensional space. This problem in some sense puts us back where we started, as the original problem motivating SVS in the first place is that the integral in the normalizing constant of the energy function is not tractable. Computing the exact integral for the base measure in SVS is obviously also not tractable (as noted up-front as a limitation the paper, which I appreciated). The assumption in 6.1 that all cells have equal base measure is a very strong assumption, given that this clearly does not hold in practice, and somewhat undermines the very careful and rigorous progression of proofs of correctness.\n\nI see a couple way to strengthen this part of the paper, and would strongly recommend doing something to close this gap for the paper to be maximally impactful. One possibility is to provide some analysis of the consequences of violating this ‘equal base measures’ assumption: is there a bound on how bad violations can get? how badly is this assumption violated in the controlled generation task being reported? Intuitively, ignoring real diff them would result in over- or under-representing certain regions. A second possibility would be to take a first step toward some approximate method to account for base measure differences. For example, just taking distances to the first k nearest neighbors of a point could efficiently put a bound on how big the cell could be? Or some kind of Monte Carlo approximation? Clearly a full resolution of this problem is a task for another paper, and I do not expect this problem to be fully solved, but charting a course (even if that course is a bit inefficient) would go a long way.\n\n(2) More analysis could be provided on the toy example with known reference distribution; for example, what do the sampled distributions actually look like? In what way is MuCoLa off when temperature is low? How does this scale with the size of the space (e.g. if instead of a 2x2 square, it were a 2^k hypercube in k dimensions)? Does MuCoLa just take longer to burn in? or gets stuck sampling {0,1,2}? I would have liked to understand these failure modes a bit better to better motivate how the proposed method overcomes them.\n\n(3) The controlled generation task could also be strengthened: the results are from a single fairly bespoke 'food' task and classifier, which limits generalizability. It would be more compelling to evaluate on a more systematic suite of controlled language generation benchmarks to better pinpoint where the benefits of Voronoi sampling are most pronounced. Additionally, since controlled generation is highly dependent on the (trained) classifier being used to guide sampling, it would be helpful to test lower or higher capacity classifiers with different uncertainty levels. I worry that using the same classifier for both generation and evaluation (the ‘success’ column of Table 1) means we may just be measuring the sampler’s ability to overfit to a bad classifier (it is quite a bit more ‘faithful’ to this classifier than MuCoLa or Fudge, which is not necessarily a good thing if it is just ‘hacking’ a bad classifier).\n\nQUESTIONS:\n1. More exposition could be provided to unpack eq. (2) which may be confusing to a typical neurips reader. in particular, I think many practitioners are more familiar writing the final layer as a softmax applied to a linear transformation, i.e. something like $P(w_n | w_{<n}) = softmax(v * f(w_{<n}))$ where $v$ is the full embedding matrix). Some readers may also have forgotten that GPT2 used the same embedding matrix at the initial and final layers (other implementation use a separate fully-connected layer here), so it may be worth a gentle reminder to avoid confusion. \n\n2. Similarly, some additional clarification could be given to interpret the denominator in Def. 2 and 3 as is essentially normalizing by the ‘volume’ of the cell (with the measure giving a suitable notion of volume). This will be obvious to those versed in measure theory but could help other ML researchers with less mathematical background. \n\n3. I was trying to think through the precise relationship to other approaches based on nearest-neighbor smoothing (e.g. Khandelwal [et](http://U.et) al, 2019. “*Generalization through memorization: Nearest neighbor language models*.” or Khandelwal et al, 2020. “*Nearest neighbor machine translation”;* El-Kishky et al, 2023. “*kNN-Embed: Locally Smoothed Embedding Mixtures for Multi-interest Candidate Retrieval”*). These approaches have a very similar flavor to the voronoi embedding, which on the surface appears to be equivalent to some kind of importance sampling based on nearest-neighbors. Does the voronoi sampler reduce to something like this as a special case? \n\n4. very minor: It was a bit confusing to have Alg. 1 prominently referred to in the main text (down to giving line numbers) but not be able to find it until I was later going through the supplemental. if it is important, it may be worth trying to fit into the main text.\n\n5. very minor: the abstract uses the acronym SVS but it doesn’t reappear until Fig. 3, at which point I was very confused (pragmatically, it’s confusing that it’s referred to as ‘voronoi sampling’ in Fig. 2, right next to it, which implies that SVS is something different?)\n\n6. very minor: is it worth pointing out in section 7.2 that there is no reason to actually prefer voronoi over regular ancestral sampling, given that ancestral sampling is already very efficient and provably matches the target distribution? And that we’re just doing this exercise (Fig. 3) as validation that the voronoi sampler is pretty close to the ‘gold standard’?\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-17"}

comments: SUMMARY:
This paper proposes Structured Voronoi Sampling: a new gradient-based algorithm to sample from a distribution (i.e. a language model). The proposed approach leads to comparably fluent text whilst being able to better follow constraints (e.g. a topic) for the desired generation.

STRENGTHS:
1. a new sampling method that is a correct MCMC scheme (unlike previous gradient-based samplers)
2. evaluation on a synthetic task shows that the proposed approach better models the tail of the true distribution
3. evaluation on constrained language modeling shows that the proposed approach achieves good fluency and diversity, whist being able to follow the control target

WEAKNESSES:
1. while the writing is generally great, it would help to have an introduction of gradient-based sampling before Section 3
2. it is not clear how the proposed approximation of the (costly!) base measure affects the correctness of the proposed method

QUESTIONS:
1. some of the key algorithms and results are reported in the appendix. It would be good to include them in the main text upon acceptance
2. in eq (2), enc(context) is a single vector. This is not the case with current models (e.g. GPT-2) which have a vector per input token. Can you comment on how this affects the formulation of LM with Embeddings?
3. in eq (3), the conditioning on V_{<n} seems wrong, as `n` is only defined in the Cartesian product

LIMITATIONS:
yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nGradient-based sampling for text generation is an important challenge, as it allows for sampling from energy-based models, such as one defined by a mixture of experts as found in classifier-guided sampling. The main challenges in gradient-based sampling for discrete distributions are encoding the discrete distribution into $R^d$ and dealing with inevitable discontinuities that arise in the encoding. The paper proposes a method that addresses these two challenges with Voronoi measures and an application of refract+reflect HMC respectively.\n\nThe method is validated in three settings: First, sampling from a non-structured and tractable discrete distribution with 4 classes and associated embeddings. Second, sampling from a language model. Third, sampling from a language model with additional constraints. In all settings, the method improves upon baselines, namely Fudge, MuCoLa, and Langevin dynamics (without reflection).\n\nSTRENGTHS:\n* The approach is well-motivated and interesting. \n* The writing is clear and easy to follow.\n* In text generation, unconstrained and constrained, the method does show improvements over the MuCoLa baseline. However, the gains from reflection and refraction seem quite small.\n\nWEAKNESSES:\nExperiment baselines: The paper only compares to MuCoLa and Fudge as text generation baselines, but other methods for controllable text-generation exist such as diffusion-LMs and other Gibbs with gradient (GwG) methods [2]. Since the claim is principled gradient-based methods for text-generation, not having a comparison to diffusion-LMs seems reasonable. However, I believe other GwG methods should be compared (see questions).\n\nQUESTIONS:\n1. How is the method related to Gibbs with gradients methods, such as [1] and [2]? SVS takes advantage of embedding geometry and I believe GwG does not, which would plausibly lead to improvements.\n2. It would have been nice to see comparisons to non-principled gradient-based sampling, such as COLD. A lack of rigorous justification for an existing method should not be enough to discount its (potential) effectiveness.\n3. Why is the constrained text generation task different from the tasks studied in MuCoLa?\n4. There are quite a few references to lines of algorithms in the Appendix.\n5. Are the small gains from reflection due to the length of the chains? Would the difference between Langevin and SVS be larger if the computational budget was smaller?\n\n[1] Grathwohl, W., Swersky, K., Hashemi, M., Duvenaud, D.K., & Maddison, C.J. Oops I Took A Gradient: Scalable Sampling for Discrete Distributions. ICML 2021.\n\n[2] Zhang, R., Liu, X., & Liu, Q. (2022). A Langevin-like Sampler for Discrete Distributions. ICML 2022.\n\nLIMITATIONS:\nThe limitations were adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}