Contextual Temperature for Language Modeling

1 INTRODUCTION . Due to the discrete nature of human language , the Softmax layer is mandatory to convert learned representations into a sequence of linguistic vocabulary . Temperature scaling is often used along with the Softmax layer in natural language processing . Generally , the temperature is applied as a denominator to output logits of the Softmax layer ( Krizhevsky et al. , 2012 ; Bahdanau et al. , 2014 ; Hu et al. , 2017 ; Caccia et al. , 2018 ) . The scale of temperature controls the smoothness of the output distribution . That is , given a temperature τ , when τ → ∞ , the distribution becomes more uniform , thus increasing the uncertainty . Contrarily , when τ → 0 , the distribution collapses to a point mass . Although temperature scaling has been justified to achieve great success , existing implementations are limited . For instance , the temperature is assumed to be constant throughout training ( Norouzi et al. , 2016 ; Ma et al. , 2017 ; Chen et al. , 2019 ) , or to be a fixed schedule ( Hu et al. , 2017 ) . Most importantly , none of the existing works studies the effects on different word tokens when the temperature changes . In reality , however , the temperature of each token can be dramatically different . Figure 1 ( a ) shows the optimized temperature for each word token during the course of training . As shown in the figure , some certain words have a distinct heating-up temperature scaling , while the majority of words have a scaling that gradually cooling down the temperature . We argue that existing methods limit themselves to some fixed schedules , and thus have great difficulty to generalize . In addition , another example can be observed in Figure 1 ( b ) , which indicates that the average temperature drops , as the length of the context increases . This suggests that the temperature mechanism helps promote stochasticity in the beginning of a sentence , then gradually suppressing the uncertainty until the end . All of these suggest the use of a more generalized temperature mechanism with the advantage of being able to deal with these phenomena . We propose contextual temperature , a generalized temperature scaling method which takes both the history of a sequence and a particular vocabulary into consideration . Through optimizing the use of temperature scaling by the change of contexts , contextual temperature has the ability to generate a unique temperature for each vocabulary . As we parameterize contextual temperature using a deep neural network , its parameters co-evolve with the rest of model parameters , making the temperature adapt to the training procedure . Experiment results on language modeling demonstrate significant improvements on both Penn Treebank and WikiText-2 datasets , with and without the fine-tune setting . Consistent improvements are shown on both validation and testing splits . In addition , we conduct comprehensive analyses and ablation studies to confirm the improvements of contextual temperature . We observe that our method is capable of controlling the uncertainties as the patterns of contexts change , allowing language models to achieve much better performances . To the best of our knowledge , this is the first systematic work that studies the role of per-token temperature changing over context for training language models . The experimental results suggest a new way of training sequential models with discrete outputs , that is , using parameterized temperatures . We have also established the link between the control of model uncertainty and the use of temperatures , paying the way for extensions on tasks that require such long-term control , for instance summarization , translation , and dialogue generation . 2 RELATED WORK . 2.1 LANGUAGE MODELING . Given a sequence of history tokens x1 .. t−1 , language modeling aims at predicting the next token xt in a sentence . In other words , the goal is to model the probability of choosing the kth word , P ( xt = k|x1 .. t−1 ) , by using an encoding function f and a Softmax layer σ. f is parameterized by θ to compress the history z = f ( x1 .. t−1 ; θ ) T ·W . P ( xt = k|x1 .. t−1 ; θ ) = σ ( z ) k ( 1 ) Here W ∈ Rd×K is a matrix that converts a d dimensional output of f into a K dimensional vector . The state-of-the-art language models adopt deep neural networks as the encoding function , including the earlier feed-forward networks ( Bengio et al. , 2003 ) , more recent recurrent-based ones ( Mikolov et al. , 2010 ; Hochreiter & Schmidhuber , 1997 ) and attention-based models ( Devlin et al. , 2018 ) . 2.2 SOFTMAX LAYER . A Softmax layer σ ( z ) normalizes a K dimension , real-valued vector z to make it sum to 1. σ ( z ) i = zi∑K j e zj ( 2 ) Recent progress in language modeling suggests that Mixture of Softmaxes ( MoS ) ( Yang et al. , 2018 ) would significantly improve the performance by computing multiple Softmax distributions and use a set of weights to sum them up as the final probability distribution . To achieve this , a set of M matrices Wm is applied to the output generated by f . That is , zm = f ( x1 .. t−1 ; θ ) T ·Wm , where Wm ∈ Rd×K has the dimension of the embedding d and the dimension of the vocabulary K. The probability distribution of the next word xt under the MoS model is thus a mixture of M Softmaxes weighted by π . Here Θ = ∪Mm=1Wm ∪ θ. PMoS ( xt = k|x1 .. t−1 ; Θ ) = M∑ m πm · σ ( zm ) k ( 3 ) 2.3 TEMPERATURE SCALING . Temperature scaling is an approach used to control the smoothness of the distribution . Instead of applying the Softmax function as suggested in Equation 1 , logits here are divided by the temperature τ before passing through the Softmax layer . P ( xt = k|x1 .. t−1 ; θ , τ ) = σ ( z/τ ) k ( 4 ) Constant Temperature . Earlier works that adopt the constant temperature can be traced back to Model Distillation ( Hinton et al. , 2014 ) . Several works have been taking advantage of a fixed temperature during training ( Norouzi et al. , 2016 ; Ma et al. , 2017 ; Chen et al. , 2019 ) . For instance , ( Norouzi et al. , 2016 ; Ma et al. , 2017 ) optimize log-likelihood on augmented outputs , which are sampled proportionally to their exponential scaled rewards , where the temperature controls the degree of augmentation . Other works incorporate the temperature to enhance calibration during inference ( Guo et al. , 2017 ) , or to achieve a great trade-off between quality and diversity ( Caccia et al. , 2018 ) . Dynamic Temperature Over Training Iterations . A wide range of works adopt a manuallycrafted schedule to tune the temperature during training . Notably in ( Hu et al. , 2017 ) , a new text generation architecture is introduced that combines VAE and a discriminator . Since text samples are discrete and non-differentiable , a continuous approximation based on the Softmax with a decreasing temperature is used to enable gradient propagation from the discriminator to the generator . Similar techniques are adopted in the gumbel-softmax trick ( Jang et al. , 2017 ) , which also allows gradients to pass through discrete sampling objectives . Studies also show that temperature with a heating up schedule makes the embedding vectors more compact ( Zhang et al. , 2018 ) . Dynamic Temperature Over Word Position . Another work that is closely related to our work , is the adaptive temperature over an attention model ( Lin et al. , 2018 ) . We note that contextual temperature further learns the temperature for each vocabulary in the output distribution . 3 METHODS . 3.1 CONTEXTUAL TEMPERATURE . Contextual temperature is a mechanism that chooses the optimal temperature by considering the “ context ” of a word xt . A context of a word includes not only the history x1 .. t−1 , but also the specific vocabulary k that we calculated the probability on . Such a mechanism allowed us to parameterize the temperature vector τ using a deep neural network and adapted the softness of the Softmax layer . Our temperature vector τ ∈ RK was generated from the mapping function f as discussed before . Although f can be any sequential models such as RNN or LSTM , we chose to parameterize it by the AWD-LSTM ( Merity et al. , 2018 ) . We omitted the details of its architecture due to the limited space in this paper . The output of AWD-LSTM was a vector with dimension D. We multiplied this vector by two matrices Wτ1 ∈ RD×Q and Wτ2 ∈ RQ×K . Please note that one can potentially use a single matrix to represent these two . However , doing so can significantly increase the number of parameters and thus in practice we factorized them into two . Finally , we scaled temperatures using a Softmax function over the dimension of the vocabulary and its range was bounded in [ αβ , 1+α β ] . τ = σ ( f ( x1 .. t−1 ; θ ) T ·Wτ1 ·Wτ2 ) + α β ( 5 ) 3.2 CONTEXTUAL TEMPERATURE MOS . We then used the Contextual Temperature Mixture of Softmaxes architecture for language modeling ( CT-MoS ) . The CT-MoS model extends Equation 3 by adding contextual temperature in Equation 5 . Here represents the element-wise division between zm and the temperature vector τ . The results after division are then sent to the Mixture of Softmaxes model . Since new parameters are added to the model , the CT-MoS now has a total of parameters Θ = ∪Mm=1Wm ∪ θ ∪Wτ1 ∪Wτ2 . PCT−MoS ( xt = k|x1 .. t−1 ; Θ ) = M∑ m πm · σ ( zm τ ) ) ( 6 ) One thing worth noticing is that comparing to prior works , the proposed contextual temperature model has the ability to adopt a different temperature for ( a ) different vocabulary in the prediction , ( b ) different position of the same vocabulary given the history and ( c ) a tunable parameter that is capable of changing as the training progresses . The detailed architecture is illustrated in Figure 2 , which highlights the difference between the proposed CT-MoS model and the MoS model . 3.3 TRAINING OBJECTIVES AND LOSS SCALING . We adopt the same regularization techniques in MoS ( Yang et al. , 2018 ) and AWD-LSTM ( Merity et al. , 2018 ) . Our loss function thus consists of four terms : Cross Entropy ( H ) , Activation Regularization ( AR ) , Temporal Activation Regularization ( TAR ) , and Weight Decay ( WD ) . AR is used to penalize high values of outputs , and TAR is used to prevent outputs from changing too much between timesteps . WD prevents the model from overfitting . Here γ1 , γ2 and γ3 are constants to scale regularization terms . And m is the dropout mask ( See in Merity et al . ( 2018 ) ) . L ( Θ ) = ∑K i τi K︸ ︷︷ ︸ LS H ( ŷ , y ) + γ1 L2 ( m f ( x1 .. t−1 ; θ ) ) ︸ ︷︷ ︸ AR +γ2 L2 ( f ( x1 .. t−1 ; θ ) − f ( x1 .. t ; θ ) ) ︸ ︷︷ ︸ TAR +γ3 L2 ( Θ ) ︸ ︷︷ ︸ WD ( 7 ) The uniqueness in our setting is the Loss Scaling term ( LS ) . The use of temperature scaling makes gradients ofH disproportional to that of the other three terms , so the scale ofH needs to be adjusted accordingly . A similar phenomenon is reported in ( Hinton et al. , 2014 ) . In our case , only gradients ofH will be influenced by temperatures . As for the other regularization terms , their gradients are not effected since they do regularization on parameters before temperature scaling . This difference leads to the unbalance between four objectives . Therefore , we scaleH in order to redress the balance . We empirically find that scalingH with the average of temperatures works well in our setting .

Summary: This work proposes a learned and context dependent way to calculate the temperatures for the softmaxes. More specifically, a low-rank affine-transformation, taking the hidden state at the current step as input, is used to calculate scalar weighting for every token in the vocabulary. The method is very general, and can be used in combination with other techniques in tasks such as language modeling and text generation. Experiments on language modeling with Penn TreeBank and WikiText-2 show that the proposed method yields strong performance.