PaperlayerView exact revision on arXiv

Source-linked AI summary

Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling

Hakan Inan, Khashayar Khosravi, Richard Socher

arXiv:1611.01462v3cs.LGcs.CLstat.ML

TL;DR

Conventional RNN language models rely on isolated one-hot input and output representations, limiting use of word relationships and requiring many parameters. The paper introduces an embedding-informed loss and derives tying the input embedding to the output projection. The framework improves language-model performance across Penn Treebank and Wikitext-2 while reducing trainable variables.

  • Problem

    Conventional RNN language models use one-hot targets and treat input and output words as isolated entities, despite meaningful distances and shared spaces defined by word embeddings.

  • Method

    The framework adds KL-divergence supervision from an embedding-based target distribution and reuses the input embedding matrix as the output projection matrix.

  • Results

    The proposed framework outperforms the conventional framework on Penn Treebank and Wikitext-2, with reused embeddings improving VD-RHN perplexity by 2.5 on Penn Treebank.

  • Takeaways & Limitations

    The framework improves learning while greatly reducing trainable variables through input-embedding reuse.

  • Takeaways & Limitations

    Large-network results on Wikitext-2 were not reported because computational resources were insufficient for some configurations.

Abstract

from arXiv · show

Recurrent neural networks have been very successful at predicting sequences of words in tasks such as language modeling. However, all such models are based on the conventional classification framework, where the model is trained against one-hot targets, and each word is represented both as an input and as an output in isolation. This causes inefficiencies in learning both in terms of utilizing all of the information and in terms of the number of parameters needed to train. We introduce a novel theoretical framework that facilitates better learning in language modeling, and show that our framework leads to tying together the input embedding and the output projection matrices, greatly reducing the number of trainable variables. Our framework leads to state of the art performance on the Penn Treebank with a variety of network models.

1 INTRODUCTION

The paper identifies inefficiencies in conventional RNN language-model training and introduces a loss framework that uses word-vector structure to improve learning while reusing model parameters.

  • Motivation: Conventional RNN language models use one-hot classification targets without an output metric, despite meaningful distances between words in embedding space.They also treat inputs and outputs as isolated entities, although both occupy identical semantic spaces.
  • Approach: The proposed framework adds a KL-divergence term matching predictions to a target distribution estimated from word-vector similarity.Its theory motivates reusing the input embedding matrix as the output classification matrix.
  • Evaluation: Extensive experiments evaluate the framework on Penn Treebank and Wikitext-2 against the conventional training framework.The experiments test whether the observed gains are consistent across datasets.
  • Results: Models trained with the proposed framework significantly outperform models trained with the conventional framework.The reported comparison covers both benchmark datasets.

2 BACKGROUND: RECURRENT NEURAL NETWORK LANGUAGE MODEL

An RNN language model maps word-token sequences and recurrent hidden states to a probability distribution over the next vocabulary word, typically using cross-entropy against a one-hot target.

  • Model formulation: The recurrent network updates its hidden state from the current input and previous hidden state, then projects that state through a softmax over vocabulary words.The recurrent unit may be a standard RNN, GRU, or LSTM.
  • Model formulation: The input embedding matrix L has dimensions dx × |V|, where dx is embedding dimension and |V| is vocabulary size.Each vocabulary word is represented by a column of L.
  • Model formulation: The output projection matrix W and bias b produce the |V|-dimensional next-word distribution from the hidden state.W has dimensions |V| × dh, and dh is the hidden-state size.
  • Conventional loss: Training almost exclusively minimizes cross-entropy between the model prediction distribution and the observed one-hot word token.Because the target is one-hot, this loss is equivalent to KL divergence up to the target entropy.

3 AUGMENTING THE CROSS-ENTROPY LOSS

The proposed loss augments one-hot supervision with a word-embedding-based target distribution, so each example supplies probability-weighted information across related vocabulary classes.

  • Loss framework: The augmented loss adds a term that minimizes KL divergence between the model prediction and an estimated target distribution based on word-embedding similarity.The framework aims to match predictions to a more accurate estimate of the true conditional data distribution.
  • Optimization effect: Each training example updates all classes with nonzero target probability, with each class’s step weighted by its conditional probability.The paper characterizes this deterministic update as less noisy and as providing improved supervision unless the target is always one-hot.
  • Relation to prior work: The framework also relates to knowledge distillation, but estimates distributions from word embeddings rather than a separately trained large model.Its knowledge transfer occurs between different parts of the same network.
  • Target distribution: The embedding-based target distribution assigns probability mass according to inner products between the target word vector and other word vectors, adjusted by temperature.Similar words therefore receive similar probability masses.

4 THEORETICALLY DRIVEN REUSE OF WORD EMBEDDINGS

The analysis shows that augmented-loss training constrains output logits toward the embedding-defined space, motivating direct reuse of the input embedding matrix in the output layer.

  • Assumptions: The theoretical argument assumes equal embedding and hidden-state dimensions, zero bias, zero training loss, and a large temperature parameter.These assumptions define the analyzed setting rather than the full practical training setup.
  • Logit matching: At sufficiently high temperature, the augmented loss acts to match the model’s logits to the logits of the more informative target distribution.The derivation uses high-temperature approximations of the prediction distributions.
  • Subspace analysis: Under the stated theoretical assumptions, the output projection and transposed embedding matrices span equivalent column spaces.This permits writing W = L^T A for a square linear mapping A.
  • Parameter tying: Reusing the transposed embedding matrix in the output projection lets the hidden state pass through a learned linear mapping before producing the same result.The neural network supplies the mapping h → Ah.
  • Practical modification: The proposed practical constraint sets W = L^T and b = 0, explicitly restricting the output probability space to a subspace governed by the embeddings.This removes a large vocabulary-dependent matrix and reduces the work required by the augmented loss.

