From Fully Trained to Fully Random Embeddings: Improving Neural Machine Translation with Compact Word Embedding Tables

One of the main challenges in NLP is to properly encode discrete tokens into continuous vector representations. The most common practice in this regard is to use an embedding matrix which provides a one-to-one mapping from tokens to $n$-dimensional, real-valued vectors (Mikolov et al., 2013; Li and Yang, 2018). Typically, values of these vectors are optimized via back-propagation with respect to a particular objective function. Learning embedding matrices with robust performance across different domains and data distributions is a complex task, and can directly impact quality Tian et al. (2014); Shi et al. (2015); Sun et al. (2016).

Embedding matrices are a key component in a seq2seq NMT model, for instance, in a transformer-base NMT model Vaswani et al. (2017) with a vocabulary size of $50k$, $36\%$ of the total model parameters are utilized by embedding matrices. Thus, a significant amount of parameters are utilized only for representing tokens in a model. Existing work on embedding compression for NMT systems Khrulkov et al. (2019); Chen et al. (2018); Shu and Nakayama (2017), have shown that these matrices can by compressed significantly with minor drop in performance. Our focus in this work is to study the importance of embedding matrices in the context of NMT. We experiment with Transformers Vaswani et al. (2017) and LSTM based seq2seq architectures, which are the two most popular seq2seq models in NMT. In Section 3.1, we compare the performance of a fully-trained embedding matrix, with a completely random word embedding (RWE), and find that using a random embedding matrix leads to drop in about $1$ to $4$ BLEU Papineni et al. (2002) points on different NMT benchmark datasets. Neural networks have shown impressive performance with random weights for image classification tasks Ramanujan et al. (2020), our experiments show similar results for embedding matrices of NMT models.

RWE uses no trainable parameters, thus it might be possible to recover the drop in BLEU score by increasing the number of parameters using additional layers. We increase the number of layers in RWE model such that the number of parameters used by the entire model is comparable to fully-trained embedding model. We find that even in comparable parameter setting RWE model performance was inferior to fully-trained model by $~{}1$ BLEU point for high and medium resource datasets. Our results suggest that even though the embedding parameters can be compressed to a large extent with only a minor drop in accuracy, embedding matrices are essential components for token representation, and cannot be replaced by deeper layers of transformer based models.

RWE assumed a fully random embedding sampled from a single Gaussian distribution; however, to control the amount of random information, and to better understand the importance of embedding matrices, we introduce Gaussian product quantization (GPQ). GPQ assumes that $k$ Gaussian distributions are required to approximate the embedding matrix for NMT models. The means and variances of the $k$ distributions are learned from a fully-trained embedding matrix trained on the same dataset as GPQ model. $k$ is a hyper-parameter which controls the amount of information distilled from a pre-trained embedding to partially random embedding in GPQ model. GPQ has the ability to move from a fully random embedding to a fully-trained embedding by increasing the number of distributions $k$, as shown in Figure 1. Our results show that only $50$ Gaussian distributions are sufficient to approximate the embedding matrix, without drop in performance. GPQ compresses the embedding matrix $5$ times without any significant drop in performance, and uses only $100$ floating point values for storing the means and variances. GPQ demonstrates effective regularization of the embedding matrix, by out-performing the transformer baseline model with fully-trained embedding by $1.2$ BLEU points for $En\rightarrow Fr$ dataset and $0.5$ BLEU for $Pt\rightarrow En$ dataset. A similar increase in performance was also observed for LSTM based models, thus further showing the effectiveness of GPQ for embedding regularization.

Product Quantization (PQ) (Jegou et al., 2010) was proposed for fast search by approximating nearest neighbour search. PQ has been adapted for compressing embedding matrices for different NLP problems Shu and Nakayama (2017); Kim et al. (2020); Tissier et al. (2019); Li et al. (2018). Our method is an extension of PQ, first, we incorporate variance information in GPQ, then we define unified partitioning to learn a more robust shared space for approximating the embedding matrix. Our extensions consistently outperform the original PQ-based models with better compression rates and higher BLEU scores. These improvement are observed for Transformer-based models as well as conventional, recurrent neural networks.

PQ is the core of our techniques, so we briefly discuss its details in this section. For more details see Jegou et al. (2010). As previously mentioned, NLP models encode tokens from a discrete to a continuous domain via an embedding matrix E, as simply shown in Equation 1:

where $V$ is a vocabulary set of unique tokens and $n$ is the dimension of each embedding. We use the notation $E_{w}\in\mathbb{R}^{n}$ to refer to the embedding of the w-th word in the vocabulary set. In PQ, first $E$ is partitioned into $g$ groups across columns with $G_{i}\in\mathbb{R}^{|V|\times\frac{n}{g}}$ representing the i-th group. Then each group $G_{i}$ is clustered using K-means into $c$ clusters. Cluster indices and cluster centroids are stored in the vector $Q_{i}\in\mathbb{N}^{|V|}$ and matrix $C_{i}\in\mathbb{R}^{c\times\frac{n}{g}}$, respectively. Figure 2 illustrates the PQ-based decomposition process.

