Doubly Sparse: Sparse Mixture of Sparse Experts for Efficient Softmax Inference

Abstract

Computations for the softmax function are significantly expensive when the number of output classes is large. In this paper, we present a novel softmax inference speedup method, Doubly Sparse Softmax (DS-Softmax), that leverages sparse mixture of sparse experts to efficiently retrieve top-k classes. Different from most existing methods that require and approximate a fixed softmax, our method is learning-based and can adapt softmax weights for a better inference speedup. In particular, our method learns a two-level hierarchy which divides entire output class space into several partially overlapping experts. Each expert is sparse and only contains a subset of output classes. To find top-k classes, a sparse mixture enables us to find the most probable expert quickly, and the sparse expert enables us to search within a small-scale softmax. We empirically conduct evaluation on several real-world tasks, including neural machine translation, language modeling and image classification, and demonstrate that significant computation reductions can be achieved at no performance loss.

References in corpus (15)

• Adam: A Method for Stochastic Optimization
• Distilling the Knowledge in a Neural Network
• Sequence to Sequence Learning with Neural Networks
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Recurrent Neural Network Regularization
• Pointer Sentinel Mixture Models
• Efficient softmax approximation for GPUs
• Breaking the Softmax Bottleneck: A High-Rank RNN Language Model
• Learning K-way D-dimensional Discrete Codes for Compact Embedding Representations
• Strategies for Training Large Vocabulary Neural Language Models
• A New Unbiased and Efficient Class of LSH-Based Samplers and Estimators for Partition Function Computation in Log-Linear Models
• Learning K-way D-dimensional Discrete Code For Compact Embedding Representations
• Navigating with Graph Representations for Fast and Scalable Decoding of Neural Language Models
• Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks
• Fast Amortized Inference and Learning in Log-linear Models with Randomly Perturbed Nearest Neighbor Search