Source-linked AI summary
Compression of Neural Machine Translation Models via Pruning
Abigail See, Minh-Thang Luong, Christopher D. Manning
TL;DR
NMT models are often over-parameterized and too large for storage-constrained deployment. The paper compares three magnitude-based pruning schemes for NMT and finds that retraining supports extreme compression while preserving or improving performance. The experiments also reveal where redundancy is concentrated in the architecture.
Problem
Over-parameterization leaves NMT models with storage sizes of hundreds of megabytes, motivating compression for models with over 200 million parameters.
Method
The paper compares class-blind, class-uniform, and class-distribution magnitude-based pruning schemes for NMT weight classes.
Results
With retraining, the state-of-the-art NMT model can be pruned by 80% with no performance loss, while a separate model supports 90% pruning and regains original performance.
Takeaways & Limitations
Weight pruning with retraining compresses the NMT model to 20% of its size without performance loss and exposes redundancy patterns across the architecture.
Takeaways & Limitations
The method uses one pruning stage; multiple pruning-and-retraining iterations are expected to improve compression and performance.
Abstract
from arXiv · showhide
Neural Machine Translation (NMT), like many other deep learning domains, typically suffers from over-parameterization, resulting in large storage sizes. This paper examines three simple magnitude-based pruning schemes to compress NMT models, namely class-blind, class-uniform, and class-distribution, which differ in terms of how pruning thresholds are computed for the different classes of weights in the NMT architecture. We demonstrate the efficacy of weight pruning as a compression technique for a state-of-the-art NMT system. We show that an NMT model with over 200 million parameters can be pruned by 40% with very little performance loss as measured on the WMT'14 English-German translation task. This sheds light on the distribution of redundancy in the NMT architecture. Our main result is that with retraining, we can recover and even surpass the original performance with an 80%-pruned model.
1 Introduction
NMT offers strong translation capabilities but remains too large for many mobile deployments because over-parameterization increases storage demands. This paper studies magnitude-based weight pruning, comparing three schemes and showing substantial compression with limited performance loss.
- NMT is an end-to-end deep neural network for translating text between languages and has achieved state-of-the-art results across several language pairs.
- Over 200 million parameters can make state-of-the-art NMT models require hundreds of megabytes of storage, limiting use on mobile devices.The paper links over-parameterization to storage size, overfitting, and long running times, while noting that runtime is outside this paper’s scope.
- The paper investigates weight pruning as a compression method for NMT and compares class-blind, class-uniform, and class-distribution magnitude-based schemes.
- 40% pruning causes negligible performance loss, while retraining enables 80% pruning with no performance loss on the state-of-the-art NMT system.
- Pruning experiments indicate that higher layers, attention, and softmax weights are more important, whereas lower layers and embedding weights contain more redundancy.
2 Related Work
Prior neural-network pruning research largely targeted CNNs and vision, using either magnitude-based or more computationally intensive saliency methods. This work extends magnitude-based weight pruning to LSTM-based NMT and contrasts it with other compression approaches.
- OBD and OBS use Hessian-based parameter saliency, while magnitude-based pruning removes low-magnitude weights with substantially lower computational complexity.
- Iterative magnitude-based retraining pruned 75% of AlexNet parameters with small ImageNet accuracy loss and 89% with no accuracy loss.
- Neuron-pruning methods require whole zero rows or near-identical row pairs, whereas weight pruning removes weights independently.
- Other compression methods reduce matrix rank, share weights, lower numerical precision, use binary weights, or distill knowledge into a smaller network.
- This paper extends magnitude-based pruning from CNNs to LSTM architectures for NMT, addressing a setting previously lacking pruning-based compression work.
3 Our Approach
The paper applies magnitude-based pruning to a deep multi-layer LSTM NMT architecture and compares three ways of distributing pruning across weight classes. It also retrains the remaining weights while preserving the induced sparse structure.
- NMT architecture: The NMT system uses an encoder-decoder architecture with deep multi-layer LSTMs, separate source and target weights, attention, and softmax output scores.Feed-forward weights connect layers, while recurrent weights connect successive time steps.
- LSTM weights: Each LSTM block produces hidden and memory vectors from the preceding block or embedding input, with vectors of length n.The core LSTM weight matrix has size 4n × 2n and is decomposed into eight weight subgroups.
- Pruning schemes: Magnitude-based pruning removes the x% smallest-magnitude parameters, but the schemes differ in how pruning is distributed across weight classes.The three schemes are class-blind, class-uniform, and class-distribution.
- Pruning schemes: Class-blind pruning sorts all parameters together, whereas class-uniform prunes x% within every class and class-distribution uses class-specific thresholds λσ_c.In class-distribution pruning, σ_c is each class’s standard deviation and λ is chosen to reach the desired overall pruning rate.
- Scheme comparison: Class-blind pruning outperforms the other two schemes in translation quality at all tested pruning percentages.The paper attributes the comparison to experiments summarized in Figure 3.
- Retraining: Retraining continues to update surviving weights while binary masks prevent updates at pruned locations, preserving the sparse structure.The zero-and-mask implementation is simple, while sparse-matrix multiplication is identified as a possible but out-of-scope route to speed improvements.
4 Experiments
Experiments show that class-blind magnitude pruning is the strongest of the three schemes, while retraining enables substantial pruning with little or no BLEU loss. The experiments also identify where redundancy and sensitivity are concentrated in the NMT model.
- Comparing pruning schemes: Class-blind pruning outperforms class-uniform and class-distribution pruning at every tested pruning percentage.The remaining experiments therefore use only class-blind pruning.
- Comparing pruning schemes: Under class-uniform pruning, target layer 4, attention, and softmax weights disproportionately cause performance loss, whereas class-blind pruning makes class contributions more uniform.The experiments suggest that higher layers, attention, and softmax weights are especially important.
- Pruning and retraining: 40% pruning leaves translation performance mostly unaffected, despite pruning’s immediate performance loss increasing exponentially with pruning percentage.This indicates substantial redundancy and over-parameterization in NMT.
- Pruning and retraining: Retraining may improve validation performance because pruning regularizes training and can enable descent toward a better local optimum.The 50% pruned and retrained model outperforms the 90% pruned and retrained model on validation performance.
- Distribution of redundancy in NMT: At 80% pruning, all parameters for 100 uncommon source words and 1,194 uncommon target words are deleted, producing zero embeddings rather than unknown-word mappings.These words already had near-zero embeddings in the unpruned model.
5 Generalizability of our results
Tests on a smaller Vietnamese-English NMT model found results similar to the main model despite differences in scale, training data, and language pair. Pruning remained effective, especially with retraining.
- The smaller model used 133,000 Vietnamese-English sentence pairs and about 11% as many parameters as the state-of-the-art model.
- 60% of parameters could be pruned with no immediate performance loss.
- 90% pruning with retraining restored the model’s original performance.
- Class-blind pruning was the most successful scheme on the smaller model as well as the main model.
- Sparse-from-the-beginning models were less successful than models pruned and retrained afterward.
6 Future Work
Future work could improve pruning through repeated pruning-retraining cycles, exploit sparsity for computational speedups, and compare magnitude-based pruning with more principled alternatives.
- Several pruning and retraining iterations could improve compression and performance.
- Sparse matrix representations and multiplications could exploit pruned-model sparsity to speed training and runtime.
- A comparative study could revisit whether Optimal Brain Damage and Optimal Brain Surgery outperform magnitude-based pruning.
7 Conclusion
Weight pruning with retraining effectively compresses and regularizes a state-of-the-art NMT system. The approach is simple, broadly applicable, and also provides insight into redundancy within the NMT architecture.
- The model was compressed to 20% of its original size with no performance loss.
- Weight pruning with retraining was effective for both compression and regularization on a state-of-the-art NMT system.
- The absolute size of parameters was the primary criterion for choosing which weights to prune.
- The resulting pruning approach was extremely simple to implement and applicable to any neural network.
- The experiments provided insight into the distribution of redundancy in the NMT architecture.