On Next-Token Prediction in LLMs: How End Goals Determine the Consistency of Decoding Algorithms

TL;DR

This work addresses how decoding choices after next-token prediction interact with different end-goals in large language models. By formalizing an asymptotic consistency framework where the predictor converges to the true distribution $p^*$, it analyzes $K_T$-lookahead, random sampling, and temperature-scaled sampling with respect to two losses: the $N$-gram Hamming loss (information retrieval) and the sequence-level cross-entropy (creative generation). Key findings include: random sampling is consistent for cross-entropy under convergence to $p^*$, while no universal polynomial-time decoder is optimal for the $N$-gram Hamming loss; $K_T$-lookahead is optimal only under a specific condition on $p^*$; and stochastic decoders generally fail to be consistent for the Hamming loss but can be consistent for cross-entropy only in particular parameter regimes. The results reveal a dichotomy between information retrieval and creative generation, underscoring that the choice of decoding strategy should be tailored to the intended goal, and motivate adaptive, goal-aware decoding in practice.

Abstract

Probabilistic next-token prediction trained using cross-entropy loss is the basis of most large language models. Given a sequence of previous values, next-token prediction assigns a probability to each possible next value in the vocabulary. There are many ways to use next-token prediction to output token sequences. This paper examines a few of these algorithms (greedy, lookahead, random sampling, and temperature-scaled random sampling) and studies their consistency with respect to various goals encoded as loss functions. Although consistency of surrogate losses with respect to a target loss function is a well researched topic, we are the first to study it in the context of LLMs (to the best of our knowledge). We find that, so long as next-token prediction converges to its true probability distribution, random sampling is consistent with outputting sequences that mimic sampling from the true probability distribution. For the other goals, such as minimizing the 0-1 loss on the entire sequence, we show no polynomial-time algorithm is optimal for all probability distributions and all decoding algorithms studied are only optimal for a subset of probability distributions. When analyzing these results, we see that there is a dichotomy created between the goals of information retrieval and creative generation for the decoding algorithms. This shows that choosing the correct decoding algorithm based on the desired goal is extremely important and many of the ones used are lacking theoretical grounding in numerous scenarios.

Figures (6)

• Figure 1: A plot of the amount of trials $K_1$-lookahead was optimal for the 1-gram Hamming loss (the Hamming loss). Each point represents the average optimality over $200$ randomly generated Markov chains with a set amount of nodes and Dirichlet parameter $\alpha$. Smaller $\alpha$s create more "peaky" distributions and thus have higher KL-divergence from the uniform distribution, while larger $\alpha$s create more uniform distributions. There were 8 nodes in each Markov chain for this figure and the sequence length goes up by two as one moves right in the plots.
• Figure 2: A plot of the amount of trials $K_1$-lookahead was optimal for the $L$-gram Hamming loss (the $0-1$ loss). The same setup as Figure \ref{['fig:klook_1gh']} otherwise.
• Figure 3: A plot of all trials of $K_1$-lookahead decoding being optimal for $N$-gram Hamming on a Markov chain with "num nodes" amount of nodes. The amount of nodes increases as one looks down the rows and the sequence length increases as one goes down the columns. Each line represents a specific $K_1$-lookahead being compared to the optimal $N$-gram Hamming. The $x$ and $y$ axes denote the same that has been used in Figures \ref{['fig:klook_1gh']} and \ref{['fig:klook_Lgh']}.
• Figure 4: A plot of the amount of trials $K_K$-lookahead was optimal for the 1-gram Hamming loss (the Hamming loss). There were 8 nodes in each Markov chain and the sequence length goes up by two as one moves right in the plots.
• Figure 5: A plot of the amount of trials $K_K$-lookahead was optimal for the $L$-gram Hamming loss (the $0-1$ loss). The same setup as Figure \ref{['fig:kklook_1gh']} otherwise.

Theorems & Definitions (28)

• Lemma 4.1
• Proposition 1
• Theorem 4.2
• Corollary 4.2.1
• proof
• Lemma 4.3
• Theorem 4.4
• Corollary 4.4.1
• proof
• Proposition 2