Title: DISTILLED EMBEDDING: NON-LINEAR EMBEDDING FACTORIZATION USING KNOWLEDGE DISTILLATION

Abstract: Word-embeddings are a vital component of Natural Language Processing (NLP) systems and have been extensively researched. Better representations of words have come at the cost of huge memory footprints, which has made deploying NLP models on edge-devices challenging due to memory limitations. Compressing embedding matrices without sacrificing model performance is essential for successful commercial edge deployment. In this paper, we propose Distilled Embedding, an (input/output) embedding compression method based on low-rank matrix decomposition with an added non-linearity. First, we initialize the weights of our decomposition by learning to reconstruct the full word-embedding and then fine-tune on the downstream task employing knowledge distillation on the factorized embedding. We conduct extensive experimentation with various compression rates on machine translation, using different data-sets with a shared word-embedding matrix for both embedding and vocabulary projection matrices. We show that the proposed technique outperforms conventional low-rank matrix factorization, and other recently proposed word-embedding matrix compression methods.

1. INTRODUCTION:

Deep Learning models are the state-of-the-art in NLP, Computer Vision, Speech Recognition and many other fields in Computer Science and Engineering. The remarkable deep learning revolution has been built on top of massive amounts of data (both labeled and unlabeled), and faster computation. In NLP, large pre-trained language models like BERT (Devlin et al., 2018) are state-of-the-art on a large number of downstream NLP problems. The largest publicly available language model to date is trained with 1.6 billion parameters (Keskar et al., 2019) . On machine translation the stateof-the-art models have parameters in the order of millions. Data privacy and server cost are some major issues, driving research towards deploying these models on edge-devices. However, running these models on edge-devices, face memory and latency issues due to limitations of the hardware. Thus, there has been considerable interest towards research in reducing the memory footprint and faster inference speed for these models (Sainath et al., 2013; Acharya et al., 2019; Shu & Nakayama, 2017; Shi & Yu, 2018; Jegou et al., 2010; Chen et al., 2018; Winata et al., 2018) .

The architecture of deep-learning-based NLP models can be broken down into three components. The first component, represents the embedding section, which maps words in the vocabulary to continuous dense vector representations of the words. For all future references, the first component includes both the source and target vocabulary mapping using a common embedding matrix with a shared vocabulary. The second component, consists of a function f , typically a deep neural-network (Schmidhuber, 2015; Krizhevsky et al., 2012; Mikolov et al., 2010) which maps the embedding representation for different NLP problems (machine-translation, summarization, question-answering and others), to the output-space of function f . The third component, is the output layer which maps the output of function f to the vocabulary-space, followed by a softmax function. Most of these NLP problems use linear transformations, in the first and the third components of the models. Since, these components depend upon a large vocabulary-size, they require large number of parameters which results in higher latency and larger memory requirements. For instance, the Transformer Base model (Vaswani et al., 2017) uses 37% of the parameters in the first and third components using a vocabulary size of 50k, and with parameter-tying between the components. The percentage of parameters increases to 54%, when parameters are not shared between the first and third components. The reader is encouraged to look at Table A.1 of the Appendix for more information.

Thus, an obvious step is to compress parametric functions used by the first and third components. Recently, many researchers have worked on compressing word-embedding matrices (Sainath et al., 2013; Acharya et al., 2019; Shu & Nakayama, 2017; Shi & Yu, 2018; Jegou et al., 2010; Chen et al., 2018; Winata et al., 2018) . These techniques have proven to perform at-par with the uncompressed models, but still suffer from a number of issues.

First, embedding compression models (Shi & Yu, 2018; Chen et al., 2018; Khrulkov et al., 2019; Shu & Nakayama, 2017) , require two hyper-parameters to be fine-tuned. These hyper-parameters influence the number of parameters in the model, and thus the compression rate. This leads to an additional layer of complexity for optimizing the model for different NLP problems. Additionally, Chen et al. (2018) requires an additional optimization step for grouping words, and lacks end-to-end training through back-propagation. Shi & Yu (2018) also requires an additional step for performing k-means clustering for generating the quantization matrix. Thus, most of the current state-of-the-art systems are much more complicated to fine-tune for different NLP problems and data-sets.

Second, all the state-of-the-art embedding compression models compress the first and third components separately. In practice, state-of-the-art NLP models like Vaswani et al. (2017) have shown to perform better with parameter sharing between the first and third components (Press & Wolf, 2016) . Thus, there is a need for an exhaustive analysis of various embedding compression techniques, with parameter sharing.

