Source-linked AI summary
Using the Output Embedding to Improve Language Models
Ofir Press, Lior Wolf
TL;DR
The paper asks whether the output matrix of a neural language model can serve as a useful word embedding and whether sharing it with the input embedding improves models. It analyzes tied updates, evaluates embedding quality and perplexity, and reports comparable translation performance with substantially fewer parameters.
Problem
The output matrix is typically overlooked as a word embedding, leaving its quality and the effects of tying it to the input embedding insufficiently established across language models.
Method
The paper evaluates input and output embeddings, ties them by setting U = V = S, analyzes their updates, and tests the approach in language modeling and translation.
Results
Weight tying improves perplexity across neural network language models, while three-way tying yields similar translation performance with about 52% fewer parameters.
Takeaways & Limitations
The tied embedding evolves more like the untied model’s output embedding, and weight tying can reduce translation-model size without harming reported performance.
Takeaways & Limitations
The paper distinguishes NNLMs from word2vec, where tying the embeddings is reported to be detrimental, and notes that its argument depends on the NNLM architecture’s embedding decoupling.
Abstract
from arXiv · showhide
We study the topmost weight matrix of neural network language models. We show that this matrix constitutes a valid word embedding. When training language models, we recommend tying the input embedding and this output embedding. We analyze the resulting update rules and show that the tied embedding evolves in a more similar way to the output embedding than to the input embedding in the untied model. We also offer a new method of regularizing the output embedding. Our methods lead to a significant reduction in perplexity, as we are able to show on a variety of neural network language models. Finally, we show that weight tying can reduce the size of neural translation models to less than half of their original size without harming their performance.
1 Introduction
Neural language models use an input embedding to represent the current word and an output matrix to score every vocabulary word. The paper argues that both matrices constitute word embeddings and studies tying them together.
- Model architecture: The input word is mapped through U to hidden activations, then V produces one score per vocabulary word before softmax converts scores into next-word probabilities.The model computes h3 = V h2 and applies softmax to obtain p.
- Two embeddings: U is the input embedding, whereas V is the output embedding.Both matrices have the same size in the described LSTM-based models.
- Two embeddings: Rows for similar words are expected to be similar in both matrices, but the desired similarity reflects different roles for inputs and output scores.Input embeddings should react similarly to synonyms, while output embeddings should assign similar scores to interchangeable words.
- Paper focus: The paper compares input and output embedding quality and proposes tying them to improve neural language models.Its introduction reports evaluations across word2vec, recurrent language models, and neural translation models.
2 Related Work
Prior work established neural language modeling, word2vec embeddings, and related uses of weight tying, but left the role and effects of output embeddings insufficiently tested. This paper positions its analysis within those developments and distinguishes outcomes across model families.
- Neural language models: Neural network language models evolved from feedforward models through recurrent and LSTM architectures, with dropout-based variants later introduced.The cited sequence includes Bengio et al., recurrent models, LSTM language modeling, and dropout methods.
- Word2vec: Skip-gram word2vec learns both input and output word embeddings, although the input vectors are typically returned after training.Some prior work instead uses both embeddings to compute word similarity.
- Weight tying: In the paper’s experiments, tying embeddings is detrimental in word2vec but improves performance in neural network language models.This contrast motivates treating weight tying as model-family dependent rather than universally beneficial.
- Neural machine translation: Neural machine translation decoders function as language models conditioned on prior target words and the source sentence.Recent systems also segment source and target words into subword units, including with byte pair encoding.
- Weight tying: Weight tying had appeared in a prior log-bilinear model without an explained decision or tested performance effect.Concurrent work independently offered an explanation for weight tying in neural network language models.
3 Weight Tying
The paper applies weight tying across neural language-model and translation settings, analyzes how tied embeddings are updated, and adds projection regularization for smaller models. For translation, it also proposes tying decoder input/output and encoder input embeddings when source and target subwords overlap.
- Weight Tying: Weight tying is evaluated in NNLMs, word2vec skip-gram, and neural machine translation models.Translation models additionally use a three-way weight-tying method.
- Weight Tying: Projection regularization inserts P before the output embedding and adds λ∥P∥2 to the small model’s loss, with λ = 0.15.The method lets one embedding serve input and output roles while adapting the output through a regularized projection.
- Update Analysis: In untied NNLMs, only the current input-word row updates in U, whereas every row of V updates at each timestep.Consequently, rare input-word vectors receive relatively few updates.
- Update Analysis: Tied NNLMs set U = V = S, so every row updates each iteration and most rows follow the untied output-embedding update more closely than the input-embedding update.The current input-word row also receives an input-embedding contribution, while other rows receive no such contribution.
- Translation Models: Up to 90% of BPE subwords are shared between English and French, and 85% between English and German.This overlap motivates three-way tying of decoder input, decoder output, and encoder input embeddings in a shared source/target vocabulary.
4 Results
The experiments evaluate embedding quality, embedding similarity, perplexity, and parameter efficiency across word2vec, neural language models, and neural translation models. Weight tying improves NNLM performance and preserves translation performance while reducing model size.
- Embedding quality: Embedding quality is evaluated with five benchmarks based on cosine-distance correlations against human judgments.The benchmarks are Simlex999, Verb-143, MEN, Rare-Word, and MTurk771.
- Embedding quality: The word2vec output embedding is almost as good as the input embedding, whereas tying embeddings produces a distinct representation.The comparison uses Spearman’s correlation on the five embedding evaluation benchmarks.
- Embedding similarity: For NNLMs, tied embeddings resemble the untied output embedding more than the untied input embedding.This pattern is reported for both small and large NNLMs using Spearman correlations between pairwise word-similarity values.
- Neural network language models: Weight tying significantly reduces validation and test perplexity on PTB without reducing training-set perplexity.The result is interpreted in the paper as indicating less overfitting after reducing the number of parameters.
- Neural network language models: Projection regularization does not help large dropout-regularized models, while experiments also evaluate tied and regularized models across text8, IMDB, and BBC.The small and large NNLM experiments use different regularization settings, and the cross-dataset evaluation reports the combined effects of projection regularization and weight tying.
- Neural machine translation: 28% fewer parameters preserves similar translation performance with weight tying, while three-way tying reaches about 52% fewer parameters with similar performance.The translation experiments report model size and performance comparisons for standard and three-way weight tying.