Self-organized Hierarchical Softmax
Yikang Shen Shawn Tan Chrisopher Pal Aaron Courville
cs.CL cs.LG
 

Shifted lattices and asymptotically optimal ellipses
    [
    December 30, 2023
====================================================



 
We propose a new self-organizing hierarchical softmax formulation for neural-network-based language models over large vocabularies. Instead of using a predefined hierarchical structure, our approach is capable of learning word clusters with clear syntactical and semantic meaning during the language model training process. We provide experiments on standard benchmarks for language modeling and sentence compression tasks. We find that this approach is as fast as other efficient softmax approximations, while achieving comparable or even better performance relative to similar full softmax models.
 



§ INTRODUCTION

The softmax function and its variants are an essential part of neural network based models for natural language tasks, such as language modeling, sentence summarization, machine translation and language generation.

Given a hidden vector, the softmax can assign probability mass to each word in a vocabulary. The hidden vector could be generated from the preceding context, source sentence, dialogue context, or just random variables. The model decides how the context is converted into the hidden vector, and there are several choices for this, including recurrent neural network <cit.>, feed forward neural network <cit.> or log-bilinear models <cit.>. In our experiments here, we use a long short-term memory (LSTM) model for language modeling, and a sequence-to-sequence model with an attention mechanism for sentence compression. Both models are simple but have been shown capable of achieving state-of-the-art results. Our focus is to demonstrate that with a well designed structure, the hierarchical softmax approach can perform as accurately as the full softmax, while maintaining improvements in efficiency.

For word-level models, the size of the vocabulary is very important for higher recall and a more accurate understanding of the input. However the training speed for models with softmax output layers quickly decreases as the vocabulary size grows. This is due to the linear increase of parameter size and computation cost with respect to vocabulary. 

Many approaches have been proposed to reduce the computational complexity of large softmax layers <cit.>. These methods can largely be divided into two categories:


  * Approaches that can compute a normalized distribution over the entire vocabulary with a lower computational cost <cit.>. Normalized probabilities can be useful for sentence generation tasks, such as machine translation and summarization. 


  * Methods that provide unnormalized values <cit.>. These methods are usually more efficient in the training process, but less accurate. 


In this paper, we propose a self-organized hierarchical softmax, which belongs in the first category. In contrast to previous hierarchical softmax methods which have used predefined clusters, we conjecture here that a hierarchical structure learned from the corpus may improve model performance. Instead of using term frequencies as clustering criteria <cit.>, we want to explore the probability of clustering words together considering their preceding context. The main contributions of this paper are as follows:


  * We propose an algorithm to learn a hierarchical structure during the language model learning process. The goal of this algorithm is to maximize the probability of a word belonging to its cluster considering its preceding context.

  * We conduct experiments for two different tasks: language modeling and sentence summarization. Results show
that our learned hierarchical softmax can achieve comparable accuracy for language modeling, and even better performance for summarization when compared to a standard softmax. We also provide clustering results, which indicate a clear semantic relevance between words in the same cluster.

  * Empirical results show that our approach provides a more than a 3× speed up compared to the standard softmax.
 



§ RELATED WORK

Representing probability distributions over large vocabularies is computationally challenging. In neural language modeling, the standard approach is to use a softmax function that output a probability vector over the entire vocabulary. Many methods have been proposed to approximate the softmax with lower computational cost <cit.>. We briefly review the most popular methods below. 



 §.§ Softmax-based approaches



  
Hierarchical Softmax (HSM): <cit.> and its variants are the most popular approximations. In general, this approach organizes the output vocabulary into a tree where the leaves are words and intermediate nodes are latent variables, or classes. The tree structure could have many levels and there is a unique path from root to each word. The probability of a word is the product of probabilities of each node along its path. In practice, we could use a tree with two layers, where we want to organize words into simple clusters. In this case, the computational complexity reduces from O(|V|) to O(√(|V|)). If we use a deeper structure like the Huffman Tree, the computational complexity could be reduced to O(log|V|). In general, the hierarchical structure is built on frequency binning <cit.> or word similarities <cit.>. In this paper, we propose another word-similarity-based hierarchical structure. But, instead of performing k-means over pre-learned word embeddings, we propose a new approach that learns hierarchical structure based on the model's historical prediction during the language model learning process. 



  
Differentiated softmax (D-softmax): <cit.> is based on the intuition that not all words require the same number of parameters: The many occurrences of frequent words allows us to fit many parameters to them, while extremely rare words might only allow us to fit relatively few parameters. D-softmax assign different dimension of vector to words according to their frequency to speed up the training and save memory.
Adaptive softmax <cit.> can be seen as a combination of frequency binning HSM and 
D-softmax. 



  
CNN-softmax: <cit.> is inspired by the idea that we could use convolution network to produce word embedding from a character level model. Aside from a big reduction in number of parameters and incorporating morphological knowledge from words, this method can also easily deal with out-of-vocabulary words, and allows parallel training over corpora that have different vocabulary size.
But this method does not decrease the computational complexity compared to the standard full softmax <cit.>.



 §.§ Sampling-based approaches

Sampling based approaches approximate the normalization in the denominator of the softmax with some other loss that is cheap to compute. However, sampling based approaches are only useful at training time. During inference, the full softmax still needs to be computed to obtain a normalized probability. These approaches have been successfully applied to language modeling <cit.>, machine translation <cit.>, and computer vision <cit.>.



  
Importance sampling (IS): <cit.> select a subset of the vocabulary as negative samples to approximate the softmax normalization. 
Originally unigram or bigram distribution of word in entire corpus are used for sampling negative samples <cit.>, but researchers found that sampling from a more carefully designed distribution could help achieve a better accuracy.
Instead, two variants of n-gram distributions are proposed:


  * an interpolated bigram distribution and unigram distribution <cit.>, 

  * a power-raised unigram distribution <cit.>.




  
Noise Contrastive Estimation (NCE): Noise Contrastive Estimation (NCE) is proposed in <cit.> as a more stable sampling method than IS. NCE does not try to estimate the probability of a word directly. Instead, it uses an auxiliary loss that works to distinguish the original distribution from a noisy one. <cit.> showed that good performance can be achieved even without computing the softmax normalization.



§ SELF-ORGANIZED HIERARCHICAL SOFTMAX



 §.§ Cluster-based Hierarchical Softmax

We employ a modified 2-layer hierarchical softmax to compute the distribution of next word in a sentence. Given vocabulary 𝐕 of size N, and pre-softmax hidden states 𝐡, we first project 𝐡 into a cluster vector 𝐡_c and a word vector 𝐡_w,

    [ 𝐡_c; 𝐡_w ]
    =Relu[ 𝐖_c 𝐡; 𝐖_w 𝐡 ]

where 𝐖_c, 𝐖_w ∈ℝ^d × d. The cluster distribution can be expressed as

    P(c|𝐡 )=exp( 𝐡_c ^𝖳𝐔^𝒞_c) /∑ _c' ∈𝒞exp( 𝐡_c ^𝖳𝐔^𝒞_c')