Lastly, embedding compression models not based on linear SVD (Khrulkov et al. (2019) ; Shu & Nakayama (2017) ; Shi & Yu (2018) ) require the reconstruction of the entire embedding matrix or additional computations, when used at the output-layer. Thus, the model either uses the same amount of memory as the uncompressed model or requires additional computation cost. This makes linear SVD based techniques more desirable for running models on edge-devices.

In this paper, we introduce Distilled Embedding, a non-linear matrix factorization method, which outperforms the current state-of-the-art methods. Our method, first compresses the vocabulary-space to a much smaller dimension compared to the original hidden-dimension, then applies a non-linear activation function, before recovering the original embedding-dimension. Additionally, we also introduce an embedding distillation method, which is similar to Knowledge Distillation (Hinton et al., 2015) but we apply it to distill knowledge from a pre-trained embedding matrix and use an L2 loss instead of cross-entropy loss. To summarize our contributions are:

• We demonstrate that at the same compression rate our method outperforms existing stateof-the-art methods. • Our proposed method is much simpler than the current state-of-the-methods, with only a single hyper-parameter controlling the compression rate. • We introduce an embedding distillation method, which out-performs standard machine translation cross-entropy loss. • Unlike the current state-of-the-art systems, we compress the embedding matrix with parameter sharing between the first and third components. We perform an exhaustive comparison of current state-of-the-art models in this setting. • Our model uses a similar factorization as SVD, with the addition of a non-linearity, thus it is able to leverage benefits of linear SVD based methods discussed above. • We also compare all the models against a simple low-rank matrix factorization using Singular Value Decomposition (SVD) and demonstrate that it can compete with more complex algorithms.

2. RELATED WORK:

We can model the problem of compressing the embedding matrix as a matrix factorization problem. There is a considerable amount of work done in this field, some of the popular works broadly belong to the domain of low-rank factorization (Singular Value Decomposition (SVD); Srebro & Jaakkola (2003) ; Mnih & Salakhutdinov (2008) , product quantization (Jegou et al., 2010) and tensor decomposition (De Lathauwer et al., 2000) . A number of prior works in embedding compression are influenced by these fields and have been applied to various NLP problems. In this Section, we will discuss some of the significant works across different NLP problems.

Low-rank Factorization Low-rank approximation of weight matrices, using SVD, is a natural way to compress deep learning based NLP models. Sainath et al. (2013) apply this to a convolutional neural network for language modeling and acoustic modeling. Winata et al. (2018) use SVD on all the weight matrices of an LSTM and demonstrate competitive results on question-answering, language modeling and text-entailment. Acharya et al. (2019) use low-rank matrix factorization for word-embedding layer during training to compress a classification model. However, they do not study the effects of applying a non-linear function before reconstructing the original dimension.

GroupReduce Chen et al. ( 2018) apply weighted low-rank approximation to the embedding matrix of an LSTM. They first create a many-to-one mapping of all the words in the vocabulary into g groups, this initial mapping is done based upon frequency. For each group g they apply weighted SVD to obtain a lower rank estimation, the rank is determined by setting the minimum rank and linearly increasing it based upon average frequency. Finally, they update the groups by minimizing the reconstruction error from the weighted SVD approximation. They demonstrate strong results on language modeling and machine translation compared to simple SVD. In their models they use different embedding matrices for input and softmax layers and apply different compression to each.

Product Quantization Jegou et al. ( 2010) introduced product quantization for compressing high dimensional vectors, by uniformly partitioning them into subvectors and quantizing each subvector using K-means clustering technique. Basically, product quantization approach assumes that the subvectors share some underlying properties which can be used to group similar subvectors together and unify their representation. That being said, this approach breaks the original matrix into a set of codebooks coming from the center of the clusters in different partitions together with a separate index matrix which refers to the index of the clusters for each subvector. Shi & Yu (2018) applied product quantization to a language model and were able to show better perplexity scores. Shu & Nakayama (2017) extended this technique by first representing the product quantization as a matrix factorization problem, and then learning the quantization matrix in an end-to-end trainable neural network. Li et al. (2018) implement product quantization through randomly sharing parameters in the embedding matrix, and show good results on perplexity for an LSTM based language model.

Tensor Decomposition De Lathauwer et al. (2000) introduced multilinear SVD, which is a generalization of SVD for higher order tensors. Oseledets (2011) introduced an efficient algorithm Tensor Train (TT) for multilinear SVD Tensor. Novikov et al. (2015) applied the Tensor Train decomposition on fully connected layers of deep neural networks. Khrulkov et al. (2019) applied Tensor Train algorithm to the input embedding layer on different NLP problems like language modeling, machine translation and sentiment analysis. They demonstrate high compression rate with little loss of performance. However, they compress only the input embedding and not the softmax layer for language modeling and machine translation.

Knowledge Distillation Knowledge distillation has been studied in model compression where knowledge of a large cumbersome model is transferred to a small model for easy deployment. Several studies have been studied on the knowledge transfer technique (Hinton et al. (2015) ; Romero et al. (2015) ). In this paper, we propose a embedding factorization of word-embedding matrix using knowledge distillation to mimic the pre-trained word-embedding representation.

3 METHODOLOGY: DISTILLED EMBEDDING

3.1. FUNNELING DECOMPOSITION AND EMBEDDING DISTILLATION:

We present an overview of our proposed method in Figure 1 . Given an embedding matrix E ∈ R |V|×d , we can decompose it into three matrices (Equation 1), using the SVD algorithm

EQUATION

)

