PaperlayerView exact revision on arXiv

Source-linked AI summary

Breaking the Softmax Bottleneck: A High-Rank RNN Language Model

Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, William W. Cohen

arXiv:1711.03953v4cs.CLcs.LG

TL;DR

The paper argues that Softmax-based language models face a rank-based expressiveness bottleneck when modeling highly context-dependent natural language. It introduces Mixture of Softmaxes, which combines multiple Softmax distributions to increase expressiveness, and reports improved results across standard and large-scale benchmarks. The approach is presented as a way to obtain higher rank without the overfitting associated with simply enlarging embeddings.

  • Problem

    Softmax-based language models are limited by word-embedding dimension, while natural language may require a high-rank representation of context-dependent distributions.

  • Method

    Mixture of Softmaxes introduces discrete latent variables and computes a weighted average of multiple Softmax distributions for next-token prediction.

  • Results

    MoS outperforms all baselines by up to 3.6 points in perplexity on PTB and WT2, and outperforms Softmax by over 5.6 points on 1B Word.

  • Takeaways & Limitations

    The results support using a high-rank model for natural language while avoiding the overfitting associated with non-parametric models and naively increasing embedding dimensions.

  • Takeaways & Limitations

    The mixture-of-contexts alternative still suffers from the same rank limitation as Softmax.

Abstract

from arXiv · show

We formulate language modeling as a matrix factorization problem, and show that the expressiveness of Softmax-based models (including the majority of neural language models) is limited by a Softmax bottleneck. Given that natural language is highly context-dependent, this further implies that in practice Softmax with distributed word embeddings does not have enough capacity to model natural language. We propose a simple and effective method to address this issue, and improve the state-of-the-art perplexities on Penn Treebank and WikiText-2 to 47.69 and 40.68 respectively. The proposed method also excels on the large-scale 1B Word dataset, outperforming the baseline by over 5.6 points in perplexity.

1 INTRODUCTION

The paper identifies a Softmax bottleneck: standard Softmax-based recurrent language models are limited by the rank of their word-embedding factorization, despite expressive recurrent encoders. It proposes Mixture of Softmaxes (MoS), which improves benchmark perplexity and learns higher-rank matrices.

  • Problem: Standard Softmax-based recurrent language models factorize context-dependent token probabilities through word embeddings, limiting their expressiveness.RNNs encode contexts into vectors, multiply them by word embeddings to produce logits, and apply Softmax.
  • Problem: The Softmax bottleneck arises because natural language may require factorizing a high-rank matrix, while Softmax models are constrained by embedding dimension.The paper frames standard Softmax language modeling as a matrix factorization problem.
  • Method: Mixture of Softmaxes introduces discrete latent variables and combines multiple Softmax distributions into the next-token probability distribution.MoS uses K component distributions with learned mixture weights.
  • Results: MoS learns matrices with much larger normalized singular values and higher rank than Softmax and other baselines on real-world datasets.This provides an empirical measure of its greater expressiveness.
  • Results: 3.6 points: MoS improves over current state-of-the-art benchmark results in perplexity, reaching 47.69 on Penn Treebank and 40.68 on WikiText-2.The paper also reports improved performance over Softmax and other baselines on a dialog dataset.

2 LANGUAGE MODELING AS MATRIX FACTORIZATION

The paper recasts language modeling as matching context-conditioned next-token distributions through a matrix factorization, exposing a rank-based capacity limit in standard Softmax models. It proposes Mixture of Softmaxes to increase expressiveness without the parameter growth and generalization concerns of naïvely increasing dimensionality.

  • 2 LANGUAGE MODELING AS MATRIX FACTORIZATION: Language modeling asks whether a parameterized distribution can match the true next-token distribution for every context.The analysis treats language as context–conditional distribution pairs and studies whether some parameter setting reproduces them exactly.
  • 2 LANGUAGE MODELING AS MATRIX FACTORIZATION: Standard Softmax language models correspond to factoring a log-probability matrix with context representations and word embeddings of dimension d.The rows of Hθ, Wθ, and A represent contexts, tokens, and true log probabilities, respectively.
  • 2 LANGUAGE MODELING AS MATRIX FACTORIZATION: If d < rank(A) −1, no Softmax-based model can match the true distribution for every context.The formal bottleneck follows from the minimum rank required among logits equivalent under row-wise shifts.
  • 2.2 HYPOTHESIS: NATURAL LANGUAGE IS HIGH-RANK: The bottleneck matters because natural language is hypothesized to have a high-rank log-probability matrix, while practical embedding dimensions are much smaller.The paper notes that d is typically around 10^2 whereas rank(A) may approach M, around 10^5, so Softmax learns a low-rank approximation.
  • 2.3 EASY FIXES?: N-gram and high-dimensional Softmax alternatives increase expressiveness but dramatically increase parameters, creating overfitting and generalization concerns.The paper frames language modeling as a tradeoff between expressiveness and generalization.
  • 2.4 MIXTURE OF SOFTMAXES: A HIGH-RANK LANGUAGE MODEL: Mixture of Softmaxes introduces discrete latent components, mixes K Softmax distributions, and can produce arbitrarily high-rank approximations while retaining similar model size.MoS can reduce d to compensate for mixture parameters; with K = 1 it reduces to Softmax.