where 𝒞 is set of clusters, 𝐔^𝒞∈ℝ^|𝒞| × d is vector representation of clusters.
The in-cluster probability function is

    P(w_t|𝐡, 𝒞(w_t))=exp( 𝐡_w ^𝖳𝐔^𝐕_w_t) /∑ _w' ∈𝒞(w_i)exp( 𝐡_w ^𝖳𝐔^𝐕 _w')

where 𝐔^𝐕∈ℝ^|𝐕| × d is vector representation of words, 

𝒞(w_t) is the cluster assigned to w_t in 𝒞.
Thus, the final probability function is

    P(w_t|𝐡) = P(𝒞(w_t)|𝐡 ) P(w_t|𝐡, 𝒞(w_t) )

If the number of cluster is in O(√(N)) and the maximum number of words in cluster is in O(√(N)), then the computational cost of normalization at each layers is only O(√(N)), (as opposed to O(N) for the standard softmax).  Thus a large matrix dot product is transformed into two small matrix dot product, which are very efficient on a GPU <cit.>.



 §.§ Cluster Perplexity

In order to evaluate the quality of a clustering over words, we propose the cluster perplexity:

    ppl_cluster(𝒞) = 2^1/M∑_w_t - log_2 p(𝒞(w_t) | 𝐰_<t)

