Temperature Scaling in Discrete Sequence (Language) Models
Abstract
Temperature scaling (TS) is widely used to control generation behavior in discrete sequence models. However, TS at the sequence level is challenging as the needed partition function is typically intractable and naive token-wise TS induces bias. 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 TS, along with metrics to quantify the success of a given TS procedure. Specifically, our contributions are threefold. First, we introduce our objective across a wide class of sequence models through a unifying formalism. Second, we design two metrics that quantify TS 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 Negative Log-Likelihood (NLL). 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.