3 EXPERIMENTS

Experiments show that Mixture of Softmaxes (MoS) improves language-modeling performance across benchmarks and dialog, while ablations and rank analyses support its proposed mechanism.

  • 3 EXPERIMENTS: MoS is evaluated on PTB, WikiText-2, 1B Word, and Switchboard using perplexity and dialog evaluation metrics.The benchmark setup includes comparable model sizes and standard baselines.
  • 3.1 MAIN RESULTS: Up to 3.6 points in perplexity separate MoS from prior state-of-the-art results on PTB and WikiText-2 with comparable parameters.MoS outperforms baselines with or without dynamic evaluation.
  • 3.1 MAIN RESULTS: Over 5.6 points in perplexity separate MoS from Softmax on the 1B Word dataset.The authors note that MoS hyper-parameters were not tuned because of limited computational resources.
  • 3.1 MAIN RESULTS: MoS outperforms MoC and Softmax on all reported Switchboard evaluation metrics.The dialog experiment tests MoS as a decoder structure for context-dependent distributions.
  • 3.2 ABLATION STUDY: MoC does not consistently improve over AWD-LSTM, while AWD-LSTM trained with MoS hyper-parameters performs substantially worse.These ablations distinguish the MoS structure from simply adding a hidden layer, using a feature-space mixture, or changing hyper-parameters.
  • 3.3 VERIFY THE ROLE OF RANK: MoS induces an almost full-rank log-probability matrix, whereas Softmax and MoC ranks remain limited by their embedding sizes.On PTB, the MoS matrix nearly reaches full rank with M = 10000.
  • 3.3 VERIFY THE ROLE OF RANK: Increasing MoS mixture components increases empirical rank and improves PTB performance before full rank is reached.The reported positive correlation between rank and performance supports the proposed high-rank explanation.
  • 3.3 VERIFY THE ROLE OF RANK: MoS and Softmax show similar generalization gaps on 1B Word, while MoS does not improve character-level modeling where Softmax lacks rank limitation.MoS also incurs a computational cost, with a two- to three-times slowdown for up to 15 mixture components in most settings.

4 RELATED WORK

Prior work approaches mixture structures and latent variables in language and sequence modeling, but differs from this paper in where mixtures act, their motivation, and how latent variables are optimized.

  • Matrix-rank methods previously improved Ngram generalization through sparse-plus-low-rank approximation, whereas this work targets high-rank neural language models.
  • Neubig and Dyer’s approach mixes Ngram and neural models, but separately training its components can limit expressiveness.
  • Mixture of Softmaxes is related to Mixture of Experts, but differs because prior Gaussian mixtures model continuous data and use mixture structure for a distinct motivation.
  • Earlier mixture architectures such as Eigen et al. and Shazeer et al. apply mixtures in feature space, targeting expressiveness or efficient conditional computation.
  • Most prior latent-variable sequence models use continuous per-step variables and optimize a variational lower bound with SGVB, unlike the discrete-latent formulation here.

5 CONCLUSIONS

The paper identifies the Softmax bottleneck as a word-embedding-dimension limit on expressiveness and proposes Mixture of Softmaxes to increase expressiveness without sacrificing generalization. Its benchmark gains support the claim that high-rank models are important for natural language.

  • Softmax-based language models are limited by word-embedding dimension under the matrix factorization framework, creating the Softmax bottleneck.
  • Mixture of Softmaxes increases expressiveness while avoiding the overfitting associated with non-parametric models and naively larger word embeddings.
  • Large benchmark improvements support the paper’s theoretical reasoning that natural language benefits from high-rank modeling.

A PROOFS

The proofs characterize when Softmax-based conditional distributions can be represented through matrix factorization, including the role of row-shift invariance and embedding dimension.

  • Softmax-equivalent logit matrices can differ by adding a scalar to every entry in each row.
  • The row-shift structure can be represented using a diagonal matrix multiplied by an all-ones matrix.
  • Adding the all-ones direction increases the possible row-space dimension by at most one in the proof’s rank argument.
  • If a target distribution is represented by a parameterized recurrent model, its logit product belongs to the Softmax-equivalent matrix family.
  • When the embedding dimension reaches the minimum required rank, universal approximation provides factors realizing the target matrix and distribution.