where M is number of words in the dataset, 𝐰_<t is context preceding w_t, p(𝒞(w_t) | 𝐰_<t) is the probability that words in cluster 𝒞(w_t) appear behind 𝐰_<t. Given a word cluster 𝒞 and 𝐰_<t, this metric evaluate the difficulty to choose correct cluster. If words that share similar context have been successfully grouped together, the ppl_cluster should be small.

In addition to ppl_cluster, we also propose the in-cluster perplexity:

    ppl_in-cluster ( 𝒞 ) = 2^1/M∑ - log_2 p(w_t | 𝐰_<t, 𝒞(w_t))

where p(w_t | 𝐰_<t, 𝒞(w_t)) is the probability of word w_t appearing after 𝐰_<t given a subset of vocabulary 𝒞(w_t). If 𝒞(w_t) contains words that share the same context with w_t, ppl_in-cluster should be large.



 §.§ Optimizing Cluster Perplexity

With the definitions in Equations <ref> and 
<ref> established, our goal is to minimize ppl_cluster(𝒞):

    argmin_𝒞ppl_cluster(𝒞) 
       = - argmax_𝒞1/M∑_w_tlog_2 p(𝒞(w_t) | 𝐰_<t) 
       = - argmax_𝒞∑_w ∈𝐕n_w/M∑_w_t = wlog_2 p(𝒞(w) | 𝐰_<t)/n_w
       = - argmax_𝒞∑_w ∈𝐕 tf(w) g(w,𝒞(w))

where tf(w)=n_w/M is term frequency of word w in the corpus, and

    g(w,𝒞(w))=1/n_w∑_w_t = wlog_2 p(𝒞(w) | 𝐰_<t)

is the average of log_2 p(𝒞(w) | 𝐰_<t) over different preceding contexts 𝐰_<t that followed by word w.

According to equation <ref>, we need a 𝒞 that maximize the weighted sum of g(w,𝒞(w)). While directly computing g(w,𝒞(w)) is intractable, 
the output of equation <ref> at training time can be considered a sample of p(c | 𝐰_<t). We propose to use exponential smoothing to estimate g(w,𝒞(w)):

    q_τ(w_t,c) =    λ(w_t) q_τ-1(w_t,c) 
       + (1 - λ(w_t)) log_2 P(c|𝐡_t)

q_τ(c|w_t) is a weighted sum over all historical samples of p(c|𝐰_<t) under different context and parameters, previously seen in training. 
The smoothing factor λ(w_t) is defined as

    λ(w_t) = 1/f_w_t

where f_w_t is the raw count of w_t in the entire dataset.



 §.§ Greedy Word Clustering

In practice, we assigned two constraints to each cluster:


  * The number of words in each cluster cannot exceed γ√(N), where 
γ > 1 is a hyperparameter; and 

  * The sum of term frequencies in each cluster should be smaller than a frequency budget f_b. This is known as the the frequency-budget trick <cit.>. 

These constraints prevent us from getting clusters that are either too large, which make computing in-cluster distributions very expensive, or too unbalanced in frequency, which will bias our word cluster distribution.


The greedy algorithm <ref> is proposed to optimize the cluster perplexity. As each cluster has limited positions for words, some words w cannot be assign to their best cluster c = argmax_c q(w,c). If we assign words according to certain order, then words at the tail end of the sequence will be less likely to be assign to their best cluster. In the algorithm, we assign words to clusters in descending order of their term frequency tf(w). In this schema, high frequency words w have priority to choose clusters, because they have higher weight in equation <ref>.






 §.§ Training Language Model with Self-organized HSM

In the training phase, we start from a randomly initialized word cluster, and update parameters using gradient descent based optimization algorithms, updating word cluster every K iterations. K is a hyperparameter that is chosen based on dataset size and vocabulary size. 
This learning process can also be considered as an EM algorithm: In the E-Step, we update the clusters; in M-Step, we update parameters based on the new clusters.



§ LANGUAGE MODELING EXPERIMENT

Language Modeling (LM) is a central task in NLP. The goal of LM is to learn a probability distribution over a sequence of tokens from a given vocabulary set 𝐕. The joint distribution is defined as a product of conditional distribution of tokens given their preceding context. Given a sequence of word w_1,...,w_T ∈𝐕, the probability distribution can be defined as:

    P(w_1,...,w_T) = ∏_t=1^T P(w_t|w_1,...,w_t-1)

