Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Decoding-Free Sampling Strategies for LLM Marginalization

David Pohl, Marco Cognetta, Junyoung Lee, Naoaki Okazaki

TL;DR

The paper tackles marginalizing over all tokenizations of a surface form in subword-based LLMs, noting that canonical tokenization is an imperfect proxy and exact marginalization is NP-hard. It introduces a decoding-free marginal estimation method based on a tokenization lattice that encodes all tokenizations via automata/transducers, avoiding LLM generation. By sampling tokenizations from the lattice and scoring them with the LM, the method yields fast, tokenizer-agnostic marginal estimates and provides provable lower bounds that can outperform proxy importance sampling in practice. Empirical results on Q&A and translation tasks show substantial runtime speedups (up to >30x) with competitive marginal accuracy, while revealing that proxy importance sampling often underestimates the true marginal.

Abstract

Modern language models operate on subword-tokenized text in order to make a trade-off between model size, inference speed, and vocabulary coverage. A side effect of this is that, during inference, models are evaluated by measuring the probability of only the specific tokenization produced as the output, despite there being many possible ways to represent the same text with a subword vocabulary. Recent studies have argued instead for evaluating LLMs by marginalization - the probability mass of all tokenizations of a given text. Marginalization is difficult due to the number of possible tokenizations of a text, so often approximate marginalization is done via sampling. However, a downside of sampling is that an expensive generation step must be performed by the LLM for each sample, which limits the number of samples that can be acquired given a runtime budget, and therefore also the accuracy of the approximation. Since computing the probability of a sequence given the tokenization is relatively cheap compared to actually generating it, we investigate sampling strategies that are decoding-free - they require no generation from the LLM, instead relying entirely on extremely cheap sampling strategies that are model and tokenizer agnostic. We investigate the approximation quality and speed of decoding-free sampling strategies for a number of open models to find that they provide sufficiently accurate marginal estimates at a small fraction of the runtime cost and demonstrate its use on a set of downstream inference tasks.

Decoding-Free Sampling Strategies for LLM Marginalization

TL;DR

The paper tackles marginalizing over all tokenizations of a surface form in subword-based LLMs, noting that canonical tokenization is an imperfect proxy and exact marginalization is NP-hard. It introduces a decoding-free marginal estimation method based on a tokenization lattice that encodes all tokenizations via automata/transducers, avoiding LLM generation. By sampling tokenizations from the lattice and scoring them with the LM, the method yields fast, tokenizer-agnostic marginal estimates and provides provable lower bounds that can outperform proxy importance sampling in practice. Empirical results on Q&A and translation tasks show substantial runtime speedups (up to >30x) with competitive marginal accuracy, while revealing that proxy importance sampling often underestimates the true marginal.

Abstract

Modern language models operate on subword-tokenized text in order to make a trade-off between model size, inference speed, and vocabulary coverage. A side effect of this is that, during inference, models are evaluated by measuring the probability of only the specific tokenization produced as the output, despite there being many possible ways to represent the same text with a subword vocabulary. Recent studies have argued instead for evaluating LLMs by marginalization - the probability mass of all tokenizations of a given text. Marginalization is difficult due to the number of possible tokenizations of a text, so often approximate marginalization is done via sampling. However, a downside of sampling is that an expensive generation step must be performed by the LLM for each sample, which limits the number of samples that can be acquired given a runtime budget, and therefore also the accuracy of the approximation. Since computing the probability of a sequence given the tokenization is relatively cheap compared to actually generating it, we investigate sampling strategies that are decoding-free - they require no generation from the LLM, instead relying entirely on extremely cheap sampling strategies that are model and tokenizer agnostic. We investigate the approximation quality and speed of decoding-free sampling strategies for a number of open models to find that they provide sufficiently accurate marginal estimates at a small fraction of the runtime cost and demonstrate its use on a set of downstream inference tasks.
Paper Structure (23 sections, 7 equations, 3 figures, 4 tables, 1 algorithm)

This paper contains 23 sections, 7 equations, 3 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background and Notation
  4. Subword Tokenization
  5. Building a Tokenization Lattice
  6. Sampling from LLMs
  7. Constrained Sampling
  8. Marginalization
  9. Full and Partial Enumeration
  10. Full and Partial Enumeration
  11. The Distribution of Probability Mass
  12. A Lot of Mass is "Near" the Canonical Sequence
  13. Downsides of a Constrained LLM as a Proxy
  14. Generating Tokenizations Efficiently
  15. Uniformly Sampling Tokenizations
  16. ...and 8 more sections

Figures (3)

  • Figure 1: Three methods for approximating the marginal of an input text. The typical approach is to compute the probability of the canonical tokenization, which is fast but inaccurate. Importance sampling allows for more accurate marginalization, but sampling from an LLM is very slow. Our lattice approach is fast and produces good estimates.
  • Figure 2: An example subword tokenization lattice for the input $\mathbf{s} = \texttt{sample}$ using the Llama2 vocabulary.
  • Figure 3: The runtimes of generating a sequence from an LLM and scoring a sequence with an LLM for Llama and Gemma. Note that, as the sequence length grows, generation becomes significantly slower than scoring.