where |V| is the vocabulary size and d is the embedding dimension. Σ is a diagonal matrix containing the singular values, and matrices U and V represent the left and right singular vectors of the embedding matrix respectively. We can obtain the reduced form of the embedding matrix by only keeping r (< d) largest singular values out of d. Our proposed approach in this work, is to apply a non-linear transformation on the matrix U, before reconstructing the original embedding dimension using V (see Figure 1a ), as shown in Equation 3,

EQUATION

EQUATION

We use the ReLU as our non-linear function f (.) throughout this paper. We train a sequence to sequence model (Sutskever et al. (2014) ; Vaswani et al. ( 2017)) with tied input and output embedding (i.e. the output embedding is the transpose of the input embedding matrix Ẽout = ẼT = V d×r [f (U |V|×r ] T . We train our model end-to-end by replacing the embedding function with Equation 3. The matrix U and V are trainable parameters, and for the output layer we use ẼT , with the parameter sharing. We train on two losses. The standard cross entropy loss defined as:

EQUATION

)

where M is the sequence length, y i is the one-hot representation for the i th label and p i is the softmax probability of the i th term generated by the decoder.

In addition to the cross-entropy loss, we introduce a novel embedding reconstruction loss (Equation 5), which we refer to as embedding distillation as we distill information from the pre-trained embedding into our model,

EQUATION

where e i and êi are the embedding vectors corresponding to the i th word in the original embedding matrix E and the reconstructed embedding matrix Ê respectively and u i refers to the i th row of the matrix U. We use Equation 6 as our final loss function.

L total = αL recon + (1 -α)L ce (6)

where α ∈ [0, 1] is a hyper-parameter, which controls the trade-off between reconstruction and cross-entropy loss. L recon acts as the knowledge distillation loss by which we try to distill information from the original pre-trained embedding layer as a teacher to the funneling decomposed embedding layer as a student. The training process of our Distilled Embedding method is summarized in Algorithm 1.

Algorithm 1 Distilled Embedding:

Step 1) Pre-training the Embedding Matrix Pre-train the sequence to sequence model with the full embedding matrix for better initialization.

Step 2) Initializing the Weights of Funneling Decomposition Layer We extract the trained embedding matrix E from Step 1 and train our decomposed matrices U and V on reconstruction loss defined in Equation 5, as shown in Figure 1a .

Step 3) Embedding Distillation The pre-trained funneling decomposition layer is plugged into the model (replacing the original embedding matrix E) and the entire model is trained based on Equation 6.

4. EXPERIMENTAL SETUP:



4.1. DATASETS AND EVALUATION:

We test our proposed method on machine translation which is a fundamental problem in NLP and challenging for embedding compression since we typically have at least two dictionaires and an input and output embedding. Comparitively language modeling uses a single dictionary and classification tasks such as sentiment analysis don't have an output embedding.

We present results on translating three language pairs: WMT English to French (En-Fr), WMT English to German (En-De) and IWSLT Portuguese to English (Pt-En). We decided that these pairs are good representatives of high-resource, medium-resource and low-resource language pairs. WMT En-Fr is based on WMT14 training data which contain 36M sentence pairs. We used Senten-cePiece (Kudo & Richardson (2018) ) to extract a shared vocabulary of 32k subwords. We validate on newstest2013 and test on newstest2014. For WMT English to German (En-De), we use the same setup as Vaswani et al. (2017) . The dataset is based on WMT16 training data and contains about 4.5M pairs. We use a shared vocabulary of 37k subwords extracted using SentencePiece.