To address this problem, much work has been done on both parametric and non-parametric approaches. In recent years, parametric models based on neural networks have became the standard method. In our experiment, we used the standard word-level Long-Short Term Memory (LSTM) model, since multiple works show it can obtain state-of-the-art performance on different datasets <cit.>.



 §.§ Dataset

We evaluate our method on the text8[http://mattmahoney.net/dc/textdata] dataset, and use the perplexity (ppl) as an evaluation metric. We also provide the training time for full softmax and our approach. Text8 is a standard compression dataset containing a pre-processed version of the first 100 million characters from Wikipedia in English. It has been recently used for language modeling (Mikolov et al., 2014) and has a vocabulary of 44k words. The dataset partitioned into a training set (first 99M characters) and a development set (last 1M characters) that is used to report performance <cit.>.




 §.§ Implementation

In our experiments, we use the same setting as the one reported in  <cit.>. A one-layer LSTM model is used. Both the dimension of hidden state and dimension of the input word embeddings is set to 512. LSTM parameters are regularized with weight decay (λ=10^-6). Batch size is set to 128. We use Adagrad <cit.> with learning rate 0.1, the norm of the gradients is clipped to 0.25, and a 20 steps gradient truncation is applied. 

For our model, we set the number of clusters to √(𝐕), and the maximum number of words in each cluster is γ√(𝐕) with γ=1.5, and frequency budget f_b is 0.1 as in <cit.>. We update the word clusters every K=1000 mini-batches.



 §.§ Baseline Methods

We compare the proposed approach with (1) full softmax, (2) importance sampling <cit.>, (3) hierarchical softmax (HSM) with frequency binning <cit.>, (4) differentiated softmax <cit.>, and (5) adaptive softmax <cit.>. As we use the same implementation settings in <cit.>, we use their experiment results for baseline methods. Instead of using torch, we use theano <cit.> to implement our approach. Thus, in order to compare computation time, we implement another full softmax language model with theano. Our full softmax has the same perplexity on the development set as the one reported in <cit.>.



 §.§ Experimental results







Table <ref> shows results on the text8 dataset. Our approach provides the best perplexity among all approximation approaches, nearly performing as well as a full softmax. Table <ref> shows that our approach is almost 4 times faster than a normal softmax, with the
speed-up continuing to increase as the vocabulary size increases.

Figure <ref> monitors learning process of our approach. At the beginning of training, we observe a high cluster perplexity, and very low in-cluster perplexity. Because we initialize clusters randomly, our model has difficulties to predict cluster given the preceding sequence context of the target word. 
As training continues, our cluster update algorithm 
considers the cluster word assignments based on the distribution given by Equation <ref>, resulting in the cluster perplexity decreasing rapidly.
In contrast, the in-cluster perplexity first increases and then decreases slowly. Because our approach assigns similar words into the same cluster that are difficult to distinguish,
the model has to explicitly learn small differences between words that share a similar context.
In the end, the model reaches a balance between cluster and in-cluster perplexity.



Table <ref> shows some examples of words that belong to the same cluster. We observe a strong syntactical similarity between these words, and semantic closeness. These examples show that our cluster update algorithm is capable of placing words with similar context into the same cluster. It is interesting to see that the unsupervised approach could learn word cluster with clear meaning. 



§ ABSTRACTIVE SENTENCE SUMMARIZATION EXPERIMENT

Summarization is an important challenge in natural language understanding. The aim is to produce a condensed representation of an input text that captures the core meaning of the original. 

Given a vocabulary 𝐕 and a sequence of M words x_1, ..., x_M ∈𝐕, a summarizer takes 𝐱 as input and outputs a shortened sentence 𝐲 of length N < M. Assuming that words in the output sentence also belong to 𝐕, we can express the output as y_1, ..., y_N ∈𝐕. The output sentence is then called a summary of the input sentence. Thus, the probability distribution of the summary can be defined as:

    P(y_1,...,y_N) = ∏_t=1^N P(y_t|𝐱, y_1,...,y_t-1)

For an extractive summarization, the probability distribution P(y_t|𝐱, y_1,...,y_t-1) is on the set of input words, while for an abstractive summarization the distribution is on the entire vocabulary. In this experiment, we focus on the abstractive summarization task, which is more difficult and computationally expensive.



 §.§ Dataset

We trained our model on the Gigaword5 dataset <cit.>. This dataset was generated by pairing the headline of each article with its first sentence to create a source-compression pair. <cit.> provided scripts to filter out outliers, resulting in roughly 3.8M training pairs, a 400K validation set, and a 400K test set. We use the most frequent 69k words in the title as input and output vocabulary, which correspond to the decoder vocabulary size used in <cit.>. Out of vocabulary words are represented with a symbol <unk>.

We evaluate our method on both the standard DUC-2004 dataset and the single reference Gigaword5 test set. The DUC corpus[http://duc.nist.gov/duc2004/tasks.html] 2004 corpus consists of 500 documents, each having 4 human generated reference titles. Evaluation of this dataset uses the limited-length Rouge Recall at 75 bytes on DUC validation and test sets. In our work, we simply run the models trained on Gigaword corpus as they are,  without tuning them on the DUC validation set. The only change we have made to the decoder is to suppress the model from emitting the end-of-summary tag, and forced it to emit exactly 30 words for every summary. <cit.> provides a random sampled 2000 title-headline pair as test set. We acquired the exact test sample used by them to make a precise comparison of our models with theirs. Like  and , we use the full length F1 variant of Rouge[http://www.berouge.com/Pages/default.aspx] to evaluate our system.



 §.§ Implementation

In this experiment, we use the standard Encoder-Decoder with attention architecture. Both encoder and decoder a consists of a single layer uni-direction LSTM model, and an attention mechanism over the source-hidden states and a softmax layer to output a distribution probability over an output vocabulary. 
The hidden state dimension and input embedding are both set to 512. All parameters are regularized with weight decay (λ=10^-6). Batch size is 128. We use Adam <cit.> with a learning rate of 0.001. No dropout or gradient clipping is used. At decode time, we use beamsearch of size 5 to generate summary. The maximum length of output summary is limited to 30.

For our approach, we learn word clusters through training a language model on the titles of our training dataset. We then use these clusters as the fixed structure for hierarchical softmax in the summarization model.



 §.§ Baseline methods

We compared the performance of our model with state-of-the-art models that are trained with teacher forcing and cross-entropy loss, including: (1) TOPIARY <cit.>, (2) ABS+ <cit.>, (3) RAS-Elman <cit.>, and (4) words-1vk5k-1sent <cit.>. We also include our implementation of normal softmax as a baseline method. There are some newly proposed summarization models that come up with different type of loss function including the reconstructive loss function <cit.>, and the minimum risk training (MRT) loss <cit.>. We did not compare with these methods, since this experiment is focused on evaluating our approach against other softmax-based approaches under similar implementations and learning settings.



 §.§ Experiment Results









Table <ref> and Table <ref> show results of our approach comparing the different methods. Our approach not only outperforms the full softmax method, but also outperforms state-of-the-art methods on most of evaluation metrics. Table <ref> also shows that our approach is 3 times faster than standard full softmax, which is widely used in all kinds of different summarization models. Figure <ref> present examples of summaries generated by self-organized HSM in comparing with true headline and full softmax outputs.

As the word cluster learned on Gigaword corpus shows similar in-cluster syntactical similarity and semantic closeness, we suggest that this hierarchical structure decomposes the difficult word generation task into 2 easier tasks. The first task is to decide correct syntactical role for next word, and the second task is to find semantically correct words in a subset of the vocabulary.



§ CONCLUSION

In this paper we have proposed a new self organizing variant of the hierarchical softmax. We observe that this approach can achieve the same performance as a full softmax approach for language modelling, and even better performance for sentence summarization. Our approximation approach is also as efficient as other hierarchical softmax approximation techniques. In particular in our experiments we observe that our self-organized HSM is at least 3 times faster than a full softmax approach.

Our approach yields self organized word clusters which are influenced by the context of words.
Examining word clusters produced by or approach reveals that our method groups words according to their syntactical role and semantic similarity. These results are appealing in that we have obtained a certain level of understanding of grammar rules without explicit part-of-speech tagged input. We therefore think that the use of this approach shows promise for other NLP tasks, including machine translation and natural language generation.


emnlp_natbib

