Temperature Scaling in Discrete Sequence (Language) Models

Temperature scaling is widely used to control generation behavior and improve reasoning performance in discrete sequence models. However, temperature scaling at the sequence level is particularly challenging since the partition function of the model is typically not tractable and naive token-wise temperature scaling induces bias in the maximum a posteriori (MAP) estimates. Discrete diffusion settings exacerbate this issue: enhanced flexibility due to multiple potential orderings for generating tokens complicates likelihood computation and incurs high variance. We propose and investigate fine-tuning objectives for sequence-level temperature scaling, along with metrics to quantify the success of a given temperature scaling procedure. Specifically, our contributions are threefold. First, we introduce our objective across a wide class of sequence models through a novel formalism that unifies autoregressive and diffusion-based language models. Second, we design two metrics that quantify temperature scaling based on likelihood ratios and thus obviate access to the partition function. Finally, we introduce fine-tuning objectives that reliably achieve the desired change in model temperature. Our experiments with language models (GPT2, BD3LM) show that the proposed approach leads to more consistent generation with lower perplexity. Furthermore, we provide empirical evidence that it can enhance the reasoning performance in language models. Code is available at github.com/Aalto-QuML/temp_lm.

Temperature Scaling is a method to change how a language model like GPT responds. With low temperatures, good answers should appear more often; with high temperatures, the diversity of answers is greater. This, for example, can improve the model's ability to solve maths problems. However, the commonly used method to achieve this (myopic temperature scaling) is only mathematically sound when producing one output token, not long answers. This can lead to poor results when used for long answers. In this paper, we define metrics that can measure how much the temperature of the model is changed for long answers, and provide training procedures that achieve this desired behaviour. We demonstrate that our method effectively scales the temperature and, unlike myopic temperature scaling, does not drastically degrade model performance.