For the IWSLT Portuguese to English (Pt-En) dataset, we replicate the setup of Tan et al. (2019) for training individual models. Specifically, the dataset contains about 167k training pairs. We used a shared vocabulary of 32k subwords extracted with SentencePiece.

For all language pairs, we measure case-sensitive BLEU score (Papineni et al. (2002) ) using Sacre-BLEU1 (Post (2018) ). In addition, we save a checkpoint every hour for the WMT En-Fr and WMT En-De language pairs and every 5 minutes for the IWSLT Pt-En due to the smaller size of the dataset. We use the last checkpoint which resulted in the highest validation BLEU and average the last five checkpoints based on this. We use beam search with a beam width of 4 for all language pairs.

4.2. EXPERIMENT DETAILS:

Hyper-Parameters For WMT En-Fr and WMT En-De we use the same configuration as Transformer Base which was proposed by Vaswani et al. (2017) . Specifically, the model hidden size d model is set to 512, the feed-forward hidden size d ff is set to 2048 and the number of layers for the encoder and the decoder was set to 6. For the IWSLT Pt-En, we use Transformer Small configuration. Specifically, the model hidden-size d model is set to 256, the feed-forward hidden size d ff is set to 1024 and the number of layers for the encoder and the decoder was set to 2. For Transformer Small, the dropout configuration was set the same as Transformer Base. All models are optimized using Adam (Kingma & Ba (2015) ) and the same learning rate schedule as proposed by Vaswani et al. (2017) . We use label smoothing with 0.1 weight for the uniform prior distribution over the vocabulary (Szegedy et al. (2015) ; Pereyra et al. (2017) ). Additionally, we set the value α of Equation 6to 0.01.

Hardware Details:

We train the WMT models on 8 NVIDIA V100 GPUs and the IWSLT models on a single NVIDIA V100 GPU. approximately 6000 source tokens and 6000 target tokens for each GPU worker. All experiments were run using the TensorFlow framework2 .

5. RESULTS:



5.1. MACHINE TRANSLATION:

We present BLEU score for our method and compare it with SVD, GroupReduce (Chen et al. ( 2018)), Structured Emedding (Shi & Yu (2018) ), Tensor Train (Khrulkov et al. (2019) ) and a smaller transformer network with the same number of parameters. We learn a decomposition for all the methods except Tensor Train since it was pointed out in Khrulkov et al. (2019) that there is no difference in performance between random initialization and tensor train learnt initialization. Once initialized we plug the decomposed embedding and fine-tune till convergence. None of the weights are frozen during fine-tuning.

Table 1 presents the results on English-French translation. The hyper-parameters for tuning the competing methods are presented in Section A.1 of the Appendix. We see that on this task our method along with Structured Embedding performs the best. Group Reduce is the next, and SVD performs better than Tensor Train, showing that SVD is a strong baseline, when fine-tuned till convergence. We also compare against a 2 layer neural network (NN) with the same parameterization as distilled embedding which has not been initialized offline. The results show that initializing the NN with the embedding weights is important.

On English-German translation, as seen in Table 2 , our method outperforms all other methods. The smaller transformer network does well and is only surpassed by GroupReduce amongst the competing methods. SVD again performs better than Tensor Train.

We present the Portuguese-English translation results on Table 3 . This task presents a problem where the embedding matrix constitutes the majority of the parameters of the neural network. The embedding dimension is smaller (256) compared to the other two tasks but embedding compression yields a BLEU score increase in all methods except Structured Embedding. This is due to a regularization effect from the compression. Our model again achieves the highest BLEU score.

On these three experiments we demonstrate that our funneling decomposition method with embedding distillation consistently yields higher BLEU scores compared to the competing methods.

5.2. ABLATION STUDY:

We present different experiments to demonstrate the effect of 1) Model Initialization, 2) Embedding Distillation, 3) Fine-tuning strategies, 4) Compression capability and 5) Extension and generality of our method. 

Initialization:

We do an ablation study on all the three language pairs defined in Section 4.1, to conclude, if random initialization is better than model based initialization. We conclude that model based initialization, consistently performs better (Table 4 ).

