Source-linked AI summary
GroupReduce: Block-Wise Low-Rank Approximation for Neural Language Model Shrinking
Patrick H. Chen, Si Si, Yang Li, Ciprian Chelba, Cho-jui Hsieh
TL;DR
Large-vocabulary language models devote much of their memory to embedding and softmax matrices, making them important compression targets. GroupReduce applies frequency-based block-wise weighted low-rank approximation; on OBW, combining it with quantization achieves 26 times compression for those matrices.
Problem
Embedding and softmax matrices can dominate language-model memory, while traditional compression methods have struggled to compress word embeddings effectively.
Source: §1 · p. 2, §2.2 · p. 3Method
GroupReduce groups words by frequency and iteratively applies weighted low-rank approximation to each block.
Source: §1 · p. 2Results
On OBW, GroupReduce combined with quantization achieves 26 times compression for both embedding and softmax matrices, saving more than 5GB of memory.
Source: §5 · p. 10Takeaways & Limitations
The method provides practical benefits when deploying neural language models on memory-constrained devices.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Model compression is essential for serving large deep neural nets on devices with limited resources or applications that require real-time responses. As a case study, a state-of-the-art neural language model usually consists of one or more recurrent layers sandwiched between an embedding layer used for representing input tokens and a softmax layer for generating output tokens. For problems with a very large vocabulary size, the embedding and the softmax matrices can account for more than half of the model size. For instance, the bigLSTM model achieves state-of- the-art performance on the One-Billion-Word (OBW) dataset with around 800k vocabulary, and its word embedding and softmax matrices use more than 6GBytes space, and are responsible for over 90% of the model parameters. In this paper, we propose GroupReduce, a novel compression method for neural language models, based on vocabulary-partition (block) based low-rank matrix approximation and the inherent frequency distribution of tokens (the power-law distribution of words). The experimental results show our method can significantly outperform traditional compression methods such as low-rank approximation and pruning. On the OBW dataset, our method achieved 6.6 times compression rate for the embedding and softmax matrices, and when combined with quantization, our method can achieve 26 times compression rate, which translates to a factor of 12.8 times compression for the entire model with very little degradation in perplexity.
1 Introduction
Large-vocabulary language models devote much of their memory to embedding and softmax matrices, making their compression important. GroupReduce uses frequency-based block-wise weighted low-rank approximation and outperforms standard approaches, with strong compression results on OBW and DE-EN NMT.
1 Introduction
- Embedding and softmax matrices grow with vocabulary size and account for over 90% of the memory footprint in the One Billion Word task.They also use around 80% of memory for DE-EN NMT.Source: §1 · p. 2
- Standard compression methods such as quantization have struggled on word embeddings; one reported method achieved less than 3 times compression for language models.The paper focuses on embedding and softmax matrices, whose token columns have frequency-dependent usage.Source: §1 · p. 2
- GroupReduce groups words by frequency and iteratively refines the clusters using weighted low-rank approximation for each block.This approach is motivated by token-column usage frequencies following Zipf’s law.Source: §1 · p. 2
- On OBW, GroupReduce compresses embedding and softmax matrices 6.6 times, rising to more than 26 times with quantization while maintaining similar perplexity.For DE-EN NMT, it achieves 10 times compression without much performance degradation, rising to 24 times with quantization.Source: §1 · p. 2
2 Related Work
Prior work applies low-rank methods, pruning, and quantization to CNNs and recurrent models. For language-model embeddings, reported compression results remain limited, motivating methods that exploit language-specific structure.
2.1 Model Compression for CNN
2.2 Model Compression for RNN/LSTM
- RNN work applied quantization and retraining [8], matrix or tensor factorization [23], and pruning [18] [16] to recurrent models.Some factorization evaluations addressed image and music classification rather than word embedding matrices.Source: §2.2 · p. 3
- For language-model embeddings, prior quantization achieved less than 3 times compression on PTB without performance loss [8], while pruning reached 87% sparsity with more than 5% performance loss [16].The paper characterizes achieving 4 times compression without sacrificing performance as extremely challenging.Source: §2.2 · p. 3
3 Proposed Algorithms
GroupReduce compresses the embedding and softmax matrices by exploiting word-frequency structure in weighted, block-wise low-rank approximation. On PTB-Small, the reported strategies each improve perplexity at the same compression rate.
3 Proposed Algorithms
- Vanilla SVD can reconstruct embedding matrices poorly because their eigenvalues remain large, and it ignores the word-specific structure of matrix rows.Figure 1 illustrates the eigenvalues and low-rank reconstruction error for PTB embedding matrices.Source: §3 · p. 4, Figure 1 · p. 4
3.1 The Word Frequency Matters
- Word frequencies follow a power law, so the method prioritizes better approximation of frequent words through frequency-weighted error.The weighted objective penalizes approximation error more for higher-frequency words.Source: §3.1 · p. 4, §3.2 · p. 5
3.2 Improved Low-rank Approximation by Exploiting Frequency
- Weighted low-rank approximation is solved by taking an SVD of QA, where Q contains square-root word-frequency weights, then transforming the factors back.This provides a direct solution to the row-weighted objective when all frequencies are nonzero.Source: §3.2 · p. 5
- GroupReduce partitions words into blocks, approximates each block separately, and jointly refines cluster assignments and block factors to reduce weighted reconstruction error.The clustering is initialized by frequency; alternating minimization updates the factors and assignments.Source: §3.2 · pp. 5–6
- At the same compression rate, every added strategy improves PTB-Small perplexity, and the authors report all are needed for the final algorithm.Table 2 evaluates the strategies with 10 blocks and 5× compression.Source: Table 2 · p. 6, §3.2 · p. 6
- During refinement, words move between clusters when doing so reduces reconstruction error, balancing frequency grouping with approximation quality.The algorithm tests each word against cluster bases to estimate the error reduction from a move.Source: §3.2 · pp. 5–6
- Block ranks scale with average word frequency, assigning more rank to high-frequency blocks so their words receive better approximations.The rank for the least-frequent block is set by the budget requirement.Source: §3.2 · p. 6
- The block representation uses O(Nk + ckD) memory, where N is vocabulary size, c is the number of clusters, and k is their average rank.The method targets embedding matrices, which account for more than 75% of overall memory in three of the four reported setups.Source: §3 · p. 3, §3.2 · p. 6
4 Experiments
Experiments evaluate GroupReduce on language-modeling and translation benchmarks, comparing it with standard compression methods and testing quantization and cluster-count choices. GroupReduce achieves over 10 times overall compression on both tasks, while OBW-bigLSTM requires a separate, more conservative treatment of its LSTM cells.
4.1 Datasets and Pretrained Models
- GroupReduce initializes frequency-based word clusters, assigns ranks using cluster-average frequency, and iteratively updates clusters using reconstruction error and weighted low-rank approximations.Words with a lowest-error cluster different from their original cluster enter a candidate set for reassignment.Source: §4.1 · p. 7
- GroupReduce compresses embedding and softmax matrices in four 2-layer LSTM models evaluated on PTB, OBW, and IWSLT 2014 German-to-English, using task-specific performance-drop criteria to select model sizes.PTB models were trained from scratch, while the OBW and NMT models used existing checkpoints.Source: §4.1 · p. 7
4.2 Comparison with Low-Rank and Pruning
- GroupReduce compresses input embedding and softmax layers 5–10 times without much accuracy loss, including 6.6 times compression on OBW that saves more than 5 GB.The comparison is against low-rank approximation and pruning; Table 3 reports embedding compression results across three datasets.Source: Table 3 · p. 8, §4.2 · p. 8
4.3 Comparison with Quantization
- After retraining, quantized GroupReduce achieves at least 24 times compression for both input embedding and softmax matrices on OBW.The method first applies GroupReduce, then 4- or 8-bit quantization to the resulting low-rank matrices.Source: §4.3 · p. 8, Table 4 · p. 9
4.4 Overall Compression
- Over 10 times overall compression was achieved on language-modeling and neural machine-translation tasks; OBW-bigLSTM's LSTM cells received only 16-bit quantization because even slight compression significantly reduced performance.The reported overall compression rates are summarized in Table 5.Source: Table 5 · p. 9, §4.4 · p. 9
4.5 Selection of the Number of Clusters
- The cluster-count hyperparameter is robust on PTB-Large at 6.6 times compression, with 5 clusters used for PTB and IWSLT and 20 for OBW.OBW uses more clusters because its vocabulary is larger.Source: §4.5 · p. 9
4.6 Comparison with Deep Compositional Coding
- GroupReduce compresses the input embedding twice as much as deep compositional coding in the evaluated setups and also applies to the softmax layer.The reported comparison covers NMT:DE-EN and PTB-Large; Table 7 reports input-embedding compression rates rather than overall model size.Source: §4.6 · pp. 9–10, Table 7 · p. 10
5 Conclusion
GroupReduce uses word statistics to form block-wise low-rank approximations of embedding and softmax layers. It achieves strong compression, including 26 times on OBW with quantization, and is intended to benefit deployment on memory-constrained devices.
5 Conclusion
- On OBW, GroupReduce with quantization achieves 26 times compression for both embedding and softmax matrices, saving more than 5GB.Source: §5 · p. 10
- GroupReduce achieves at least 6.6 times compression without losing prediction accuracy.Source: §5 · p. 10
- The method forms block-wise low-rank approximations of embedding and softmax layers using statistical properties of words.Source: §5 · p. 10
- GroupReduce significantly outperforms traditional compression methods such as low-rank approximation and pruning.Source: §5 · p. 10
- The authors identify end-to-end training of the block low-rank parameterization as future work.Source: §5 · p. 10