Source-linked AI summary
Efficient softmax approximation for GPUs
Edouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, Hervé Jégou
TL;DR
Large-vocabulary neural language models make full softmax computation expensive because scores and normalization span the vocabulary. The paper introduces adaptive softmax, which uses frequency-aware hierarchical clusters and GPU-aware computation, and reports 2× to 10× speed-ups with accuracy close to the exact model.
Problem
Full softmax has O(dk) complexity and often dominates computation when the vocabulary contains many words.
Method
Adaptive softmax builds an approximate hierarchical model using unbalanced frequency-based clusters and a computation-time model for modern GPU matrix multiplications.
Results
2× to 10× speed-ups over exact softmax are reported across GPU experiments while maintaining low perplexity and performance on par with the exact model.
Takeaways & Limitations
The reported efficiency enables training on extremely large corpora in reasonable time without requiring many GPUs.
Takeaways & Limitations
The computation-time model is described as very crude, although it explains empirical observations well.
Abstract
from arXiv · showhide
We propose an approximate strategy to efficiently train neural network based language models over very large vocabularies. Our approach, called adaptive softmax, circumvents the linear dependency on the vocabulary size by exploiting the unbalanced word distribution to form clusters that explicitly minimize the expectation of computation time. Our approach further reduces the computational time by exploiting the specificities of modern architectures and matrix-matrix vector operations, making it particularly suited for graphical processing units. Our experiments carried out on standard benchmarks, such as EuroParl and One Billion Word, show that our approach brings a large gain in efficiency over standard approximations while achieving an accuracy close to that of the full softmax. The code of our method is available at https://github.com/facebookresearch/adaptive-softmax.
1. Introduction
The paper targets the computational cost of neural language modeling with very large vocabularies. It introduces adaptive softmax, a GPU-oriented approximate hierarchical model evaluated on standard benchmarks.
- Motivation: Very large vocabularies make neural language models computationally intensive despite their strong benchmark performance.The One Billion Word benchmark contains around 800K words, and computing probabilities over the whole vocabulary can become prohibitively slow.
- Approach: Adaptive softmax forms an approximate hierarchical model designed for efficient processing on GPUs.The design is motivated by increasingly performant GPUs relative to CPUs.
- Approach: The model explicitly accounts for computation time in matrix-matrix multiplications on modern architectures.Unlike previous hierarchical models, it does not assume computation is trivially linear in matrix dimensions.
- Results: 2× to 10× speed-ups are reported against regular softmax.The paper also evaluates the method on Text8, Europarl, and One Billion Word datasets.
- Results: On the largest corpus, the higher efficiency comes at no cost in accuracy for a given amount of training data.This contrasts with concurrent approaches improving efficiency.
2. Related work
Prior work approximates softmax through hierarchical, sampling-based, self-normalized, or other output-space methods. Adaptive softmax differs by optimizing hierarchy structure for realistic computation on modern GPUs.
- Hierarchical softmax: Hierarchical softmax reduces softmax cost by organizing words into tree-based classes.Existing hierarchies use word similarities, frequency binning, or Huffman coding based on frequency.
- Hierarchical softmax: Frequency-based Huffman coding does not account for complexity reductions from matrix-matrix multiplication and distributed computation on modern GPUs.This motivates a hierarchy designed around hardware-aware computation.
- Hierarchical softmax: A prior complexity-oriented hierarchy places words only at leaves, causing a reported 5−10% performance drop and requiring word-similarity optimization.Adaptive softmax allows classes at internal nodes, which leads to almost no performance drop in the cited comparison.
- Alternative approximations: Other approaches include frequent-word short-lists, frequency-dependent classifier capacity, negative-target sampling, and self-normalized classifiers.These methods approximate or avoid full softmax computation in different ways.
- Alternative approximations: The cited exact high-dimensional-output method cannot be directly applied to softmax, limiting its potential application to language modeling.
3. Preliminaries on language modeling
Neural language models predict the next word from preceding context and normalize scores over the vocabulary with softmax. Class-based hierarchical softmax reduces this cost by factoring word probabilities through classes.
- Language modeling: Language modeling learns a probability distribution over word sequences by conditioning each token on its past.Parametric neural models follow earlier nonparametric counting-based approaches such as smoothed N-grams.
- Feedforward networks: Feedforward language models concatenate one-hot representations of a fixed context window to predict the next word.The hidden state and next-token scores are computed by a two-layer neural network.
- Feedforward networks: The feedforward model uses weight matrices P, A, and B and can be trained with stochastic gradient descent or Adagrad.Its hidden transformation uses a nonlinearity such as the pointwise sigmoid function.
- Recurrent networks: Recurrent networks update the hidden state using both the current input and the previous hidden state.Efficient approximation uses truncated back-propagation through time and norm clipping.
- Softmax computation: Exact softmax computes and normalizes a score for every vocabulary word, requiring O(k) operations and O(dk) overall neural-network complexity.For large vocabularies, this computation often dominates the whole model.
- Hierarchical softmax: Class-based hierarchical softmax assigns each word to a class and factorizes its probability into class and within-class probabilities.With equal-sized classes, the computational cost is reduced relative to full softmax, although the supplied passage truncates the resulting expression.
4. Our approach: the adaptive softmax
Adaptive softmax organizes words into GPU-oriented hierarchical clusters whose sizes and capacities are chosen to reduce expected computation while preserving frequent-word accuracy. The method exploits unbalanced word frequencies and measured matrix-multiplication behavior, achieving substantial speedups in the two-cluster case.
- Method: Adaptive softmax builds a hierarchical probability model designed specifically to reduce computation time on GPUs.The first level contains frequent words and cluster vectors; rare words are handled in second-level clusters.
- GPU-aware design: GPU matrix multiplication is inefficient when one dimension is small, motivating wide clusters rather than binary hierarchies or tiny rare-word clusters.The computation-time model is crude but explains empirical timing observations well.
- Two-cluster case: The two-cluster design uses a small, high-probability head and a large, low-probability tail, with root and tail costs modeled as C = g(kh + 1, B) + g(kt, ptB).The head contains frequent words, while the tail contains many rare words and is evaluated on a reduced effective batch size.
- Accuracy trade-off: Frequent words remain directly accessible through a shortlist because assigning every word to a leaf cluster can reduce performance by around 5−10%.Direct root access gives frequent words sharper inference than multiplying cluster and within-cluster probabilities.
- Results: More than 5× speedup over full softmax is observed for the optimized two-cluster split on Bulgarian Europarl.The optimal head size need not make head and tail probabilities equal.
- General case: The general model assigns different classifier capacities across clusters and finds that 10–15 clusters give the best modeled computation time, with limited gains beyond five clusters.Rare words receive less capacity because they are observed less often and are harder to learn reliably.
5. Experiments
The experiments evaluate adaptive softmax across datasets with different vocabulary sizes, using perplexity over training time or epochs and comparing against standard approximation baselines. Results show strong efficiency while retaining competitive perplexity, including on the largest benchmark.
- Evaluation protocol: Experiments cover Text8, Europarl, and One Billion Word, spanning vocabularies from 44k to almost 800k words.Perplexity is measured as a function of training time or training data.
- Baselines: The comparison includes full softmax, hierarchical softmax, importance sampling, and differentiated softmax variants.The improved D-softmax [*] uses the whole hidden representation with projection matrices of different sizes.
- Text8: On Text8, adaptive softmax is fastest and comes within 3 perplexity points of full softmax.The improved D-softmax [*] achieves similar results but is slower by a factor ×1.8.
- Europarl: 2× to 3× faster than improved D-softmax [*], adaptive softmax significantly outperforms competitors on Europarl's perplexity-versus-training-time tradeoff.Similar conclusions are reported for other Europarl languages.
- One Billion Word: 43.9 test perplexity after five epochs is achieved on One Billion Word in less than three days on a single GPU.The paper reports this as lower than 50 on a single GPU, while the only lower reported perplexity uses a model 8× bigger, 32 GPUs, and three weeks of training.
6. Conclusion
The paper presents an efficient softmax approximation that maintains low perplexity while achieving 2× to 10× GPU speed-ups over the exact model. This efficiency supports training on extremely large corpora without requiring many GPUs.
- 2× to 10× speed-ups over the exact model are achieved while maintaining low perplexity in GPU experiments.The method is described as performing on par with the exact model.
- The method explicitly accounts for matrix-multiplication time on parallel systems and combines this with a frequent-word shortlist and reduced rare-word capacity.
- The resulting efficiency allows extremely large corpora to be handled in reasonable time without a large number of GPUs.
- The authors suggest that the approach may extend to other parallel architectures, losses, and domains with unbalanced class distributions.