Entropy-Aligned Decoding of LMs for Better Writing and Reasoning

TL;DR

EPIC addresses decoding quality gaps by directly aligning the LM's sampling entropy with the irreducible data entropy. It defines an entropy-tilted next-token distribution via $q_{t,\alpha}(y_t|\boldsymbol{y}_{<t}) \propto q_t(y_t|\boldsymbol{y}_{<t}) \exp(-\alpha H_{t:t+k}(y_t))$ and proves that an optimal $\alpha^*$ minimizes $H(p,q_\alpha)$, achieving entropy calibration without sacrificing accuracy. The method couples Entropy-Aware Lazy Gumbel-Max with tight entropy bounds and a Rao–Blackwellized entropy estimator to enable exact, sublinear sampling from the entropy-aligned distribution. Empirically, EPIC yields more diverse yet coherent generations across creative writing, abstractive summarization, and mathematical reasoning, improving LM-as-judge preferences and standard metrics over strong baselines like top-$k$, top-$p$, min-$p$, and typical decoding.

Abstract

Language models (LMs) are trained on billions of tokens in an attempt to recover the true language distribution. Still, vanilla random sampling from LMs yields low quality generations. Decoding algorithms attempt to restrict the LM distribution to a set of high-probability continuations, but rely on greedy heuristics that introduce myopic distortions, yielding sentences that are homogeneous, repetitive and incoherent. In this paper, we introduce EPIC, a hyperparameter-free decoding approach that incorporates the entropy of future trajectories into LM decoding. EPIC explicitly regulates the amount of uncertainty expressed at every step of generation, aligning the sampling distribution's entropy to the aleatoric (data) uncertainty. Through Entropy-Aware Lazy Gumbel-Max sampling, EPIC manages to be exact, while also being efficient, requiring only a sublinear number of entropy evaluations per step. Unlike current baselines, EPIC yields sampling distributions that are empirically well-aligned with the entropy of the underlying data distribution. Across creative writing and summarization tasks, EPIC consistently improves LM-as-judge preference win-rates over widely used decoding strategies. These preference gains are complemented by automatic metrics, showing that EPIC produces more diverse generations and more faithful summaries. We also evaluate EPIC on mathematical reasoning, where it outperforms all baselines.

Figures (4)

• Figure 1: Entropy trajectories during text generation conditioned on "The cat" prompt. Figure (a) shows a reference sample with entropy values staying in the "typical" band (gray region). Figure (b) illustrates degeneration into repetition, where entropy collapses below the typical range. Figure (c) shows gibberish, where entropy rises well above the typical range. Figure (d) demonstrates Epic decoding, which maintains entropy close to the typical band and yields coherent text.
• Figure 2: Interaction of Epic and truncation. Points denote candidate tokens, represented by probability ($\rightarrow$) and lookahead entropy ($\uparrow$). Gray crosses denote truncated tokens. Tokens considered by truncation algorithms are denoted using blue circles, while Epic-candidate tokens are denoted by violet pluses. The vertical dashed line denotes the probability cutoff, and the horizontal dashed line denotes the entropy cutoff. Epic favors tokens that balance both probability and future uncertainty, avoiding degenerate low-entropy repetitions or incoherent high-entropy expansions.
• Figure 3: Tighter Entropy Bounds. Making use of \ref{['lemma:miscalibration']} in tandem with the empirical model error lends to much tighter entropy bounds than theoretical.
• Figure 4: Calibration of Decoding Methods to Model Cross-Entropy. We plot the mean conditional entropy of generated sequences as a function of the reference cross-entropy, along with the resulting calibration error for each decoding strategy. Epic closely tracks the target cross-entropy across the full range, yielding substantially lower calibration error and a markedly smoother trajectory than competing methods. In contrast, standard heuristics such as top-$p$, top-$k$, and min-$p$ exhibit systematic bias and high variance, leading to persistent miscalibration despite aggressive tuning.

Theorems & Definitions (9)

• Lemma 3.0
• Lemma 3.0
• Lemma 4.0
• Lemma A.0
• proof
• Lemma A.0
• proof
• Lemma A.0
• proof