Embedding Distillation Table 5 presents different compression rates on the Pt-En task, and embedding distillation performs better across all of them. In Table 4 , we see that across all language pairs when we initialize our model using weights from the funneling decomposition, we improve when using Embedding Distillation during finetuning. We performed embedding distillation with random initialization only on the smaller Pt-En dataset and observed that Embedding Distillation improves BLEU score. drop in BLEU score, thus, we need to re-train all the weights. The use of a non-linearity and adding embedding distillation also improves BLEU score after finetuning.

Compression Rate We demonstrate in:

Extension We experimented with applying two key lessons from our method, namely, using a non-linear function and embedding distillation, to a model initialized with group partitions of the GroupReduce method (Chen et al., 2018) , we refer to this method as GroupFunneling. Table 7 shows that, GroupFunneling achieves a higher BLEU score on Pt-En compared to GroupReduce.

6. DISCUSSION:

Importance of Non-linearity We postulate that only a subset of word vector dimensions, explains most of the variance, for most word vectors in the embedding matrix. Thus, using ReLU activation might help in regularizing the less important dimensions for a given word vector.

Importance Reconstruction Loss We propose that the embedding reconstruction might suffer from adding the ReLU activation function. Thus, adding a loss for embedding reconstruction helps in grounding the embedding and not loose a lot of information. Thus, the amount of regularization is controlled by the hyper-parameter α. Our intuition is partly justified by results shown in Table A .2, as reconstruction loss performs worse without the ReLU activation function.

7. CONCLUSION AND FUTURE WORK:

In this paper we proposed Distilled Embedding, a low-rank matrix decomposition with non-linearity in the bottleneck layer for a shared word-embedding and vocabulary projection matrix. We also introduce knowledge distillation of the embedding during fine-tuning using the full embedding matrix as the teacher and the decomposed embedding as the student. We compared our proposed approach with state-of-the-art methods for compressing word-embedding matrix. We did extensive experiments using three different sizes of datasets and showed that our approach outperforms the state-ofthe art methods on the challenging task of machine translation. For future work, we will apply our approach to compress feed-forward and multi-head attention layers of the transformer network.

:

https://github.com/mjpost/sacreBLEU

https://www.tensorflow.org/

Figures and Tables:

Figure 1: Funneling Decomposition method to compress the shared embedding matrix of a transformer based sequence to sequence model.

Each training batch contained a set of sentence pairs containing Machine translation accuracy in terms of BLEU for WMT En-Fr on newstest2014.

Table content: <table><tr><td>Model</td><td>Param</td><td>Emb. Params</td><td>Emb. Compression Rate</td><td>WMT En-Fr (BLEU)</td></tr><tr><td>Transformer Base</td><td>60M</td><td>16.3M</td><td>1.0x</td><td>38.12</td></tr><tr><td>Smaller Transformer Network (416)</td><td>46M</td><td>13.3M</td><td>1.23x</td><td>37.26</td></tr><tr><td>SVD with rank 64</td><td>46M</td><td>2.08M</td><td>7.87x</td><td>37.44</td></tr><tr><td>End-to-End NN with non-linearity</td><td>46M</td><td>2.08M</td><td>7.87x</td><td>37.23</td></tr><tr><td>GroupReduce (Chen et al. (2018))</td><td>46M</td><td>2.10M</td><td>7.79x</td><td>37.63</td></tr><tr><td>Structured Embedding (Shi &amp; Yu (2018))</td><td>46M</td><td>2.07M</td><td>7.90x</td><td>37.78</td></tr><tr><td>Tensor Train (Khrulkov et al. (2019))</td><td>46M</td><td>2.12M</td><td>7.72x</td><td>37.27</td></tr><tr><td>Distilled Embedding (Ours)</td><td>46M</td><td>2.08M</td><td>7.87x</td><td>37.78</td></tr></table>

Machine translation accuracy in terms of BLEU for WMT En-De on newstest2014.

