DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

By Minsik Cho, Keivan Alizadeh-Vahid, Saurabh N. Adya, Mohammad Rastegari. 2021

40 citations recorded by the source corpus.

Deep neural network (DNN) model compression for efficient on-device inference is becoming increasingly important to reduce memory requirements and keep user data on-device. To this end, we propose a novel differentiable k-means clustering layer (DKM) and its application to train-time weight clustering-based DNN model compression. DKM casts k-means clustering as an attention problem and enables joint optimization of the DNN parameters and clustering centroids. Unlike prior works that rely on additional regularizers and parameters, DKM-based compression keeps the original loss function and model architecture fixed. We evaluated DKM-based compression on various DNN models for computer vision and natural language processing (NLP) tasks. Our results demonstrate that DKM delivers superior compression and accuracy trade-off on ImageNet1k and GLUE benchmarks. For example, DKM-based compression can offer 74.5% top-1 ImageNet1k accuracy on ResNet50 DNN model with 3.3MB model size (29.4x model compression factor). For MobileNet-v1, which is a challenging DNN to compress, DKM delivers 63.9% top-1 ImageNet1k accuracy with 0.72 MB model size (22.4x model compression factor). This result is 6.8% higher top-1accuracy and 33% relatively smaller model size than the current state-of-the-art DNN compression algorithms. Additionally, DKM enables compression of DistilBERT model by 11.8x with minimal (1.1%) accuracy loss on GLUE NLP benchmarks.

Subfields: Computer Vision and Pattern Recognition, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

1. Ali Farhadi
2. Luke S. Zettlemoyer
3. Luc Van Gool
4. Mitchell Wortsman
5. Trevor Darrell
6. Ming-Hsuan Yang
7. Ludwig Schmidt
8. Dacheng Tao

In Efficiency & Compression

1. Luke S. Zettlemoyer
2. Tim Dettmers
3. Mitchell Wortsman
4. Ali Farhadi
5. Ludwig Schmidt
6. Ari S. Morcos
7. Dan Alistarh
8. Elias Frantar

In Optimization & Training

1. Luke S. Zettlemoyer
2. Tim Dettmers
3. Ali Farhadi
4. Ludwig Schmidt
5. Ari S. Morcos
6. Mitchell Wortsman
7. Torsten Hoefler
8. Tal Ben-Nun

Citations in the corpus

• Stable and low-precision training for large-scale vision-language models (2023, background): “Moreover, while int8 and int4 are used for inference [17, 66, 16], and there is earlier work exploring 8 bit training for convnets [61, 79, 10], these formats are not commonly used for training transformer models at scale.”