Structured Voronoi Sampling

Summary

The 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. The 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.

Strengths

The 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. The 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.

Weaknesses

My 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. Another 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. Finally, 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.

Questions

In 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?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

I appreciate the authors discuss the broader impact as well as transparently give insightful limitations. The 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.

Summary

Authors 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.

Strengths

# Originality this 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. # Significance Research 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.

Weaknesses

# Experiments * 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. * 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. # Efficiency The 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.

Questions

General 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, Questions: sec 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. Table 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Authors 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.

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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Please refer to summary and weakness section

Summary

The 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.

Strengths

Developing 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. A 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. Related 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. The 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.

Weaknesses

(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. I 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. (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. (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).

Questions

1. 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. 2. 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. 3. 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? 4. 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. 5. 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?) 6. 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’?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

See above.

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

yes.

Summary

Gradient-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. The 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).

Strengths

* The approach is well-motivated and interesting. * The writing is clear and easy to follow. * 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.

Weaknesses

Experiment 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).

Questions

1. 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. 2. 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. 3. Why is the constrained text generation task different from the tasks studied in MuCoLa? 4. There are quite a few references to lines of algorithms in the Appendix. 5. 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? [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. [2] Zhang, R., Liu, X., & Liu, Q. (2022). A Langevin-like Sampler for Discrete Distributions. ICML 2022.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations were adequately addressed.

Thank you for the detailed answers

My comments have been addressed in the response, I hope that the presentation can be further improved when the paper is accepted to make the reading easier and more accessible by other readers.

Thanks!

I appreciate the thoughtful response and the improvements described, which strengthen an already very-strong paper.

thanks!

thanks for response! I am satisfied with the provided answers and extra experiments and going to increase my score!

Response to Authors

The rebuttal reads well. There seems to be a misunderstanding by me in Langevin in the paper. Thank you for the rebuttal, and I will change my score accordingly.

Feedback on rebuttal

Thank you for clarifying my concerns and engaging with the comments raised by the other reviewers!