Contextual Temperature for Language Modeling

TL;DR

Contextual Temperature introduces per-token adaptive temperature τ learned from context to control Softmax softness in language modeling, addressing limitations of fixed or globally scheduled scaling. The CT-MoS framework extends Mixture of Softmaxes by applying token- and context-aware temperatures, enabling fine-grained uncertainty management across sequences. Experiments on Penn Treebank and WikiText-2 show significant perplexity improvements over strong baselines, and analyses reveal distinct temperature patterns by token and position. This work establishes a general mechanism for learned temperature control with potential applications to other discrete-output NLP tasks requiring long-term uncertainty management.

Abstract

Temperature scaling has been widely used as an effective approach to control the smoothness of a distribution, which helps the model performance in various tasks. Current practices to apply temperature scaling assume either a fixed, or a manually-crafted dynamically changing schedule. However, our studies indicate that the individual optimal trajectory for each class can change with the context. To this end, we propose contextual temperature, a generalized approach that learns an optimal temperature trajectory for each vocabulary over the context. Experimental results confirm that the proposed method significantly improves state-of-the-art language models, achieving a perplexity of 55.31 and 62.89 on the test set of Penn Treebank and WikiText-2, respectively. In-depth analyses show that the behaviour of the learned temperature schedules varies dramatically by vocabulary, and that the optimal schedules help in controlling the uncertainties. These evidences further justify the need for the proposed method and its advantages over fixed temperature schedules.

Figures (4)

• Figure 1: (a) Temperature of each word learned using the proposed method over training epochs. Each line in the figure represents a distinct token ranked by frequency. The y-axis shows the tokens' temperature values. Note that the actual value equals to the value on the y-axis plus 2; we eliminate the integer part (plus 2) for a better visualization. (b) Means and confidence intervals of the temperature vector (y-axis) over positions in a sentence (x-axis). As in this figure, the average temperature is high at the beginning of a sentence and gradually decreases towards latter positions in a sentence.
• Figure 2: The architecture of the proposed CT-MoS model. Black components are those the same as the MoS model, while the blue ones are the newly added ones in the proposed approach.
• Figure 3: Gradients of loss with respect to (a) logit $z_0$ and (b) logit $z_1$. In each figure, the x-axis is the probability $p_i$ of class $i$, y-axis is the temperature value $\tau_i$ of class $i$, and z-axis is the gradient $\frac{\partial L}{\partial z_i}$. The colorful mesh represents the gradients when the contextual temperature is applied, while the red mesh represents the case without temperature.
• Figure 4: Gradients of loss with respect to $z_{\tau_i}$. The x-axis is the probability $p_i$ of class $i$, y-axis is the temperature value $\tau_i$ of class $i$, and z-axis is the gradient $\frac{\partial L}{\partial z_{\tau_i}}$.