5 RELATED WORK

Prior work improved recurrent language models through architectures, regularization, and auxiliary components, but did not change the loss structure. Related approaches to informed distributions and shared embeddings differ in training strategy or lack theoretical justification.

  • Earlier RNNLM improvements included dropout methods, novel recurrent units, and pointer networks, but none addressed the loss structure.
  • The paper’s loss framework resembles Hinton et al.’s KL-divergence augmentation but estimates target distributions from word-vector similarity rather than large teacher networks.
  • Press and Wolf concurrently reused the word-embedding matrix in output projection, but their approach was empirical and lacked theoretical justification.
  • Earlier language models also shared input and output representations implicitly, whereas this work proposes sharing as a supplement to a baseline.

6 EXPERIMENTS

Experiments on PTB and Wikitext-2 evaluate augmented loss and reused embeddings across LSTM sizes, showing complementary gains and strongest performance when combined.

  • Experimental setup: The experiments use PTB and Wikitext-2, with 2-layer variational-dropout LSTMs sized at 200, 650, and 1500 hidden units.PTB uses a 10k-word vocabulary; Wikitext-2 provides a larger alternative benchmark.
  • Theory validation: The augmented loss drives W toward the input-embedding column space, reaching subspace distance around 0.06 when used alone.At temperature τ = 10, increasing the augmented-loss proportion reduces the distance from almost 1 to around 0.06.
  • Experimental setup: Four models compare the baseline against augmented loss, reused embeddings, and both modifications together.The variants are VD-LSTM, VD-LSTM +AL, VD-LSTM +RE, and VD-LSTM +REAL.
  • PTB and Wikitext-2 results: All AL, RE, and REAL models significantly outperform the baseline, while combining RE and AL gives the best performance on both datasets.The comparison covers final validation and test perplexities across network sizes.
  • PTB and Wikitext-2 results: AL helps smaller networks more, whereas RE significantly outperforms AL for larger networks; RE also reduces model size while preserving representational power.Partitioned Wikitext-2 experiments support the explanation that AL contributes more when training data are limited.
  • Comparison with prior work: The best LSTM model, VD-LSTM+REAL, outperforms prior conventional-framework work on PTB, while VD-RHN +RE improves VD-RHN by 2.5 perplexity.The paper reports the VD-RHN +RE result as the best overall performance in the comparison.
  • Qualitative results: Qualitatively, VD-LSTM +REAL assigns lower probability to <unk> and frequent function words while assigning more probability to words close to the target.The proposed predictions sometimes select semantically close and plausible alternatives.

7 CONCLUSION

The framework links embedding-informed loss training with reusing the input embedding matrix in the output projection. Experiments support its effectiveness across datasets and its applicability beyond vanilla language modeling.

  • The framework uses word-embedding structure to provide more informed supervision than one-hot targets and reuses the input embedding in the output projection.Together, these changes reduce trainable variables while improving learning in language modeling.
  • Experiments on Penn Treebank and Wikitext-2 show that the framework outperforms the conventional language-modeling framework.
  • The improvements are applicable to neural machine translation, speech recognition, and text summarization, especially with large vocabularies.The framework also offers the benefit of greatly reducing the number of parameters to train.

A MODEL AND TRAINING DETAILS

Training uses scheduled learning-rate decay, truncated backpropagation, gradient clipping, and recurrent dropout configurations that vary by network size. Augmented-loss training uses temperature and weighting settings selected separately for each dataset.

  • Training begins with learning rate 1, then decays at epochs 5, 10, and 1 for small, medium, and large networks, respectively.The decay rate is 0.9 for small and medium networks and 0.97 for large networks.
  • Networks are unrolled for 35 steps during backpropagation on both Penn Treebank and Wikitext-2.
  • Gradient norms are clipped at 5 for small and medium networks and 6 for large networks on both datasets.
  • Dropout uses shared masks through the unrolled network, with probabilities of 0.7, 0.5, and 0.35 for small, medium, and large Penn Treebank networks.Dropout is not applied to the input embedding layer.
  • Augmented-loss training uses temperature τ = 20 and sets α = γτ, with γ between 0.5–0.8 on Penn Treebank and 1.0–1.5 on Wikitext-2.Moderate variations in τ or α did not cause sudden performance deterioration.

B METRIC FOR CALCULATING SUBSPACE DISTANCES

The subspace-distance metric compares matrices through their column spaces: orthonormal bases are constructed, one subspace is projected onto the other, and the residual norm is measured. The resulting distance is related to principal angles and ranges from zero to one.

  • The metric distance between matrices X and Y is computed by comparing their spanned column spaces.
  • Orthonormal bases U and V are obtained for the column spaces of X and Y, potentially using QR decomposition.
  • The method projects V onto U using S = UU^T V and forms the residual R = V − S.
  • The distance uses the Frobenius norm of the residual and the number of residual columns in its calculation.
  • The distance is zero for identical column spaces, one for orthogonal subspaces, and is related to the sines of principal angles.The principal-angle expression describes d² as the average of sin²(θ_i), a quantity between 0 and 1.
Loading 1611.01462v3…