DKM: D IFFERENTIABLE k -M EANS C LUSTERING L AYER FOR N EURAL N ETWORK C OMPRESSION
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep neural network (DNN) model compression for efﬁcient 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 ﬁxed.
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.4 x 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.
I NTRODUCTION
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 efﬁcient 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 regression 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 been shown to deliver a high DNN compression ratio where weights are clustered into a few shareable weight values (or centroids) based on well-known 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 ﬂoating-point values.
Designing a compact DNN architecture and enabling weight-clustering together could provide the best solution in terms of efﬁcient 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 signiﬁcant 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 ﬁxed 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 activation layer 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 activation layer.
• We demonstrate that multi-dimensional k -means clustering 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/tasks.
R ELATED W ORKS
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-propagation, the shared weight is used for each weight.
In the backward-propagation, the gradient for each shared weight is calculated and used to update the shared value.
This approach might degrade model quality because it cannot formulate weightcluster assignment during gradient back propagation (Yin et al., 2019).
ESCQ (Choi et al., 2017) extends DeepCompression by proposing weighted clustering based on Hessian-weighted distortion measure to minimize the impact on the loss function due to 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 (Han et al., 2016)
yet with ﬂexible 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 ﬂexible conﬁguration. "
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 & Moshovos, 2020)
ﬁrst 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 ﬁts 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 modiﬁed loss cannot capture the dynamic interaction between weight distributions and cluster centroids within a batch, thus requiring an additional training ﬂow 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, when combined with (Stock et al., 2020), (Fan et al., 2021) showed good compression vs. accuracy trade-off on
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.
Efﬁcient networks : Efﬁcient architectures such MobileNet (Howard et al., 2017; Sandler et al., 2018), Efﬁcient Net (Tan & Le, 2019; 2021) and ESPNet (Mehta et al., 2019) are designed to be memory-efﬁcient.
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 while delivering competitive accuracies.
A LGORITHM 3. M OTIVATION
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 approaches based on clustering have two critical drawbacks:
• The weight-to-cluster assignment in conventional approaches are based on a distance metric without being aligned with the 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.(a).
In the conventional approaches, i and j are assigned to the centroids C 2 and C 1 respectively, simply because of their marginal difference in a distance metric.
However, assigning i to C 0 and j to C 2 could be better for the training loss as their difference in distance is so small (Nagel et al.,
[Figure]
Since DKM is a differentiable activation layer, backward-propagation will run through the iterative loop and the gradients for the weights will be computed against the task loss function.
Such lost opportunity cost is especially higher with a smaller number of centroids (or fewer bits for quantization), as each unfortunate assignment can degrade the training loss signiﬁcantly, yet without recoverable solutions.
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.
Note that such attention mechanism naturally allows differentiable and iterative k -means clustering as in Fig.
2 to be an activation layer, and can be highly effective in weight-clustering for DNN compression.
3. D IFFERENTIABLE K- MEANS CLUSTERING LAYER FOR W EIGHT - CLUSTERING
DKM can perform a differentiable train-time weight-clustering iteratively for k clusters as shown in Fig.
2 for the DNN model compression purpose.
Let C ∈ R
k be a vector of cluster centers and W ∈ R N be a vector of the weights, and then DKM performs as follows:
• In the ﬁrst iteration C can be initialized either by randomly selected k weights from W or k -means++.
For all subsequent iterations, the last known C from the previous batch is used to accelerate the clustering convergence.
• A distance matrix D is computed for every pair between a weight w i and a centroid c j using a differentiable metric f
(i.e., the Euclidean distance) using d ij = − f ( w i , c j ) .
• We apply softmax with a temperature τ on each row of D to obtain attention matrix A where a i,j = exp ( di,jτ ) P k exp ( di,kτ ) represents the attention from w i and
c j . •
Then, we obtain a new centroid ˜C by gathering all the attentions from W for each centroid by computing ˜ c
j = P i a
i,j w i P i a i,j . • We repeat this process till ˜C is close enough to C at which point k -means has converged, and we compute AC to get ˜W for forward-propagation (as in Fig. 2).
The iterative process will be dynamically executed imperatively in PyTorch (Paszke et al., 2019) and Tensorﬂow-Eager (Agrawal et al., 2019) and is differentiable for backward-propagation, as ˜W is based on the attention between weights and centroids.
DKM uses soft weight-cluster assignment which could be hardened in order to impose weight-clustering constraints.
The level of hardness can be controlled by the temperature τ in the softmax operation.
During inference we use the last attention matrix (i.e., A in Fig. 2) from a DKM layer to snap each weight to the closest centroid of the layer and ﬁnalize weight-clustering as in prior arts (i.e., no more attention), but such assignment is expected to be tightly aligned with the loss function, as the weights have been annealed by shuttling among centroids.
Using DKM for model compression allows a weight to change its cluster assignment during traintime, but eventually encourages it to settle with the best one w.r.t the task loss.
Optimizing both weights and clusters simultaneously and channeling the loss directly to the weight-cluster assignment is based on the attention mechanism.
Since each step in DKM is without additional learnable parameters and transparent to a model and loss function, we can reuse the existing training ﬂow and hyper-parameters.
Therefore, the critical differences between DKM-based compression and the prior works can be summarized as follows:
• Instead of hard weight-cluster assignment and approximated gradient (Han et al., 2016; Wang et al., 2019b; J. Lee, 2021; Stock et al., 2020), DKM uses ﬂexible and differentiable attention-based weight-clustering and computes gradients w.r.t the task loss without approximation.
• Instead of modifying the loss function with regularizers to enforce clustering (Ullrich et al., 2017; Wu et al., 2018), DKM (as a differentiable activation) can be inserted into forward functions, making the optimization fully aligned with the task objective (i
e., no interference in loss).
• DKM requires no additional learnable parameters (Ullrich et al., 2017; J. Lee, 2021), thus making the training ﬂow simple.
For example, DKM-base approach does not need to substitute a convolution layer with a specialized version with additional learning parameters.
• DKM requires no additional computation such as Hessian trace (Dong et al., 2020; Choi et al., 2017) or SVD (J. Lee, 2021; Wu et al., 2018) for gradient approximation, because DKM uses a differentiable process.
• DKM-based compression does not need a complex training ﬂow such as freezing/progress (Park & Yoo, 2020) or distillation (Stock et al., 2020), keeping training ﬂow simple and unchanged.
3. M ULTI -D IMENSIONAL DKM
DKM can be naturally extended into multi-dimensional weight-clustering (Stock et al., 2020) due to its simplicity, and is highly effective due to its differentiability.
We split N elements of weights into
Nd contiguous d dimensional sub-vectors and cluster the sub-vectors ( W ∈ R Nd ∗ d ).
For example, we simply ﬂatten all the convolutional kernels into a ( Nd , d ) matrix across both kernel and channel boundaries and apply multi-dimensional DKM to the matrix for clustering in our implementation.
Accordingly, the cluster centroids will become d -dimensional as well ( C ∈ R k ∗ d ) and the metric calculation is done in the d -dimensional space.
With the multi-dimensional scheme, the effective bit-per-weight becomes bd for b -bit/ d -dim clustering.
Such multi-dimensional clustering could be ineffective for conventional methods (i.e., DNN training not converging)
(Stock et al., 2020; Wang et al., 2019b; J. Lee, 2021), as now a weight might be on the boundary to multiple centroids, and the chance of making wrong decisions grows exponentially with the number of centroids.
For example, there are only two centroids for 1bit/1dim clustering, while there are 16 centroids in 4bit/4dim clustering, although both have the same effective bit-perweight .
Intuitively, however, DKM can work well with such multi-dimensional conﬁgurations as
DKM naturally optimizes the assignment w.r.t the task objective and can even recover from a wrong assignment decision over the training-time optimization process.
For a given sufﬁciently large N ≫ d in 32 bits, the compression ratio is 32 db and its entropy is b , assuming all the clusters are at the same size.
Since higher entropy indicates more ﬂexibility in the weight distribution and better model quality (Park et al., 2017), it can be expected that increasing b and d at the same ratio as much as possible can improve the model quality for a given target compression ratio (see Section 4.1 for results).
The storage complexity of a DKM layer with N parameters is O ( r Nd 2 b ) where r is the number of iterations per Fig.
2, as all the intermediate results such as D and A at each iteration need to be [Figure] kept for backward-propagation.
Therefore, when there are a large number of weights and the target precision is high (i.e., more clusters), the training ﬂow could suffer from out-of-memory error.
One solution to such memory overhead is to use a sparse representation for A by keeping top- k centroids for each weight.
E XPERIMENTAL R ESULTS
We compared our DKM-based compression with multiple state-of-the-art quantization or compression schemes on various computer vision and natural language tasks and models.
To study the trade-off between model compression and accuracy regression, we disabled activation quantization in every experiment for all approaches.
All our experiments with DKM were done on two x86 Linux machine with eight NVIDIA V100 GPUs each in a public cloud infrastructure, and activation layers stayed unchanged, as our main goal is the model size reduction through compression.
We used a
SGD optimizer with momentum 0.9, and ﬁxed the learning rate at 0.008 (without individual hyperparameter tuning) for all the experiments for DKM.
Each compression scheme starts with publicly available pre-trained models.
4. I MAGE N ET 1 K
We compared our DKM-based compression with DeepCompression (or DC )
(Han et al., 2016),
HAQ (Wang et al., 2019b), and "
And The Bit Goes Down" (or ATB ) (Stock et al., 2020) combined with Quantization-noise (Fan et al., 2021) which are all designed for model compression, and summarized results in Table 1.
We set the mini-batch size 128 per GPU (i.e., global mini-batch size of 2048) and ran for 200 epochs for all DKM cases.
Since the public ATB implementation does not include MobileNet-v1/v2 cases (Howard et al., 2017; Sandler et al., 2018), we added the support for these two by following the paper and the existing ResNet18/
(He et al., 2016) implementations.
Instead of using a complex RL technique as in HAQ (Wang et al., 2019b), for DKM experiments, we ﬁxed conﬁgurations for all the convolution layers (noted as cv) and the last fully connected layer (noted as fc), except that we applied 8 bit clustering to a layer with fewer than 10,000 parameters.
Even using such a simple method, DKM offers a Pareto superiority to other schemes as visualized in Fig.
For ResNet50 and MobileNet-v1/v2, DKM delivered compression conﬁgurations that yielded both better accuracy and higher compression ratio than the prior arts.
For ResNet18, DKM was able to make a smooth trade-off on accuracy vs. compression, and ﬁnd Pareto superior conﬁgurations to ATB as in Fig. 3: DKM can get 64.5% Top-1 accuracy with 1MB which is superior to ATB-large, 66.8% Top-1 accuracy with 1.49 MB which is superior to ATB-small, and 65.8% Top-1 accuracy with 1.09 MB as a superior balance point.
For MobileNet-v1/v2, ATB failed to converge, but DKM outperforms DC and HAQ in terms of both accuracy and size at the same time. For ATB cases,
[Table]
[Table] adding quantization noise improves the model accuracy (Fan et al., 2021). Since the quantization noise technique is complementary to any compression algorithm, we believe that it would further enhance the DKM results as well. We also compared DKM with a well-known regularization-based clustering method on GoolgeNet in Table 2, RPS (Wu et al., 2018) which has demonstrated superior performance to another regularization approach (Ullrich et al., 2017).
Note that only convolution layers are compressed, following the setup in RPS (Wu et al., 2018).
Table 2 clearly indicates that DKM can allow both much better compression and higher accuracy than RPS even with 1 bit-per-weight.
For a comprehensive study, we also compared our DKM-based algorithm with the latest scalar weight quantization approaches, PROFIT (Park & Yoo, 2020) and EWGS (J. Lee, 2021) (which have outperformed the prior arts in the low-precision regimes) by running their public codes on our environments with the recommended hyper-parameter sets.
Table 3 summarizes our comparison results on ResNet18, ResNet50, and MobileNet-v1/v2 for the ImageNet1k classiﬁcation task.
Following the experimental protocol in (Zhang et al., 2018; J. Lee, 2021; Rastegari et al., 2016), we did not compress the ﬁrst and last layers for all the experiments in Table 3.
It clearly shows that our approach with DKM can provide compression comparable to or better than other approaches, especially for the low-bit/high-compression regimes.
We denote clustering with b bits and d dimensions as b/d as it will assign bd bits in average to each weight, and the number of weight clusters is 2 b .
Especially with multi-dim clustering such as 4/4 or 8/8 bits, our DKM-based compression outperforms other schemes at 1 bit, while PROFIT cannot make training converge for MobileNet-v1/v2.
One notable result is 64.3% Top-1 accuracy of MobileNet-v1 with the 8/ bit conﬁguration (which is 1 bit-equivalent).
DKM with 8/16 bits (effectively 0.5 bit per weight) shows regression from the 8/8 bit conﬁguration, but still retains a good accuracy level.
We also tried
PROFIT+EWGS as proposed in (J. Lee, 2021), which showed good results on MobileNet-v1/v2 for 3 bits but failed to converge for 2 and 1 bits. [Table] [Figure]
v2 convergence with DKM 1/1: DKM delivers 50.8% top-1 accuracy withbit compression by gradually clustering the weights into two centroids using the task objective only.
With the overall compression ratio (or bit-per-weight) ﬁxed, our experiments with DKM conﬁrm that a higher d can yield a better quality training result.
For the example of MobileNet-v2, DKM 8/16 yielded 24% better top-1 accuracy than DKM 4/8 although both have the same 12 bit-per-weight, and the same trend is observed in other models.
However, DKM 8/8 failed to train ResNet50 due to the memory limitation, while DKM 8/16 successfully trained the same model, because the larger dimension (i.e., 8 vs 16) reduces the memory requirement of the attention matrix as discussed in
Section 3.3.
For additional discussion, please refer to the Appendix.
Fig.
4 lastly shows that DKM-based compression can offer a smooth convergence and gradual weight-clustering based on task loss back-propagated through DKM layers, without any extra regularization or custom training ﬂows.
4. GLUE NLP B ENCHMARKS
We compared our compression by DKM with GOBO (Zadeh & Moshovos, 2020) and
EWGS (J. Lee, 2021) for BERT models on NLP tasks from the GLUE benchmarks (Wang et al., [Table] [Table]
2019a) , QNLI (Question-answering NLI) and MNLI (Multi NLI). We ﬁxed the learning rate as 1efor all the experiments which worked best for EWGS
[Table]
and all experiments used mini-batch sizeper GPU (i.e., global mini-batch size of 1024) with the maximum seq-length 128.
We compared our DKM-based compression against EWGS (J. Lee, 2021) on the QNLI dataset, and
Table 4 demonstrates that DKM offers better predictability across all the tested models (Lan et al., 2019; Sanh et al., 2019; Turc et al., 2019; Sun et al., 2020) than EWGS.
Note that the embedding layers were excluded from compression in QNLI experiments.
As in ImageNet1k experiments, the 4/4 bit conﬁguration delivers better qualities than the 1 bit conﬁguration on all four BERT models, and especially performs well for the hard-to-compress MobileBERT.
Table 4 also indicates that different transformer architectures will have different levels of regression for a given compression target.
For the example of 1 bit, MobileBERT regressed most due to many hard-to-compress small layers, yet recovered back to a good accuracy with DKM 4/4.
When DKM compared against GOBO (Zadeh & Moshovos, 2020) (which has outperformed the prior arts on BERT compression) on DistilBERT with the MNLI dataset, our results in Table 5 clearly show that DKM offers a better accuracy-compression trade-off than GOBO, and also enables ﬁnegrained balance control between an embedding layer and others: using 2.5 bits for Transformer and 3 bits for embedding is better than 2 bits for Transformer and 4 bits for embedding for DistilBERT.
C ONCLUSION
In this work, we proposed a differentiable k -means clustering layer, DKM and its application to model compression.
DNN compression powered by DKM yields the state-of-the-art compression quality on popular computer vision and natural language models, and especially highlights its strength in low-precision compression and quantization.
The differentiable nature of DKM allows natural expansion to multi-dimensional k -means clustering, offering more than 22x model size reduction at 63.9% top-1 accuracy for highly challenging MobileNet-v1.
We plan to extend DKM to handle weight-clustering and weight-pruning simultaneously for more aggressive DNN model optimization, and apply sparse-format for the internal tensors to reduce the memory overheads.  R EPRODUCIBILITY S TATEMENT
Our universal setup for experiments is disclosed in the ﬁrst paragraph of Section 4 and per-datasetsetups are also stated in the ﬁrst paragraphs of Sections 4.1 and 4.2 along with key hyper-parameters in Appendix.
For Imagenet1k, we used a standard data augmentation techniques: RandomResizedCrop(224), RandomHorizontalFlip, and Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) as in other papers (Park & Yoo, 2020; J. Lee, 2021).
During evaluation, we used the following augmentations: Resize(256), CentorCrop(224), and Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]).
For GLUE benchmark, we used the default setup from HuggingFace
A A BLATION S TUDY : DKM S YSTEM P ERFORMANCE AND O VERHEAD
Table 6 shows the GPU memory utilization and per-epoch runtime for the DKM cases from Table 3.
While DKM layers have negligible impacts on training speed, the GPU memory utilization increases with more bits (i.e., more clusters) and with smaller dimensions, as discussed in Section 3.3: the storage complexity of a DKM layer with N parameters and d dimension is O ( r Nd 2 b ) where r is the number of iterations.
[Table]
Therefore, one straightforward way to reduce GPU usage and avoid the out-of-memory exception is to limit r .
In order to overcome the out-of-memory error for ResNet50 with DKM 8/8, we applied a constraint r < 3 to all the DKM layers, and could reach the top-1 accuracy of 74.0 % with the 98.2% GPU memory utilization.
B A BLATION S TUDY : H YPER -P ARAMETER τ SEARCH
In the current DKM implementation, we use a global τ to control the level of softness in the attention matrix.
The selection of τ affects the model predictive power as shown in Fig.
5 where there appears to be an optimal τ for a given DNN architecture.
For examples of ResNet18/80, τ = 2 e − 5 is the best value for the 2 bit clustering.
In our experiments, we used a binary search to ﬁnd out the best τ values w.r.t.
the top-1 accuracy, which are listed in Table 7.
In general, one can observe that a complex compression task (i.e., higher compression targets, more compact networks) tends require a larger τ to provide enough ﬂexibility or softness.
• For MobileNet-v1/v2, it requires about 10x larger τ values then for ResNet18/50, because they are based on a more compact architecture and harder to compress.
• When the number of bits decreases, the compression gets harder because there are fewer centroids to utilize, hence requiring a larger τ value.
• When the centroid dimension increases, the larger τ value is required, as the compression complexity increases (i.e., need to utilize a longer sequence).
[Table]
It could be possible to cast τ as a learnable parameter for each layer or apply some scheduling to improve the model accuracy further (as a future work)
[Table]
but still both approaches need a good initial point which can be found using a binary search technique. For the BERT experiments with GLUE benchmarks, we used the following τ regardless of the compression level: 5.0e-5 for ALBERT, 8.0e-5 for DistilBERT, 1.5e-4 for BERT-tiny [Table] and 4.0e-4 for
MobileBERT. We found the BERT models are less sensitive to the τ than ImageNet classiﬁers. C R ELATION TO E XPECTATION -M AXIMIZATION (EM)
Special sub-case of DKM where gradients are not propagated can be related to a standard EM formulation.
Here is the formulation-level correspondence: Suppose
[Equation]
is the Gaussian mixture model over cluster centers.
Referring to Fig.
2, Maximizing log likelihood of weights ln P
( W | C )
= P N i =1 ln (cid:8)P K j =1 1 K N ( w i | c j , σ 2 = τ/ 2) (cid:9) using the EM algorithm is equivalent to DKM for the case of d i,j = − ( w i − c j ) 2 .
Then, the attention matrix A is equivalent to the responsibilities calculated in the E step, and updating C is equivalent to the M step.
However, unlike
EM where ﬁnding C every M step is the objective, DKM focuses on generating a representative ˜W for the train-time compression for DNN.
Even though there is formulation-level similarity between DKM and EM, the way both are optimized is signiﬁcantly different.
While EM iteratively optimizes a speciﬁc likelihood function for a set of ﬁxed observations, DKM needs to adjust (i.e., optimize) the observations (which are weights) without leading to a trivial solution such as all observations collapsing to a certain point.
Hence,
DKM can neither assume any statistical distribution nor optimize a speciﬁc likelihood function (i.e., the observations are dynamically changing).
Therefore, DKM uses a simple softmax and rides on the back-propagation to ﬁne-tune the observations w.r.t.
the task loss function after unrolling multiple attention updates.
When we propagate gradients, then this will turn into a stochastic nonconvex joint optimization where we simultaneously optimize observations and centroids for the task loss function, which is shown to offer better accuracy vs. compression trade-offs according to our experiments.