Entropy-Aligned Decoding of LMs for Better Writing and Reasoning

TL;DR

EPIC is a novel entropy-aware decoding method for language models that improves the quality, diversity, and faithfulness of generated text by explicitly regulating uncertainty during sampling.

Contribution

EPIC introduces a hyperparameter-free, entropy-aligned decoding approach that manages uncertainty effectively, outperforming existing methods in various language generation tasks.

Findings

EPIC yields higher preference win-rates in creative writing and summarization.

EPIC produces more diverse and faithful generations according to automatic metrics.

EPIC outperforms baselines in mathematical reasoning tasks.

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.