B EXPERIMENT SETTING AND HYPER-PARAMETERS

The appendices summarize the hyper-parameters used for MoS language modeling and dynamic evaluation, including variational dropout terminology and references to detailed prior descriptions.

  • The MoS language-modeling hyper-parameters are summarized in an appendix table.
  • V-dropout denotes variational dropout in the MoS hyper-parameter table.
  • The dynamic-evaluation hyper-parameters are summarized separately in another appendix table.

B.2 1B WORD DATASET

This section evaluates model rank through singular-value distributions and expected pairwise KLD, using normalized singular values to compare models across scales.

  • B.2 1B WORD DATASET: Rank is estimated from the non-zero singular values of the empirical log-likelihood matrix.The analysis also studies singular-value distributions because roundoff can make direct rank computation error-prone.
  • B.2 1B WORD DATASET: Normalized singular values are plotted cumulatively as the percentage below each threshold in Figure 1.Normalization maps singular values from different models to [0, 1].
  • B.2 1B WORD DATASET: Softmax and MoC concentrate most singular values at very low values, whereas MoS has a different concentration pattern.The concentration pattern is used as an indicator of relative matrix rank.
  • B.2 1B WORD DATASET: Expected pairwise KLD measures how dissimilar next-step conditional distributions are across contexts as another rank indicator.The metric is E_c,c′∼C[KLD(P(X | c)∥P(X | c′))], where C denotes all possible contexts.

C.2 AN INVERSE EXPERIMENT ON CHARACTER-LEVEL LANGUAGE MODELING

The inverse experiment tests whether MoS helps when Softmax is not rank-limited. Character-level language modeling provides this setting because its vocabulary is small.

  • C.2 AN INVERSE EXPERIMENT ON CHARACTER-LEVEL LANGUAGE MODELING: Character-level language modeling has a log-likelihood matrix whose rank is bounded by the vocabulary size.Because CharLM usually uses tens of characters and embedding sizes are hundreds, Softmax is no longer a bottleneck in this setting.
  • C.2 AN INVERSE EXPERIMENT ON CHARACTER-LEVEL LANGUAGE MODELING: Softmax and MoS obtain the same test-set BPC and comparable validation BPC on text8.The experiment uses a 1-layer 1024-unit LSTM baseline and MoS variants with 7 or 10 mixtures under matched capacity.
  • C.2 AN INVERSE EXPERIMENT ON CHARACTER-LEVEL LANGUAGE MODELING: The CharLM results match the hypothesis that MoS and Softmax perform similarly when the Softmax rank limitation is absent.The authors use this contrast with word-level modeling to support their Softmax-bottleneck analysis.

C.3 MOS COMPUTATIONAL TIME

MoS introduces additional computational cost, but GPU matrix multiplication makes its wall-time growth sub-linear in the number of Softmax components.

  • C.3 MOS COMPUTATIONAL TIME: MoS is generally two to three times slower than Softmax in the controlled experiments.The models are controlled to have comparable model sizes across two training-time comparisons.
  • C.3 MOS COMPUTATIONAL TIME: Under the same-batch-size setting, MoS causes 1.9x, 2.5x, and 3.8x slowdowns on PTB, WT2, and 1B, respectively.The “bs” setting measures computational cost with enough memory.
  • C.3 MOS COMPUTATIONAL TIME: The authors consider best-1 and best-3 settings to better reflect computational cost when each model uses its own best-performing hyperparameters.These settings use one and three GPUs, respectively.

C.4 QUALITATIVE ANALYSIS

The qualitative analysis compares MoS and MoC on matched contexts, showing how their next-token predictions differ in context sensitivity and use of less common words.

  • C.4 QUALITATIVE ANALYSIS: MoS adjusts top predictions for different topics following the same preceding word, whereas MoC produces similar choices across contexts.The examples use contexts with the preceding word “N” and compare prediction steps where MoS has lower negative log loss.
  • C.4 QUALITATIVE ANALYSIS: MoS uses political context to predict country-related completions after “south” and to include “ual” among top predictions.MoC instead gives generic predictions in both examples.
  • C.4 QUALITATIVE ANALYSIS: MoS accurately exploits less common words according to context, while MoC fails to produce such choices.The authors relate these examples to MoS’s capacity for modeling context-dependent language.
  • C.4 QUALITATIVE ANALYSIS: Table 14 compares next-token predictions on Penn Treebank test data using contexts containing the previous and current sentences.In the table, N denotes a number introduced by preprocessing.
Loading 1711.03953v4…