[Paper Review] Tensorized Embedding Layers for Efficient Model Compression

Oleksii Hrinchuk, Valentin Khrulkov+3 more|arXiv (Cornell University)|Jan 30, 2019

Advanced Data Compression Techniques64 references50 citations

TL;DR

The paper introduces TT–embedding, a Tensor Train-based parameterization of embedding layers that dramatically reduces memory footprint while maintaining or improving NLP task performance, trained end-to-end within standard models.

ABSTRACT

The embedding layers transforming input words into real vectors are the key components of deep neural networks used in natural language processing. However, when the vocabulary is large, the corresponding weight matrices can be enormous, which precludes their deployment in a limited resource setting. We introduce a novel way of parametrizing embedding layers based on the Tensor Train (TT) decomposition, which allows compressing the model significantly at the cost of a negligible drop or even a slight gain in performance. We evaluate our method on a wide range of benchmarks in natural language processing and analyze the trade-off between performance and compression ratios for a wide range of architectures, from MLPs to LSTMs and Transformers.

Motivation & Objective

• Motivate embedding layer compression due to large NLP vocabulary size.
• Propose TT–embedding to replace standard embedding matrices with a tensorized representation.
• Analyze trade-offs between compression ratios and performance across NLP tasks.
• Demonstrate end-to-end trainability of TT–embedding within common architectures (MLP, LSTM, Transformer).
• Provide guidance on TT–shape and TT–rank hyperparameters for practical use.

Proposed method

• Represent embedding matrices as TT–matrices by reshaping input/output dimensions into N factors and learning TT cores.
• Compute word embeddings by contracting TT cores via slice selections and matrix multiplications in a sequence of operations.
• Initialize TT cores with a modified Glorot-like scheme ensuring controlled variance, as per equation (3).
• Train TT–embedding end-to-end with backpropagation alongside standard layers (LSTM, Transformer, etc.).
• Optionally compare TT–embedding with low-rank factorization and Tensor Ring variants; justify TT via full-rank expressivity (Theorem 1).
• Provide practical guidelines for choosing TT shapes and ranks, and discuss generalization/regularization implications.

Experimental results

Research questions

• RQ1 Can TT–embedding achieve significant compression of embedding layers without degrading task performance?
• RQ2 How do TT–embedding trade-offs between compression ratio and accuracy compare to low-rank factorization?
• RQ3 Do TT–embeddings maintain full expressivity in the softmax bottleneck setting and avoid rank collapse?
• RQ4 What is the impact of TT-shape/rank choices on training efficiency and generalization across NLP tasks?
• RQ5 Is TT–embedding beneficial for large-vocabulary tasks like CTR prediction and language modeling?

Key findings

• TT–embedding achieves substantial compression (e.g., 441× on IMDB, 15× on WMT En-De, 3.8× on WikiText-103) with equal or improved accuracy on some tasks.
• In CN experiments, TT–embedding can outperform low-rank embeddings in compression-accuracy trade-offs across sentiment analysis, NMT, and language modeling.
• The method remains trainable end-to-end within standard architectures and can be integrated with minimal changes to existing pipelines.
• Theoretical analysis shows TT–embedding preserves full rank in the resultant embedding matrix for almost all TT configurations, mitigating softmax bottleneck concerns.
• CTR experiments demonstrate compression up to 2011× with negligible loss in performance and total model size below 20 MB.