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 .