Table content: <table><tr><td>Model</td><td>Param</td><td>Emb. Params</td><td>Emb. Compression Rate</td><td>WMT En-De (BLEU)</td></tr><tr><td>Transformer Base</td><td>63M</td><td>18.94M</td><td>1.0x</td><td>27.08</td></tr><tr><td>Smaller Transformer Network (400)</td><td>46M</td><td>14.8M</td><td>1.28x</td><td>26.72</td></tr><tr><td>SVD with rank 64</td><td>46M</td><td>2.40M</td><td>7.89x</td><td>26.32</td></tr><tr><td>End-to-End NN with non-linearity</td><td>46M</td><td>2.40M</td><td>7.89</td><td>26.14</td></tr><tr><td>GroupReduce (Chen et al. (2018))</td><td>46M</td><td>2.40M</td><td>7.88x</td><td>26.75</td></tr><tr><td>Structured Embedding (Shi &amp; Yu (2018))</td><td>46M</td><td>2.40M</td><td>7.89x</td><td>26.34</td></tr><tr><td>Tensor Train (Khrulkov et al. (2019))</td><td>46M</td><td>2.44M</td><td>7.75x</td><td>26.19</td></tr><tr><td>Distilled Embedding (Ours)</td><td>46M</td><td>2.40M</td><td>7.89x</td><td>26.97</td></tr><tr><td>Model</td><td>Param</td><td>Emb. Params</td><td>Emb. Compression Rate</td><td>IWSLT Pt-En (BLEU)</td></tr><tr><td>Transformer Small</td><td>11M</td><td>8.19M</td><td>1.0x</td><td>41.43</td></tr><tr><td>Smaller Transformer Network (136)</td><td>5M</td><td>4.35M</td><td>1.88x</td><td>40.71</td></tr><tr><td>SVD with rank 64</td><td>5M</td><td>2.06M</td><td>3.96x</td><td>42.37</td></tr><tr><td>End-to-End NN with non-linearity</td><td>5M</td><td>2.06M</td><td>3.96</td><td>42.27</td></tr><tr><td>GroupReduce (Chen et al. (2018))</td><td>5M</td><td>2.06M</td><td>3.96x</td><td>42.13</td></tr><tr><td>Structured Embedding (Shi &amp; Yu (2018))</td><td>5M</td><td>2.06M</td><td>3.97x</td><td>41.27</td></tr><tr><td>Tensor Train (Khrulkov et al. (2019))</td><td>5M</td><td>2.06M</td><td>3.96x</td><td>42.34</td></tr><tr><td>Distilled Embedding (Ours)</td><td>5M</td><td>2.06M</td><td>3.96x</td><td>42.62</td></tr></table>

Machine translation accuracy in terms of BLEU for IWSLT Pt-En.

Table content: <table/>

Table 5 that it is possible to compress the embedding up to 15.86x with only a 2% drop in BLEU score for Pt-En.Re-training Fine-tuning is an important component in our method and we demonstrate through our experiments that at convergence most of the techniques are close in performance. Table6shows that freezing embedding weights and re-training the network weights or vice versa leads to a sharp Comparison of different methods for Funneling (64), CR refers to the compression rate.

Table content: <table><tr><td colspan="2">Model Params</td><td>Emb. Params</td><td>Emb. CR</td><td colspan="2">Random Initialization No Emb.</td><td colspan="2">Model Initialization No Emb.</td></tr><tr><td/><td/><td/><td/><td>Distillation</td><td>Distillation</td><td>Distillation</td><td>Distillation</td></tr><tr><td>En-Fr</td><td>46M</td><td>2M</td><td>7.87x</td><td>37.04</td><td>-</td><td>37.54</td><td>37.78</td></tr><tr><td>En-De</td><td>46M</td><td>2M</td><td>7.89x</td><td>26.07</td><td>-</td><td>26.7</td><td>26.97</td></tr><tr><td>Pt-En</td><td>5M</td><td>2M</td><td>3.96x</td><td>42.29</td><td>42.36</td><td>42.5</td><td>42.62</td></tr></table>

GroupFunneling (i.e. GroupReduce + Funneling) on IWSLT Pt-En.

Table content: <table/>

Peer review (iclr2020):
There are many ways to reduce the memory footprint and increase speed of a neural network: weight quantisation, compression, coarse-to-fine, knowledge distillation, etc. The method proposed in this work is a specific case of knowledge distillation that focuses on the discrete-input-to-first-layer and output-layer-to-discrete-output transformations, which represent a large portion of the parameters. The authors propose to use a variant of SVD (which can be viewed as 2 linear transformation, with a middle dimension that represents an embedding), where the first transformation is linear with a ReLu, and the second is linear. By approximating the learned matrices of the model, the experiments show that using the proposed variant of SVD gives similar predictive performance compared to the original model, with a fraction of the parameters. However, it seems that the authors could have simply replaced the input by a 2-layer NN (first a linear+ReLu, then a Linear) to obtain the same parametrisation, but they could have learned the parameters in a end-to-end fashion. It is not clear to me why using a surrogate L2 loss within the model should give better predictive performance than a fully end-to-end trained neural network. Without this comparison, I do not think the proposed experiments are conclusive enough.