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

Deep neural network ( DNN ) model compression for efficient on-device inference becomes 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 for 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 parameters and regularizers , 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 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-1 accuracy and 33 % relatively smaller model size than the current state-of-the-art DNN compression algorithms . DKM also compressed a DistilBERT model by 11.8x with minimal ( 1.1 % ) accuracy loss on GLUE NLP benchmarks . 1 INTRODUCTION . Deep neural networks ( DNN ) have demonstrated super-human performance on many cognitive tasks ( Silver et al. , 2018 ) . While a fully-trained uncompressed DNN is commonly used for server-side inference , on-device inference is preferred to enhance user experience by reducing latency and keeping user data on-device . Many such on-device platforms are battery-powered and resource-constrained , demanding a DNN to meet the stringent resource requirements such as powerconsumption , compute budget and storage-overhead ( Wang et al. , 2019b ; Wu et al. , 2018 ) . One solution is to design a more efficient and compact DNN such as MobileNet ( Howard et al. , 2017 ) by innovating the network architecture or by leveraging Neural Architecture Search ( NAS ) methods ( Liu et al. , 2019 ; Tan et al. , 2019 ) . Another solution is to compress a model with small accuracy degradation so that it takes less storage and reduces System-on-Chip ( SoC ) memory bandwidth utilization , which can minimize power-consumption and latency . To this end , various DNN compression techniques have been proposed ( Wang et al. , 2019b ; Dong et al. , 2020 ; Park et al. , 2018 ; Rastegari et al. , 2016 ; Fan et al. , 2021 ; Stock et al. , 2020 ; Zhou et al. , 2019 ; Park et al. , 2019 ; Yu et al. , 2018 ; Polino et al. , 2018 ) . Among them , weight-clustering/sharing ( Han et al. , 2016 ; Wu et al. , 2018 ; Ullrich et al. , 2017 ; Stock et al. , 2020 ) has shown a high DNN compression ratio where weights are clustered into a few shareable weight values ( or centroids ) based on k-means clustering . Once weights are clustered , to shrink the model size , one can store indices ( 2bits , 4bits , etc . depending on the number of clusters ) with a lookup table rather than actual floating-point values . ∗equal contribution Designing a compact DNN architecture and enabling weight-clustering together could provide the best solution in terms of efficient on-device inference . However , the existing model compression approaches do not usefully compress an already-compact DNN like MobileNet , presumably because the model itself does not have significant redundancy . We conjecture that such limitation comes from the fact that weight-clustering through k-means algorithm ( both weight-cluster assignment and weight update ) has not been fully optimized with the target task . The fundamental complexity in applying k-means clustering for weight-sharing comes from the following : a ) both weights and corresponding k-means centroids are free to move ( a general k-means clustering with fixed observations is already NP-Hard ) , b ) the weight-to-cluster assignment is a discrete process which makes k-means clustering non-differentiable , preventing effective optimization . In this work , we propose a new layer without learnable parameters for differentiable k-means clustering , DKM , based on an attention mechanism ( Bahdana et al. , 2015 ) to capture the weight and cluster interactions seamlessly , and further apply it to enable train-time weight-clustering for model compression . Our major contributions include the following : • We propose a novel differentiable k-means clustering layer ( DKM ) for deep learning , which serves as a generic neural layer to develop clustering behavior on input and output . • We demonstrate that DKM can perform multi-dimensional k-means clustering efficiently and can offer a high-quality model for a given compression ratio target . • We apply DKM to compress a DNN model and demonstrate the state-of-the-art results on both computer vision and natural language models and tasks . 2 RELATED WORKS . Model compression using clustering : DeepCompression ( Han et al. , 2016 ) proposed to apply kmeans clustering for model compression . DeepCompression initially clusters the weights using kmeans algorithm . All the weights that belong to the same cluster share the same weight value which is initially the cluster centroid . In the forward-pass , the shared weight is used for each weight . In the backward-pass , the gradient for each shared weight is calculated and used to update the shared value . This approach might degrade model quality because it can not formulate weight-cluster assignment during gradient back propagation ( Yin et al. , 2019 ) . ESCQ ( Choi et al. , 2017 ; 2020 ) is optimizing the clusters to minimize the change in the loss by considering hessian . Therefore , it is to preserve the current model state , instead of searching for a fundamentally better model state for compression . HAQ ( Wang et al. , 2019b ) uses reinforcement learning to search for the optimal quantization policy on different tasks . For model compression , HAQ uses k-means clustering similar to DeepCompression yet with flexible bit-width on different layers . Our work is orthogonal to this work because the k-means clustering can be replaced with our DKM with a similar flexible configuration . `` And The Bit Goes Down '' ( Stock et al. , 2020 ) algorithm is based on Product Quantization and Knowledge Distillation . It evenly splits the weight vector of N elements into N/d contiguous d dimensional sub-vectors , and clusters the sub-vectors using weighted k-means clustering to minimize activation change from that of a teacher network . GOBO ( Zadeh et al. , 2020 ) first separates outlier weights far from the average of the weights of each layer and stores them uncompressed while clustering the other weights by an algorithm similar to k-means . Model compression using regularization : Directly incorporating k-means clustering in the training process is not straightforward ( Wu et al. , 2018 ) . Hence , ( Ullrich et al. , 2017 ) models weightclustering as Gaussian Mixture Model ( GMM ) and fits weight distribution into GMM with additional learning parameters using KL divergence ( i.e. , forcing weight distribution to follow k Gaussian distributions with a slight variance ) . ( Wu et al. , 2018 ) proposed deep k-means to enable weightclustering during re-training . By forcing the weights that have been already clustered to stay around the assigned center , the hard weight-clustering is approximated with additional parameters . Both ( Ullrich et al. , 2017 ) and ( Wu et al. , 2018 ) leverage regularization to enforce weight-clustering with additional parameters , which will interfere with the original loss target and requires additional updates for the new variables ( i.e. , singular value decomposition ( SVD ) in ( Wu et al. , 2018 ) ) . Also , relying on the modified loss can not capture the dynamic interaction between weight distributions and cluster centroids within a batch , thus requiring an additional training flow for re-training . Enhance Model compression using dropout : Quant-Noise ( Fan et al. , 2021 ) is a structured dropout which only quantizes a random subset of weights ( using any quantization technique ) and thus can improve the predictive power of a compressed model . For example , ( Fan et al. , 2021 ) showed good compression vs. accuracy trade-off on ResNet50 for ImageNet1k . Model quantization : Besides clustering and regularization methods , model quantization can also reduce the model size , and training-time quantization techniques have been developed to improve the accuracy of quantized models ( Li et al. , 2019 ; Zhao et al. , 2019 ) . EWGS ( J. Lee , 2021 ) adjusts gradients by scaling them up or down based on the Hessian approximation for each layer . PROFIT ( Park & Yoo , 2020 ) adopts an iterative process and freezes layers based on the activation instability . Efficient networks : Memory-efficient DNNs include MobileNet ( Howard et al. , 2017 ; Sandler et al. , 2018 ) , EfficientNet ( Tan & Le , 2019 ; 2021 ) and ESPNet ( Mehta et al. , 2019 ) . MobileNet-v1 ( Howard et al. , 2017 ) on ImageNet1k dataset has top-1 accuracy of 70.3 % with 16.1 MB of memory in comparison to a ResNet18 which has 69.3 % accuracy with 44.6 MB of model size . Our method can be applied to these compact networks to reduce their model sizes further . 3 ALGORITHM . 3.1 MOTIVATION . Popular weight-clustering techniques for DNN model compression ( J. Lee , 2021 ; Han et al. , 2016 ; Dong et al. , 2020 ; Stock et al. , 2020 ) are based on k-means clustering along with enhancements such as gradient scaling/approximation . Using k-means clustering , the weights are clustered and assigned to the nearest centroids which are used for forward/backward-propagation during training as illustrated in Fig . 1 ( a ) . Such conventional methods with clustering have two critical drawbacks : • The weight-to-cluster assignment in conventional approaches is not optimized through back-propagation of training loss function . • Gradients for the weights are computed in an ad-hoc fashion : the gradient of a centroid is re-purposed as the gradient of the weights assigned to the centroid . These limitations are more pronounced for the weights on the boundary such as i and j in Fig . 1 ( a ) . In the conventional approaches , i and j are assigned to the centroids C2 and C1 respectively , simply because of their marginal difference in a distance metric . However , assigning i to C0 and j to C2 could be better for the training loss as their difference in distance is so small ( Nagel et al. , 2020 ) . Such lost opportunity cost is especially higher with a smaller number of centroids ( or fewer bits for quantization ) , as each unfortunate hard assignment can degrade the training loss significantly . We overcome such limitations with DKM by interpreting weight-centroid assignment as distancebased attention optimization ( Bahdana et al. , 2015 ) as in Fig . 1 ( b ) and letting each weight interact with all the centroids . Such attention mechanism naturally cast differentiable and iterative k-means clustering into a parameter-free layer as in Fig . 2 . Therefore , during backward-propagation , attention allows a gradient of a weight to be a product of the attentions and the gradients of centroids , which in turn impact how the clustering and assignment will be done in the next batch . overall our weight assignment will align with the loss function , and can be highly effective for DNN compression .

Summary: The paper is concerned with reducing the size of deep neural networks using weight sharing. The paper proposes a new building block that performs a soft k-means algorithm where each weight is assigned a convex combination of the cluster centers. At test-time, the weights are assigned to their closest cluster center such that a real (hard) weight sharing is obtained. The method achieves state-of-the-art accuracy using various architectures on several tasks.