We can obtain the quantization matrix $Q$ (also known as the index matrix in the literature) and codebook $C$ by concatenating $Q_{i}$ and $C_{i}$ for all $g$ groups along the columns, as shown in Equation 2.

The decomposition process applied by PQ is reversable, namely any embedding table decomposed by PQ can be reconstructed using the matrices $Q_{i}$ and $C_{i}$. The size of an embedding table compressed via PQ can be calculated as shown in Equation 3:

where $f_{p}$ is the number of bits defined to store parameters with floating points. In our experiments $f_{p}$ is $32$ for all settings. We use Equation 3 to measure the compression rate of different models.

In this section we first explain RWE which stands for Random Word Embeddings. In RWE, all embeddings are initialized with completely random values where there is no syntactic and semantic information is available.

Next, we move from completely random embeddings to weakly supervised embeddings by incorporating some knowledge from a pre-trained embedding matrix. We propose our Gaussian Porduct Quantization (GPQ) technique in this regard, which applies PQ to a pre-trained embedding matrix and models each cluster with a Gaussian distribution. Our GPQ empowers the PQ with taking intra-cluster variance information into account. This approach is particularly useful when PQ clusters have a high variance, which might be the case for large embedding matrices.

We evaluate our model for machine translation task on WMT-2014 English to French ($En\rightarrow Fr$), WMT-2014 English to German ($En\rightarrow De$), and IWSLT Portuguese to English ($Pt\rightarrow En$) datasets. We chose these pairs as they are good representatives of high, medium, and low resource language pairs.

For $En\rightarrow Fr$ experiments we used Sentence-Piece Kudo and Richardson (2018) to extract a shared vocabulary of 32k sub-word units. We chose newstest2013 as our validation dataset and used newstest2014 as our test dataset. The dataset contains about 36M sentence-pairs.

For $En\rightarrow De$ experiments we use the same setup as Vaswani et al. (2017). The dataset contains about 4.5M sentence-pairs. We use a shared vocabulary of 37k sub-word units extracted using Sentence-Piece.

For $Pt\rightarrow En$ experiments, we replicate dataset configuration of Tan et al. (2019) for individual models. Specifically, the dataset contains 167k training pairs. We used a shared vocabulary of 32k subwords extracted with Sentence-Piece.

There are different embedding compression techniques in the literature, these methods are based on either singular value decomposition (SVD) or product quantization (PQ). In general, SVD-based methods Chen et al. (2018); Khrulkov et al. (2019) approximate embedding matrices using fewer parameters by projecting high-dimensional matrices to lower dimensions. Therefore, the learning process is modified to work with new and smaller matrices where a multiplication of these matrices reconstruct original values. The model proposed by Chen et al. (2018) is a well-known example of SVD-based model in this field. In this technique, words are first divided into groups based upon frequency, then weighted SVD is applied on each group, which allows a compressed representation of the embedding matrix. Khrulkov et al. (2019) used a multi-linear variation instead of regular SVD. The embedding dimensions are first factorized to obtain indices of smaller tensor embedding. The tensor embedding has fewer parameters and can be used to reconstruct the original embedding matrix.

Shu and Nakayama (2017) proposed to compress the embedding matrix by learning an integer vector for all words in vocabulary (code matrix), using the Gumbel softmax technique. The code matrix is equivalent to the quantization matrix in PQ. An embedding vector is reconstructed by looking up all the values in the code matrix from the corresponding set of vector dictionaries. The vectors are added (in PQ the vectors are concatenated) to obtain the final embedding vector.

Svenstrup et al. (2017) proposed a hashing based technique with a common memory for storing vectors. The method is similar to Shu and Nakayama (2017), with two key differences. First, the code matrix is not learned but assigned uniformly using a hashing function. Second, instead of a simple summation of component vectors, weighted summation is applied to obtain the final embedding vector.

Kim et al. (2020) learn the code matrix by utilizing binarized code learning introduced in Tissier et al. (2019). The key novelty introduced by Kim et al. (2020), is to learn different length of code vector for each word.

Our work points towards the need of rethinking the process of encoding tokens to real valued vectors for machine translation. We show that Transformer model is capable of recovering useful syntactic and semantic information from a random assignment of embedding vectors. Our variant of product quantization was able to approximate the embedding matrix in an almost $~{}100$% discrete space, with better performance than the baseline model. A discrete space is easier to interpret, compared to a continuous space, and can motivate future research to handle unknown tokens through optimum cluster selection.