Efficient Softmax Explained

Softmax encounters large computing cost when the output vocabulary size is very large. Some feasible approaches will be explained under the circumstance of skip-gram pretraining task.

Softmax

The training objective of Skip-gram model is predicting the surrounding words in the sequences. Given a sequence of words w₁,w₂, …, w_T. The final objective is to minimize the average log probability:

$$J(\theta) = \frac{1}{T} \sum_{t=1}^T \sum_{-c \leq j \leq c,\, j\neq 0} \log p(w_{t+j \vert w_t})$$

The output of skip-gram formula is:

$$p(w_O \vert w_I) = \frac{\exp({v'}_{wO})^\top v_{wI}}{\sum_{w=1}^W\exp({v'_w}^\top v_{wI})}$$

where $v_w$ and $v’_w$ are the “input” vector representation (i.e., the hidden layer output h_i in previous fig) and the “output” vectors (i.e., the weights of W’_{N x V}).

Solution

Hierarchical Softmax

Hierarchical Softmax (HSM)^[1] reduce the computing cost from $O(V)$ to $O(\log V)$, where V denotes the output vocab size. The V leaf nodes in the binary tree represent the vocabulay words wilst the (V-1) inner nodes denote the probability mass.

For each leaf node, there exits a unique path from the root node to that unit, on which path the probability is estimated along the nodes in the path. As the figure below shown, the highlighted path is from root to the leaf node w₂, denoted $n(w,4)$, where $n(w,j)$ denotes the j^th node on the path from root to the word $w$.

$$p(w=w_O) = \prod_{j=1}^{L(w)-1} \sigma \bigg( \mathbb{I} \big[ n(w, j+1) = ch(n(w,j)) \big] {v'}_{n(w,j)}^\top h \bigg)$$

where $ch(n)$ is the left child of node $n$, $v’_{n(w,j)}$ denote the vector representation of the inner node $n(w,j)$, $\mathbf{h}$ is the output of hidden layer:

$$\mathbf{h}=\left\{ \begin{array}{ll} \mathbf{v}_{wI} \quad \text{in skip-gram}\\ \frac{1}{C} \sum_{c=1}^C \mathbf{v}_{wc} \quad \text{in CBOW} \end{array} \right.$$

$\mathbb{I} \big[ \cdot \big]$ is defined as:

$$\mathbb{I} \big[ \cdot \big]=\left\{ \begin{array}{ll} 1 \quad \text{if x is true}\\ -1 \quad \text{otherwise} \end{array} \right.$$

The marked path can be computed as:

$$\begin{align} p(w_2 = w_O) &= p(n(w_2, 1), \text{left}) \cdot p(n(w_2, 2), \text{left}) \cdot p(n(w_2, 3), \text{right}) \\ &= \sigma \big( {v}_{n(w_2,1)}^{\prime \top} h \big) \cdot \sigma \big( {v}_{n(w_2,2)}^{\prime \top} h \big) \cdot \sigma \big( -{v'}_{n(w_2,3)}^\top h \big) \end{align}$$

It is obvious that

$$\sum_{i=1}^V p(w_i = w_O) = 1$$

Negative Sampling

Negative Sampling(NEG)^[2] selects a subset of negative targets as well as the target to approximate the softmax normalization.
The objective is defined as:

$$E = \log \sigma (\mathbf{v}^\prime_{wO} \mathbf{h}) - \sum_{w_j \in W_\text{neg}} \log \sigma (-\mathbf{v}^{\prime\top}_{w_j} \mathbf{h})$$

where w_O is the output word, $\mathbf{v}^\prime_{wO}$ is the output vector. $h$ is the hidden layer output, and the values on Skip-gram and CBOW are as descibed in the previous section.

Adaptive Softmax

Adaptive Softmax partition the vocab into clusters, V_h and V_t, representing the head of distribution consisting of the most frequent words and the tail denoting rare words. Each cluster is regarded as independent, and the capacity of clusters are adaptive. ^[2] found that the small number of clusters can obtain comparable results as the original softmax on the very large vocabulary size.

References

---

1. 1.Rong, X. (2014). word2vec Parameter Learning Explained. ArXiv, abs/1411.2738. ↩
2. 2.Mikolov, T., Sutskever, I., Chen, K., Corrado, G.S., & Dean, J. (2013). Distributed Representations of Words and Phrases and their Compositionality. NIPS. ↩
3. 3.Mohammed, A.A., & Umaashankar, V. (2018). Effectiveness of Hierarchical Softmax in Large Scale Classification Tasks. 2018 International Conference on Advances in Computing, Communications and Informatics (ICACCI), 1090-1094. ↩
4. 4.Grave, E., Joulin, A., Cissé, M., Grangier, D., & Jégou, H. (2016). Efficient softmax approximation for GPUs. ArXiv, abs/1609.04309. ↩
5. 5.Baevski, A., & Auli, M. (2018). Adaptive input representations for neural language modeling. arXiv preprint arXiv:1809.10853. ↩