DKM: Differentiable k-Means Clustering Layer for Neural Network Compression
Authors: Minsik Cho
Keivan Alizadeh-Vahid
Saurabh Adya
Mohammad Rastegari
Venue: ICLR 2022 Poster (2022)
Keywords: Deep learning
neural network
compression

Abstract: 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.

Published as a conference paper at ICLR 2022
DKM: DIFFERENTIABLE k-MEANS CLUSTERING LAYER FOR NEURAL NETWORK COMPRESSION
Minsik Cho
Keivan Alizadeh-Vahid'
Saurabh Adya
Mohammad Rastegari
Apple
(minsik, kalizadehvahid, sadya, mrastegari] @apple com
ABSTRACT
Deep neural network (DNN) model compression for efficient on-device inference becomes increasingly important reduce memory requirements and keep user data on-device. To this end, we propose a novel differentiable k-means clus- tering layer (DKM) and its application train-time weight- for DNN model compression. DKM casts k-means clustering as an attention problem and enables joint optimization of the DNN parameters and clustering centroids. Un- like 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 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 ImageNetlk and GLUE benchmarks. For example. DKMebased compression can offer 74.5% top-1 Ima- geNet1k accuracy on ResNet50 with 3.3MB model size (29.4x model compression factor). For MobileNet-v1, which a challenging DNN to compress, DKM de- livers 63.9% top-1 ImageNetIk 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 DistiBBERT model by 11.8x with minimal (1. %) accu- racy loss on GLUE NLP benchmarks.
INTRODUCTION
Deep neural networks (DNN) have demonstrated super-human performance on many cognitive tasks (Silver et all 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 la- tency and keeping user data on-device. Many such on-device platforms are battery-powered and sscurce-constrainee demanding DNN to meet the stringent resource requirements such as power- consumption, compute budget and storage-overhead (Wang et al., 2019h; Wu et al., 2018).
One solution is to design a more efficient and compact DNN such as MobileNet (Howard et all 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 compress model with small accuracy degradation so that it takes less storage and reduces System-on-Chip (SoC) memory band- width utilization, which
minimize power-consumption and latency. To this end. various DNN compression techniques have been proposed (Wang et al.. 2019;: Dong et all 2020 Park et al 2018: Rastegari et all (2016; Fan et al. 2021: Stock all 2020 Zhou et all 2019 Park et all 2019
Yu et all 2018: Polino et al. 2018). Among them, weight-clssterinnssaaae
(Han et al.
2016
Wu et al. 2018 Ullrich et al. 2017: Stock et al., (020) has shown high DNN compression ratio
where weights are clustered into a few shareable weight values (or centroids) based on k-means clus- tering. Once weights are clustered, to shrink the model size, one can store indices (2bits, 4bits, etc. depending on the number of clusters) with lookup table rather than actual floating-poii point values.
"qqual contribution

---

Published as a conference paper at ICLR 2022
Designing 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 ap- proaches do not usefully compress an already 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 obser- vations already NP-Hard), b) the weight-to-cluster assignment is 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 clus- tering, DKM, based on an attention mechanism (Bahdana et all, (015)) 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 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 high-quality model for given compression ratio target.
We apply DKM to compress DNN model and demonstrate the state-of-the-ard results on both computer vision and natural language models and tasks.
RELATED WORKS
Model compression using clustering: DeepCompression (Han et all 2016) proposed to apply k- means clustering for model compression. DeepCompression initially clusters the weights using k- means algorithm. All the weights that belong 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 update the shared value. This approach might degrade model quality because it cannot formulate weight-cluster assignment during gradient back propagation (Yin all 2019). ESCQ (Choi et all 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. 2019h) uses reinforcement learning to search for the optimal quantization policy on different tasks. For model compression, HAQ uses k-means clustering similar to DeepCompres. sion 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. (020) algorithm is based on Product Quantization and Knowledge Distillation. 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 all 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 train- ing process is not straightforward (Wu et al. 2018). Hence, (Ullrich et al], (017) models weight- clustering as Gaussian Mixture Model (GMM) and fits weight distribution into GMM with addi- tional learning parameters using KL divergence (i.c., forcing weight distribution to follow k Gaus- sian distributions with a slight variance). (Wu et all (018) proposed deep k-means to enable weight- clustering 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 all 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 up- dates for the new variables (i.e., singular value decomposition (SVD) in (Wu et all 2018)). Also, relying on the modified loss cannot capture the dynamic interaction between weight distributions and cluster centroids within batch, thus requiring an additional training flow for re-training.
2

---

Published as a conference paper at ICLR 2022
Enhance Model compression using dropout: Quant-Noise (Fan al 1021 is structured dropout which only quantizes a random subset of weights (using any quantization technique) and thus improve the predictive power of compressed model. For example, (Fan et al 1 showed good compression vs. accuracy trade-off on ResNer50 for ImageNet1k.
Model quantization: Besides clustering and regularization methods, model quantization also reduce the model size, and training-tima quantization techniques have been developed to improve the accuracy of quantized models (Li et all 2019 Zhao all 2019). EWGS 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 & Ld 2019: 2021) and ESPNet (Mehta et al. 2019).
MobileNet-vi (Howard et all (2017) on ImageNetlk dataset has top-1 accuracy of 70.3% with 16.1 MB of memory in comparison 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
Han al. 2016 Dong et all 2020; (Soock et al., (2020) are based on k-means clustering along with enhancements such as gradient scaling/approximatioo. Using k-means clustering, the weights are clustered and assigned to the nearest centroids which are used for pawwardbbckwarrlprpppgg
Popular weight-clusterin techniques for DNN model compression (J. Lec, 2021:
during training as illustrated in Fig. (a). Such conventional methods with clustering have two critical drawbacks:
The weight-to-cluster assignment in conventional approaches is not optimized through buck-propagation of training loss function.
Gradients for the weights are computed in ad-hoc fashion: the gradient of 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 i and in Fig. (a). In the conventional approaches. and j are assigned to the centroids C2 and C₁ respectively, simply because of their marginal difference in distance metric. However, assigning to Co and Co could be better for the training loss as their difference in distance is small (Nagel et al. 2020).
CHrvic
Cerood
Caunnd
Altention Matris
AsiignasT Grddeel
CC
Ci
C2
0,3
0.2
0.3
0.5
(a) Conventional weight- (Han et all 2016:
(b) Attention based
in DKM
Wang et all 2019h Stock et al. 2020 J. Lee 2021)
Figure In conventional weight-clustering algorithms, the boundary weights i and j are assigned to clusters C2 and C1 based on the distance metric respectively, which is neither necessarily suitable for the task nor differentiable against the loss function. DKM instead applies soft assignment using attention mechanism during nowward-propggate and enables differentiable backward propaaapooonniionn which allows weights to consider other non-nearest clusters (especially helpful for boundary weights) shuttle among multiple clusters order to directly optimize their assignments based on the task loss function.
3

---

Published as a conference paper at ICLR 2022
Iteration until converges using imperative-executie works OR the iteration limit reached
t=O, use k-means++
Update centroids
Otherwise, use last
C
C Č
Σ, 4. 2
if|c-C12c
expl
Σ,ag
W
d, f(w s)
Centroid Matrix C
W
Σ, expl
if
AC Activation Layer for Clustering
Distance Matrix
Attention Matrix
weighted sum over centroids
Figure 2: Weight-sharing using attention matrix A is iteratively performed in a DKM layer until the centroids (C) converge. Once converged, compressed weight, W used for forward-por wwoarddpropagation.. Since DKM is differentiable layer, baccward-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 hard assignment can degrade the training loss significantly.
We overcome such limitations with DKM by interpreting weight-centroid assignment as distance- based attention optimization (Bahdana al_ 2015) as Fig. ① (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 allows gradient of weight to be 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.
waokkwawdd-rrppgatiiii attention
3.2 DIFFERENTIABLE K-MEANS CLUSTERING LAYER FOR WEIGHT-CLUSTERING
DKM can perform differentiable train-time weight-clustering iteratively for clusters shown in Fig. 2 for the DNN model compression purpose. Let C R be a vector of cluster centers and W E RN be vector of the weights, and then DKM performs follows:
In the first iteration C can be initialized either by randomly selected weights from W or k- k-means++. For all subsequent iterations, the last known C from the previous batch used to accelerate the clustering convergence.
A distance matrix D is computed for every pair between weight w; and centroid e; using differentiable metric f (i.e., the Euclidean distance) using d;- = - Cj).
We apply softmax with a temperature т on each row of to obtain attention matrix A where a.jj
emp(
represents the attention from wy and Cj-
Σ. ezpl
Then, we obtain a centroid candidate, Č by gathering all the attentions from W for each centroid by computing é
E. w. and update C with Č if the iteration continues. E. aig
We repeat this process till |C č < at which point k-means has converged or the itera- tion limit reached, and compute AC to get W for forward- peraard-rrppgattion (as in Fig. 2
The iterative process will be dynamically executed imperatively PyTorch (Paszke et al., 2019)
and Tensorflow-Eagee (Agrawal et all (219)) and differentiable for backward wwwoaddpropggation as is based on the attention between weights and centroids. DKM uses soft weight-cluster assignment which could be hardened in order to impose weight
constraints. The level of hardness be controlled by the temperature in the softmax operation. During inference attention matrix (i.e., A in Fig. from a DKM layer to snap each weight to the closest centroid the layer and finalize weight-clustering in prior arts (i.e., more attention), but such assignment is
use the last
4

---

Published as a conference paper at ICLR 2022
expected to be tightly aligned with the loss function, the weights have been annealed by shuttling among centroids. A theoretical interpretation of DKM is described in Appendix
Using DKM for model compression allows weight to change its cluster assignment during training, 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 the weight-cluster assignment is by the attention mechanism. Since DKM is without additional learnable parameters and transparent to model and loss function, we can reuse the existing training flow and hyper-parameters, The key 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 all 2016 Wang et al., 2019b: Lee, [202] Stock all 2020), DKM uses flexible and differentiable attention-based weight-clustering and computes gradients w.r.t the task loss.
Instead of modifying the loss function with regularizers to enforce clustering (Ullrich et all 2017: Wu et all 2018), DKM can be inserted into forward pass, making the optimization fully aligned with the task objective (i.c., no interference in loss).
DKM requires no additional learnable parameters (Ullrich et all 2017: J. Lee, 202 l), thus making the training flow simple. For example, DKM-base approach does not need to sub- stitute a convolution layer with specialized version with additional learning parameters.
DKM requires no additional computation such as Hessian trace (Dong et all 2020
Choi et al. 2017) or SVD Les. 202 Wu et al. 2018) for gradient approximation, be- cause DKM uses a differentiable process.
DKM-based compression does not require a complex training techniques such as freez- ing/progress (Park & Yoo, 2020) or distillation (Stock et al. 2020).
3.3 MUL.T-DIMENSIONAL DKM
DKM can be naturally extended into multi-dimensionn 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 d contiguous d dimensional sub-vectors and cluster the sub-vectors (W E *d) For example, simply flatten all the convolutional kernels into a ( d) matrix across both kernel and channel boundaries and apply multi-dimennional DKM to the matrix for clustering in our implementation as in Fig. 3 Accordingly, the cluster centroids will become d-dimensional
N
well (C R(*d) and the metric calculation is done in the d-dimensional space. With the multi-dimensional scheme, the effective bil-per-weight becomes 4 for b-bit/d-dim clustering. The memory complexity of DKM layer with N parameters is Orr 22) where r is the number of iterations per Fig. 2 (i.e., all the intermediate results such D and A at each iteration need to be kept for backward-propagatee
Such multi-dimensional clustering could be ineffective for conventional methods (i.e., DNN training not converging) (Stock et al. 2020; Wang et al. 2019h: J. Lec, 2021), 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 only two centroids for 1bit/Idin clustering, while there are 16 centroids in 4bit/4dim clustering, although both have the same effective bit-per- weight. Intuitively, however, DKM can work well with such multi-dimensional configurations as DKM naturally optimizes the assignment w.r. the task objective and can even recover from a wrong assignment decision over the training-time optimization process.
Cp
compression over 32bit probability of using 4
cr 32N d b
Pt Plc E W] 2
[W) Wi w 0.6, 0.3
ew Σ.,ec -p; loga Pi
entropy of W
Attention Matrix
= 20 1 log2 J.1 =
2:
Figure 3: Multi-dimensionn DKM increase the weight entropy (ew) with fixed compression ratio target (cr) by increasing dimension (d), which may improve the model accuracy (Park et al. (01
5

---

Published as a conference paper at ICLR 2022
50
32
ATB-large
DKM
cY:3/8,fc:844
cV:6/6.fc:614
Compression Ratio
45
cv:3/88.c:8444
Compression Ratio
28
ATB-large with
ATB-large
DKM
quant-noise
40
35
LiR-Net-Ibit
cv:688
DKM
20
ATB-1mall
Meta Quant
ABC-Net-MI
fv:444,
DC-2hit
HAQ-flx
30
16
ATB-small with quant-moise
GBS
25
cv:4/4,fc:814
ATB-small
59%
62%
65%
68%
66%
68%
70%
72%
74%
76%
Top-1 Accuracy
Top-1 Accuracy
(a) ResNet18
(b) ResNet50
Base
DC
HAQ
ATB
ATB large
DKM configuration
Model
Metrics
32bit
2bit
flex
small
b/w"
Top-1 (%)
69.8
65.8
61.1
65.1
ev1.5/8
ResNet18
0.717
Size (MB)
44.6
1.58
1.07
1.00
fct:8/4
73.8
68.2
Top-1 (%)
76.1
68.9
70.6
74.5
cv:6/6
ResNet50
74.3
68.8
1.077
Size (MB)
97.5
6,32
6.30
5.34
3.43
3.32
Ic:6/4
MobileNet
Top-1 (%)
70.9
37.6
57.1
63.9
cv:4/4
v1
Size (MB)
16.1
1.09
1.09
ne
ne
0.72
fo:4/2
1.427
MobileNet v2
Top-I (%)
71.9
58.1
66.8
68.0
ev:2/1
Size (MB)
13.3
0.96
0.95
ne
nc
0.84
fc:4/4
2.010
effective wei-ht (see Section 3.. 2 not converging
the convolution layers the last fully connected layer; clustering with 6 hits and 8 dimensions; ATB with quantization-noise (Fan et al 2021
also, 65.7% Top-1 accuracy and 1.09 MB with ev:6/8, fc:6/10 66.7% Top-1 accuracy and 1.49 MB with ev:4/4, fc:8/4
Figure 4: Accuracy vs. compression trade-offs: DKM-powered compression white box markers) deliv- ers the Pareto superiority to the other schemes (i.e., the top-right corner is the best trade-off) for ResNet18, ResNet50, and 1112 ImageNetlk.
The key benefit of limei-dimmnsional DKM is captured in Fig. For a given >> d in 32 bits, the compression ratio is d (i.e., function of both d and b). Assuming the number of sub-vectors assigned to each centroid is same, the entropy W is simply b. Since higher entropy in the weight distribution indicates larger learning capacity and better model quality (Park al., 2017), increasing b and d at the same ratio as much as possible may improve the model quality for a given target com- pression ratio (see Section 4. for results). However, making and extremely large maximize the entropy of W might be impractical, as the memory overhead grows exponentially with 6.
EXPERIMENTAL RESULTS
We compared our DKM-based compression with state-of-the-art quantization or compression schemes on various computer vision and natural language models. To study the trade-off between model compression and accuracy, we disabled activation quantization in every experiment for all approaches, as main goal is the model compression. All experiments with DKM were done on two x86 Linux machine with eight NVIDIA V100 GPUs each in a public cloud infrastructure.
We used SGD optimizer with momentum 0.9, and fixed the learning rate 0.008 (without individ- ual
tuning) for all the experiments for DKM. Each compression scheme starts with publicly available pre-trained models. The is set le-4 and the iteration limit 5.
4.1 IMAGENETIK
We compared our DaMed compression with prior arts: DeepCompression (or DC Han et all 2016), HAQ (Wang et all 2019H), and "And The Bit Goes Down" (or ATB) (tock al 2020

---

Published as a conference paper at ICLR 2022
ResNet18 69.8
ResNetSO 76.1
MobileNe-vl
MobileNet-12
Base (32 bit)
70.9
71.9
PROFIT
69.6
69.6
bit
EWGS
70.5
76.3
64.4
64.5
PROFIT+EWGS
68.6
69.5
DKM
69.9
76.2
69.9
70_3
PROFIT
63.4
61.9
bit
EWGS
69.3
75.8
52.0
49.1
DKM
68.9
75.3
66.4
66.2
PROFIT
nc"
nc
EWGS
66.6
73,8
8.5
23.0
1 bit
DKM 1/1
65.0
72.1
5.9
50,8
DKM 4/43
67.0
73.8
60.6
55.0
DKM 8/8
67.8
opm
64.3
62,4
bit
DKM 478
62.1
70.6
46.5
34.0
DKM 8/16
65.5
72.1
59.8
58.3
converging: out of memory 8 clustering with 4 hits and 4 dimensions
Table 1: When compared with the latest weight quantization algorithms, DKM-based algorithm shows superior Top-1 1 accuracy when the network hard to optimize (i.c., MobileNet-v1/v22 or when a low precision is required
bit). Further, with multi-dimensional DKM (see Section [ DKM delivers 64.3 % Top-1 accuracy for
MobileNet-vl with the configuration which is equivalent bit-per-weight.
Metries
Base (32bit)
RPS
DKM 4/1
DKM 4/2
DKM 8/8
Top-1 69.8
67.9
70.9
70.3
68.5
e
8
compression ratio: clustering with bits and dimensions
Table 2: GoogleNet training performance for ImageNetlk: DKM-based compression offered 2x better pression ratio with 3% higher top-1 accuracy than RPS (Wu all 2018).
bined with Quantization-noiss Fan et al. 202 1) ABC-Net (Lin et al. 2017). BWN (Rastegari et all
2016), LR-Net (Shayer et all, 2018). and Meta- Quant (Chen et all 2019). We also compared with GBS which uses the same flow as DKM except that Gumbel-softmax is used to generate stochastic soft-assignment as attention (Jang et al. 2017). the GBS implementation, we iteratively perform drawing to mitigate the large variance problem reported in (Shaver et al. 2018).
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 et al. 201 8), we added the support for these two by following the paper and the existing ResNet18/50 (He al 2016) implementations. Instead of using complex RL technique as in HAQ (Wang et al. 2019b), for DKM experiments, we fixed configurations for all the convolution layers (noted cv) and the last fully connected layer (noted fc), except that we applied 8 bit clustering to layer with fewer than 10,000 parameters.
Using DKM layer, our compression method offers Pareto superiority other schemes visual- ized in Fig. 4 For ResNet50 and MobileNet-v1//2. DKM delivered compression configurations 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 find Pareto superior configu- rations all others: DKM can get 65.1 5..1% Top-1 accuracy with IMB which is superior 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 superior balance point. For MobileNet-v 1/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, adding quantization noise improves the model accuracy (Fan et all 2021) only moderately. GBS in fact shows better performance than ATB, but still worse than the proposed method, even after reduc- ing variance through iteration, especially for the high compression configurations. For ResNet18, GBS with the same bit/dimension targets delivered the following Top-1 accuracies ranging from 61.7% to 65.7%. For details, please refer to Appendix
We also compared DKM with well-known regularization-bie based clustering method GoolgeNet in T RPS (Wu et al. 2018) which has demonstrated superior performance to another regular- ization approach (Ullrich et al. 2017). Note that only convolution layers compressed, following
7

---

Published as a conference paper at ICLR 2022
ALBERT 90.6
DistiBERTT 88.2
BERT-tiny 78.9
MobileBERT
Base (32 bit)
89.6
EWGS
83.3
87.6
78.3
87.8
DKM
85.1
88.2
80.0
89.0
EWGS
79.6
85.4
77.9
81.6
DKM
81.7
87.4
80.0
83.7
EWGS
62.0
60.9
74.5
60.2
DKM
79.0
82.8
77.4
69.8
DKM 4/48
80.0
84.0
77.2
78.3
clustering with 4 bits and 4 dimensions
Table 3: Training performance for QNLI: DKM-based scheme outperforms EWGS in compressing various transformed-bast architectures, Also, multi multi-dimmnnsooal DKM (see Section (33) largely improved the accu- racy of MobileBERT with 1 bit-per-weight target using the 4/4 configuration.
Base
GOBO
DKM xform4/2 .em44
DKM Xform5/2.em
Metries
32bit
xform ',emb'4 4
Top-1
82.4
81.3
81.3
81.3
Size (MB)
255.4
23.9
21.8
21.5
the transformer layers: ; the embedding layer
$ clustering with 4 bits and 2 dimensions
Table 4: DistillBert training performance for MNLI: DKMM-ased compression offered 10% smaller model size with the same accuracy target than GOBO.
the setup in RPS (Wu all 2018). Table Table@cleart clearly indicates that DKM can allow both much better compression and higher accuracy than RPS even with I bit-per-weight.
We also compared our DKM-based algorithm with the latest sealar weight quantization approaches, PROFIT (Park & Yoo 2020) and EWGS Lee. 2021) (which have outperformed the prior arts in the low-precision regimes) by running their public codes on our environments with the recom- mended
sets. Table summarress our comparison results ResNet18, ResNet50, and MobileNet-v1/v2 for the ImageNetlk classification task. Following the experimental proto- col in (Zhang et ald 2018 J. Lee, 2021: Rastegari et al., 2016). we did not compress the first and last layers for all the experiments in Table [ It clearly shows that our approach with DKM can provide compression comparable
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 å bits in average to each weight, and the number of weight clusters is 2h. Especially with multi-dim clustering such as 4/4 or 8/8 bits, our DKMMbased compression outperforms other schemes at I bit, while PROFIT cannot make training converge for MobileNet-v1/v2. One notable result is 64.3% Top-1 accuracy of MobileNet-vl with the 8/8 bit configuration (which is bit-equivalent). DKM with 8/16 bits (effectively 0.5 bit per weight) shows degradation from the 8/8 bit configuration. but still retains a good accuracy level. We also tried PROFIT+EWGS as proposed in (J. Lee. 202 which showed good results MobileNet-v 1/v2 for bits but failed to converge for and bits.
With the overall compression ratio (or bit-per-weight) fixed, our experiments with DKM confirm that a higher d can yield better quality training result. For the example of MobileNet-v2, DKM 8/16 yielded 24% better lop-1 accuracy than DKM 4/8 although both have the same Z 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.. vs 16) reduces the memory requirement of the attention matrix as discussed in Section 3.3
For additional discussion, please refer to the Appendix
4.2 GLUE NLP BENCHMARKS
We compared our compression by DKM with GOBO (Zadeh et al. 2020) and EWGS Les 10 for BERT models on NLP tasks from the GLUE benchmarks (Wang et all 2019a), QNLI (Question- answering NLI) and MNLI (Multi NLI). We fixed the learning rate as le-4 for all the experiments which worked best for EWGS, and all experiments used mini-batch batch size 64 per GPU (i.e., global mini-batch size of 1024) with the maximum seq-length 128.
8

---

Published as a conference paper at ICLR 2022
Network
ResNet18
ResNet50
MobileNet-vi
MobileNet-v2
ev:5/8
cv/4/4
cv:6/8
cv:6/6
ev:4/4
cv:2/1
configuration
fc:8/4
fc:8/4
fc:6/10
fc:6/4
fc:4/2
fc:4/4
Inference-time
65.1
66.7
65.7
74.5
63.9
68.0
Train-time
66.0
67.0
66.5
74.7
65.6
68.8
Table 5: Top-1 Accuracy with Train-tine and inference-time Weights for Fig.
4
Top-1 Accurucy
40.0%
Tiwe
50.0%
10.0%
0.0%
Epochs
Top-1 I accuracy
(h) Weight distribution of the largest conv layer
Figure 5: MobileNet-v2 convergence with DKM 1/1: DKM delivers 50.8% top-1 accuracy with I bit compres- sion by gradually clustering the weights into two centroids using the task objective only.
We compared our DKM- based compression against EWGS J. Lec, 2021) on the QNLI dataset, and Table 3 demonstrates that DKM offers better predictability across all the tested models (Lan et all 2019: Sanh et al.. 2019: Ture et all 2019 Sun et all (020) than EWGS. Note that the embedding layers were excluded from compression in QNLI experiments. As in ImageNet experiments, the 4/4 bit configuration delivers better qualities than the 1 bit configuration on all four BERT models, and especially performs well for the hard-to-compress MobileBERT. Table 3 also indicates that different transformer architectures will have different levels of accuracy degradation for given compression target. For the example of 1 bit, MobileBERT degraded most due to many hard-to- compress small layers, yet recovered back to good accuracy with DKM 4/4.
When DKM compared against GOBO (Zadeh all 2020) (which has outperformed the prior arts on BERT compression) on DistilBERT with the MNLI dataset, our results in Table clearly show that DKM offers a better
trade-off than GOBO, and also enables fine-grained balance control between embedding layer and others: using 2.5 bits for Transformer and 3 bits for embedding is better than bits for Transformer and 4 bits for embedding for DistilBERT.
4.3 DKM-BASED COMPRESSION ANALYSIS
Since DKM-based compression uses attention-driven clustering as in Fig. dluring training but snaps the weights to the nearest centroids, there is a gap between train-time and inference-time weights which in turn leads to accuracy difference between train and validation/test accuracies. Therefore, we measured the Top-1 accuracies with both weights for the DKM cases from Fig. 4 as in Table 5 We observed the accuracy drop is about 0.2% 1.7% and gets larger with comppess DNNs.
Fig. 5 shows that our compression approach can offer a smooth convergence and gradual weight- clustering based task loss back- poopppaop through DKM layers, without any custom training flow and extra regularization. For additional details and results, please refer to Appendix
CONCLUSION
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 compres- sion quality on popular computer vision and natural language models, and especially highlights its strength in low ow-precision compression and quantization. The differentiable nature of DKM al- lows natural expansion to multi-dimenssonall k-means clustering, offering more than 22x model size reduction 63.9% top-1 accuracy for highly challenging MobileNet-v1.

---

Published as a conference paper at ICLR 2022
6 REPRODUCIBILITY STATEMENT
Our universal setup for experiments is disclosed in the first paragraph of Section per-dataset- setups are also stated in the first paragraphs of Sections and 4.2 allon with key hyper-parameters in Appendix.
For Imagenetlk, we used a standard data augmentation techniques: RundomResizeddrrpp(2.4) Ran- domHorizontalFlipp, and (ormal [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 aug- mentations: Resize(256), CentorCrop(224), and Normalize(mean=(OAa
0.456, 0.406), std=[00.29,, 0.224, 0.225]). For GLUE benchmark, we used the default setup from HuggingFace.
REFERENCES
Akshay Agrawal, Akshay Naresh Modi, Alexandre Passos, Allen Lavoie, Ashish Agarwal, Asim Shankar, Igor Ganichev, Josh Levenberg, Mingsheng Hong, Rajat Monga, et al. Tensorflow eager: A multi-stage, python-embedded dsl for machine learning. arXiv preprint 1903:1903301155 2019.
Dzmitry Bahdana, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference Learning Representations, 2015.
Shangyu Chen, Wenya Wang, and Sinno Jialin Pan. Metaquant: Learning to quantize by learning to
penetrate nonfdifferentiable quantization. In Advances in Neural Information Processing Systems, 2019.
Yoojin Choi, Mostafa El-Khamy, and Jungwon Lee. Towards the limit of network quantization. In International Conference Learning Representations, 2017.
Yoojin Choi, Mostafa El-Khamy, and Jungwon Lee. Universal deep neural network compression. IEEE Journal of Selected Topics in Signal Processing, 14(4):715-726, 2020, doi: 10.1109/SSTP. 2020.2975903.
Zhen Dong, Zhewei Yao, Daiyaan Arfeen, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Hawq-v2: Hessian aware trace-weighted quantization of neural networks. In Advances in Neural Information Processing Systems, 2020.
Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Rémi Gribonval, Hervé Jégou, and Armand Joulin. Training with quantization noise for extreme model compression. In International Conference on Learning Representations, 2021.
Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural network with pruning, trained quantization and huffman coding. In International Conference Learning Representations, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 04861, 2017.
B. Ham J. Lee, D. Kim. Network quantization with element-wise gradient scaling. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteriization with gumbel-softmax, In Internationaal Conference on Learning Representations, 2017.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. ALBERT: A lite BERT for self-supervised learning of language representations. In Interna- fional Conference on Learning Representations, 2019.
10

---

Published as a conference paper at ICLR 2022
Yuhang Li, Xin Dong, and Wei Wang. Additive powers-of-two quantization: An efficient non- uniform discretization for neural networks. International Conference on Learning Representa- tions, 2019.
Xiaofan Lin, Cong Zhao, and Wei Pan. Towards accurate binary convolutional neural network. In Advances in Neural Information Processing Systems, 2017.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In International Conference on Learning Representations, 2019.
Sachin Mehta, Mohammad Rastegari, Linda Shapiro, and Hannaneh Hajishirzi. Espnetv2: A light- weight, power efficient, and general purpose convolutional neural network. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, 2019.
Markus Nagel, Rana Ali Amjad, Mart Van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? Adaptive rounding for post-training quantization. In International Conference on Machine Learning, 2020.
Eunhyeok Park and Sungjoo Yoo. Profit: A novel training method for sub-4-bit mobilenet models. In European Conference Computer Vision, 2020.
Eunhyeok Park, Junwhan Ahn, and Sungjoo Yoo.
Weighted-entropy-baaeed quantization for deep neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recog- nition, 2017.
Eunhyeok Park, Sungjoo Yoo, and Peter Vajda. Value-aware quantization for training and inference of neural networks. In European Conference on Computer Vision, 2018.
Sejun Park, Jaeho Lee, Sangwoo Mo, and Jinwoo Shin. Lookahead: A far-sighted alternative of magnitude-based pruning. In International Conference Learning Representations, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Ed- ward Yang. Zach De Vito. Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 2019.
Antonio Polino, Razvan Pascanu, and Dan-Adrian Alistarh. Model compression via distillation and quantization. In International Conference on Learning Representations, 2018
Mohammad Rastegari, Vicente Ordonez, Joseph Redmon, and Ali Farhadi. Xnor-net: Imagenet classification using binary convolutional neural networks. In European Conference on Computer Vision, pp. 525-542. Springer, 2016.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang Chich Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, June 2018.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of BERT; smaller, faster, cheaper and lighter. In Advances in Neural Information Processing Systems, 2019.
Oran Shayer, Dan Levi, and Ethan Fetaya. Learning discrete weights using the local reparameteri- zation trick. In International Conference Learning Representations, 2018.
David Silver, Thomas Hubert, Julian Schritwieser,, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self self-play. Science, 362(6419)::144 140- 1144, 2018
Pierre Stock, Armand Joulin, Rémi Gribonval, Benjamin Graham, and Hervé Jégou. And the bit goes down: Revisiting the guantization of neural networks. In International Conference on Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2022
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: a compact task lask-agnostic bert for resource- -imurree devices. In Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics, pp.
2020.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International Conference on Machine Learning, pp. 6105-6114. PMLR, 2019.
Mingxing Tan and Quoc V Le. Efficientnetv2: Smaller models and faster training. arXiv preprint
arXiv:2104.00 00298, 2021.
Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoe V Le. Mnasnet: Platform-aware neural architecture search for mobile. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2820-2828, 2019.
Julia Ture, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better:
The impact of student initialization on knowledge distillation. arXiv preprint arXiv:1908.000 089662. 13, 2019.
Karen Ullrich, Edward Meeds, and Max Welling. Soft weight-sharing for neural network compres- sion. In International Conference on Learning Representations, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Interna- tional Conference on Learning Representations, 2019a
Kuan Wang, Zhijian Liu, Yujun Lin, Ji Lin, and Song Han. Haq: Hardware-aware automated quan- tization with mixed precision. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019b.
Junru Wu, Yue Wang, Zhenyu Wu, Zhangyang Wang, Ashok Veeraraghavan, and Yingyan Lin. Deep k-means: Re-training and parameter sharing with harder cluster assignments for compressing deep convolutions. In International Conference Machine Learning, 2018
Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley Osher, Yingyong Qi, and Jack Xin. Under- standing straight-through estimator in training activation quantized neural nets. In International Conference on Learning Representations, 2019.
Ruichi Yu, Ang Li, Chun-Fu Chen, Jui-Hsin Lai, Vlad Morariu, Xintong Han, Mingfei Gao, Ching- Yung Lin, and Larry S Davis. Nisp: Pruning networks using neuron importance score propagation. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 9194- 9203, 2018.
Ali Hadi Zadeh, Isak Edo, Omar Mohamed Awad, and Andreas Moshovos. Gobo: Quantizing attention-basee nlp models for low latency and energy efficient inference. 2020 53rd Annual JEEE/ACM International Symposium on Microarchitecture, pp. 81 811-824. IEEE, 2020.
Dongqing Zhang, Jiaolong Yang, Dongqiangzi Ye, and Gang Hua. Lq-nets: Learned quantization for highly accurate and compact deep neural networks. In European Conference on Computer Vision, 2018
Xiandong Zhao, Ying Wang, Xuyi Cai, Cheng Liu, and Lei Zhang. Linear symmetric quantization of neural networks for low-precision integer hardware. In International Conference on Learning Representations, 2019.
Daquan Zhou, Xiaojie Jin, Qibin Hou, Kaixin Wang, Jianchao Yang, and Jiashi Feng. Neural epit- ome search for architcttue--agnotic network compression. In International Conference on Learn- ing Representations, 2019.
12

---

Published as a conference paper at ICLR 2022
A ABLATION STUDY: DKM SYSTEM PERFORMANCE AND OVERHEAD
Table the GPU memory utilization and per-epoch runtime for the DKM cases from Table II 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 [33] the memory complexity of a DKM layer with N parameters and d dimension is O(r 24) where y' is the number of iterations.
ResNet18
ResNet50
GPU memory utilization (%)
Per-epoch runtime (sec)
GPU memory utilization (%)
Per-epoch runtime (sec)
bit 23.8
414.1
55.3
425.0
2 bit 21.7
401.2
49.2
429.3
bit 20.4
413.3
45.3
426.3
4/48 bit
22.3
414.4
51.8
430.4
4/8 bit
20.7
423.6
46.7
410.4
8/8 bit
51.8
409.4
oom
8/1 bit
35.1
421.9
78.3
454.6
8 clustering with bits and dimensions
Table 6: Memory and Runtime overheads from DKM on R/5Net18/50.
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, applied a constraint < 3 to all the DKM layers, and could reach the top-1 accuracy of 74.0 % with the 98.2% GPU memory utilization. We believe the fundamental solution this bottleneck is use sparse representation for A by keeping top-k centroids for each weight which is one of our future works.
B ABLATION STUDY: HYPER-PARAMETER SEARCH
In the current DKM implementation, we use a global 7 to control the level of softness the attention matrix. The selection of affects the model predictive power shown in Fig. where there appears be an optimal for given DNN architecture. For examples of ResNet1 2e the best value for the 2 bit clustering.
76.0%
Top-1 Accuracy
75.0%
74.0%
73.0%
ResNet50 2bit
72.0%
ResNetl1 8 2bit
71.0%
70.0%
69.0%
68.0%
0.0E+00
1.0E-05
2.0E-05
3.0E-05
4.0E-05
5.0E-05
6.0E-05
7.0E-05
T
Figure 6: ResNet18/50 compression using 2 bits with varying values.
In our experiments, we used a binary search to find out the best values w.r.t. the top-1 accuracy, which are listed in 7 In general, observe that complex compression task (i.e., higher
13

---

MobileNet-v2 Published as a conference paper at ICLR 2022
compression targets, more compact networks) tends require a larger to provide enough flexibility or softmess.
For Mobi it requires about 10x larger T values then for ResNetl 18850, 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 fewer centroids to utilize, hence requiring a larger 7 value.
When the centroid dimension increases, the larger value is required, as the compression complexity increases (i.e., need to utilize a longer sequence).
ResNet18
ResNet50
MobileNet-vi
3 bit
8.0e-6
8.0e-6
5.0e-5
5.0e-5
2 bit 2.0e-5
2.0e-5
1.0c-4
1.0c-4
5.0e-5
5.0e-5
3.0e-4
1.5e-4
4/4 bit 5.0e-5
4.0e-5
1.0c-4
1.0c-4
4/8 bit 5.0e-5
5.0e-5
1.0e-4
1.0e-4
8/8 8 bit 8.0e-5
oom
1.0c-4
1.0c-4
8/10 16 bit
1.3e-4
6.0e-5
1.2e-4
1.4e-4
clustering with 4 bits and 4 dimensions
Table 7: for the DKM experiments in Table Din Section
could be possible cast as a learnable parameter for each layer or apply some scheduling to improve the model accuracy further (as future work), but still both approaches need a good initial point which can be found using 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, and 4. 4.0e-4 for MobileBERT. We found the BERT models are less sensitive to the than ImageNet classifiers.
C TRAIN-TIME VS. INFERENCE-TING WEIGHT DIFFERENCE/ERROR
As discussed Section 4.3 DKM- compression requires snap the train-time weights to the nearest centroids for accurate validation or inference, which creates minor accuracy degradation during inference. To understand the behavior better, we measured the Frobenius norm of the weight difference (i.e., warrhhnorm(treeg (train_weight, in inffernnc (feencceereiihh))) weight)) as error for each layer for every batch from ResNet 8 with training with DKM cv:6/8 and fc:6/10 (from Fig. 4 on ImageNetlk. The error changes of the five representative layers for over 120,000 batches or 200 epochs are plotted in Fig. 7 from which we can make the following observations:
Although every layer starts with different level of error, the errors get smaller over the training time, and eventually at the scale of le-4.
Aggressive compression makes layer to begin with high level of emror For example, the last fc layer starts with an error of 4 (because it targets 0.6 bit per weight).
the later layers get stabilized better than the earlier layers after enough epochs have been passed.
Our observations aligned with Fig. 5 (b) in the sense that DKM will encourage weights to be clustered tightly over time, decreasing the difference between trainttimo and inference- time weights, thus can be very effective in compressing the model and minimizing the model accuracy degradation.
D THE NUMBER OF ITERATIONS IN DKM LAYER
In a DKM layer, we iteratively perform & means clustering using the attention mechanism until the clustering process is converged the maximum iteration number is reached. In experiments in
14

---

Published as a conference paper at ICLR 2022
layer1.0.conv
Error
1,00-00
il blu bell w.s.
Batch Index
(a) the first conv layer in the ResBlockl
layer2.0.convl
Error
2,0E:00
1.00 00
1.05-04
Batch Index
(b) the first conv layer in the ResBlock2
layer3.0.convl
Error
2,0H000
1.0E-000
1. 1.00-04
100
1000
10000
Batch Index
(c) the first conv layer in the ResBlock3
layer4.0.convl
Error
2.00-00
1.06:00
1.0F-04
100
Batch Index
(d) the first conv layer in ResBlock4
fc
4,00:000
Error
3,0E=00
2,06=000
Batch Index
(e) the layer ResNet18
Figure 7: The between train-time and inference-time weights Frobenius for every batch for the 200 epochs are plotted, while compressing ResNet18 with DKM cv:6/8 and fc:6/10 on ImageNetlk. After enough number epochs, such error is reduced to the 1e-4 level.
Section 4 we set the maximum number of iteration as 5 to avoid the out-of-memory error. In order to understand how many iterations are needed before each forward pass, we collected the iteration count per batch from each layer for the first eight epochs (the trend holds true for the remaining epochs) from ResNet18 with training with DKM cv:6/8 and fc:6/10 (from Fig. 4 on ImageNetlk, and plotted the graphs in Fig. & Our observations include the following:
15

---

Published as a conference paper at ICLR 2022
layers combined
The Itrnatons
Batch Index
(a) The number of iterations across all layers combined.
I.y.or.
layer 1.0.cony2
The number of Iterations
The number of Iterations
Batch Index
(b) The 1st conv layer in the ResBlock1
(c) The 2nd conv layer in the ResBlock1
layer2.0.conv1
fc
The number of Ieeaaions
The number of Iterations
180
300
2000
200
1400
2000 1380 200
Batch Index
Batch Index
(d) The conv layer ResBlock2
(e) The last layer in ResNet18
Figure 8: The number of iterations in ResNet18 training with DKM cv:6/8 and fc:6/10 ImageNetlk.
(a) shows the number of iteration changes the 8 epochs across all the compressed layers. As one can see, the number of iterations hits the maximum limit initially, implying that the weights are being clustered aggressively.
As the training continues however, the number of iterations decreases slowly with sporadic spikes, implying DKM-based compression helps learn good weight sharing.
While the convolutional layers in (b), (c), and (d) get stabilized after only a few dozens of batches, the last fe layer in (c) requires much more iterations throughout the training. We partially believe that this because we used the DKM 6/10 configuration which more challenging than the DKM 6/8 for the convolution layer.
E IN DKM
The in Fig. determines when to exit the iterative process. Note that the default value for the experiments Section 4 le-4 based on from Alleam..lusterrKMMenn (https://scikit http://wwwweeeeeeeeeeeeeeneeeeee
Therefore, performed the sen-
16

---

Published as a conference paper at ICLR 2022
7.0E400
0.01
0.001
0.0001
5.0E+00
0.00001
4.0E100
0.000001
-0.0000000
3.0E+00
2.0H+00
1.0E+00
10
100
1000
10000
Batch Index
Figure plotted the train vs. inference weight emror from last fe layer ResNet18 with DKM cv:6/8 and fc:6/10 on ImageNetlk with varying from 1e-2 to 1e-7 for 200 epochs. Although the initial emror is larger with bigger e, the final errors all converged to a similar level after a sufficient number of epochs.
sitivity analysis by varying the from le-2 to le-7 for ResNet18 with DKM cv:6/8 and fc:6/10 on ImageNetlk and increase the maximum iteration limit to 15 not to make it bottleneck. From the results, we found that the final Top-1 1 accuracies for all the cases were the range of 65.7 -65.9%. To understand why and develop insights, plotted the train-time
inference-time weight dif- ference of the fe layer in Fig. 9 as similarly with Fig. 7 (d). From the Fig. 9 we could make the following observations:
With a large e, there is wider gap between training and inference weights, implying the clustering is not fully optimized. For example, the initial error with le-2 is about 2x larger than one with 1e-6.
However, the gap is closing over time, and after sufficient number of epochs, the errors are all in the similar range of le-4 regardless of the value.
Larger e values make DKM layers iterate fewer, decreasing the peak memory consumption. For example, the iteration count was throughout training when 1e-2.
Therefore, apparently, long the training with DKM layers can run long enough, the selection of e might not affect the final result. We believe this is because DKM ensures the clustering con- tinuity by resuming from the last known centroid (i.e., from the previous batch). In case that the planned training time is short, smaller value would be preferred but at the cost of larger memory requirement.
F GUMBEL-SOFTMAX FOR SOFT ASSIGNMENT
Gumbel-Softmal distribution is continuous distribution that approximates samples from egorical distribution and also works with back-propagation. Therefore, it is feasible to use the drawing from Gumbel-softmax as soft assignment to generate attention. Therefore. we experi- mented Gumbel--Sfftmx--assse DKM (GBS) along with the hard assignment scheme (Hard) with ResNet18 on ImageNetlk. We ran clustering iteratively for both GBS and Hard, and such iteration helps GBS reduce the variance (Shayer et all 20118).
We applied the same compression configurations from Fig. and kept all hyper-parameters and training flow intact. The comparison results are in Fig. 10 which shows that Hard is much worse than both GBS and DKM on all the cases with over 10% drop in Top-1 accuracy, proving that soft assignment is superior way of clustering weights for model compression. Although GBS outper- forms Hard, GBS is still worse than DKM on cases, and the degradation can as significant as 3.4% drop in Top-1 accuracy for the most aggressive compression target.
17

---

Published as a conference paper at ICLR 2022
DKM ev:5/8,fc:8/4
Compression Ratio
A Hard ev:518,fc:814
GBS ev:5/8,fc:8144
A Haard cv:6//,fc:6/110
40
GBS ev:6/88ffc:6/10
DKM cv:6/8,fcc6/110
DKM cv:4/4,fc:814
30
Hard cv:4/4,
cv:4/4,fe:8/4
48%
50%
52%
54%
56%
58%
60%
62%
64%
66%
68%
Top-1 Accuracy
configuration
cv:4/4,fc:84
cv:6/8,fc:6/100
ev:5/8.c:884
DKM
66.7
65.7
65.1
GBS
65.7
62.9
61.7
Hard
54.9
52.
49.4
Figure 10: When DKM is compared with another soft assignment method, GBS based on Gumbel-softmax and a hard assignment method. Hard with ResNet18 on ImageNetIk, DKM outperforms both. Soft assignment techniques outperform hard wide margin, and DKM is superior to GBS.
G RELATION EXPECTATION--AXIIIIAATOO (EM) AND THEORETICAL INTERPRETATION
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
K
p(r) Σ o' 1=1
the Gaussian mixture model over cluster centers. Referring Fig.
Maximizing log likelihood of weights
In P(W
Σ In
o2
using the EM algorithm equivalent to DKM for the case of dij = (w; e2)
exp( (
The attention matrix A where ai
is equivalent to the responsibilities
Σ. exp( exp( Ti caleulated in the E. step: Tiy Nio
7/2)
dz.j
Updating using Cj Σg o is equivalent to step EM algorithm. Notice that ang
variance is fixed, therefore M step in EM is only updating cluster centers as DKM does.
However, unlike EM where finding every M step is the objective, DKM focuses on generating 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 significantly different. While EM iteratively optimizes a specific likelihood function for a set of fixed observations, DKM needs adjust (i.e,, optimize) the observations (which 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 specific likelihood function (i.e.,
18

---

Published as a conference paper at ICLR 2022
the observations are dynamically changing). Therefore, DKM uses simple softmax and rides on the back-proe propagaaion to fine-tune the observations w.rt. the task loss function after unrolling multiple attention updates. When we propagate gradients, then this will turn into stochastic non- convex joint optimization where we simultaneously optimize observations and centroids for the task loss function, which is shown offer better accuracy vs. compression trade-off's according to experiments.
19


Reviews:
[{"review_id": "jy5W-y7fBRQ", "reviewer": "Reviewer_yNNv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well-written and easy to follow. Using an attention matrix to optimize cluster centers with updates from all the weights is a promising and, to the best of my knowledge, a novel idea. I like Figure 4-b since it shows very clearly that the weight distribution is trained to be centered around clusters as desired. Here are some of questions and concerns:\n\n- The authors claim that weight cluster assignments have not been optimized in prior work. However, although the clusters assignments are not optimized directly through back propagation, there is still an iterative process to optimize the assignments in many model quantization works such as [1] and [2]. I believe making this distinction more clear would be fairer to the prior work.  \n\n-  If I understand it correctly, the proposed method DKM is applied to each layer separately with different weight cluster centers and attention matrices. It seems like one bottleneck for the proposed method is that the storage complexity increases a lot for deep neural networks. Do the authors have an idea of to what extent these cluster centers and/or attention matrices might be shared across layers? Also, the cluster centers and attention matrices change at every iteration (before doing a forward pass). Can the authors provide an approximate number for the number of iterations required before each forward pass (this corresponds to $r$ in the paper)? \n\n- I am a bit confused about the third paragraph in Section 3.3. Multi-Dimensional DKM. The authors make a connection between entropy and model quality and they say \"... its entropy is $b$.\". I am not sure what \"it\" refers to here. I thought $2^b$ was the number of clusters. I do not think entropy is the right term to use here. Can the authors clarify what they mean in that paragraph? \n\n[1] Choi, Yoojin, Mostafa El-Khamy, and Jungwon Lee. \"Towards the limit of network quantization.\" arXiv preprint arXiv:1612.01543 (2016).\n\n[2] Choi, Yoojin, Mostafa El-Khamy, and Jungwon Lee. \"Universal deep neural network compression.\" IEEE Journal of Selected Topics in Signal Processing 14.4 (2020): 715-726.", "summary_of_the_paper": "This paper studies k-means clustering for deep neural network (DNN) compression with several improvements over prior work. The authors provide a differentiable k-means clustering layer (DKM) to optimize the cluster centers and cluster assignments throughout the training without separating the problem of training and compression. ", "summary_of_the_review": "The authors propose a simple but novel clustering method for DNN compression. Overall, I believe it has useful findings but also has some flaws that I mentioned in the previous section. Hence, my score is 6: marginally above the acceptance threshold.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YvR2_KKTQpv", "reviewer": "Reviewer_JTKT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is mostly well written. The experiments on benchmarks are strong and it is reported that architectures that are already resource-efficient (e.g., MobileNet) can be further compressed using the proposed method.\n\nMy main concern is that the paper does not present many insights into the method itself besides good results on benchmark datasets. It is stated that the clustering is adopted to the loss function and not just based on a metric. However, when looking at Figure 2, it appears that the weights are clustered only based on a distance metric. It is stated that it is expected that the weight assignments align well with the loss function. I would require some more intuition, proofs, or experimental evidence to be convinced that this can be expected.\n\nSince the cluster centers are not stored explicitly, they are implicitly stored in the set of weights. Is there some intuition about how the proposed method affects those weights and, therefore, also the implicitly stored cluster centers? Do the weights---before they enter the k-means block---develop some kind of clustering behavior? Or is it just that the weights after the k-means block exhibit clustering.\n\nHow confident (i.e., how close to one) are the weight assignments (i.e., attention) to different cluster centers during training? If a weight gets assigned to multiple clusters, I would expect that the weight used for testing might be quite different than the soft weight assigned during training. How does the accuracy change when going from the soft clustered weights (as used during training) to the hard clustered weights (as used during testing).\n\nHow important is the initialization of the weights using a pre-trained model? How well does the proposed method work when trained from scratch?\n\nI understand that the proposed method aims to reduce the memory footprint. What about the computational overhead at test time? Does the weight sharing induce some computational overhead at test time? Or can the weight sharing even be exploited to compute predictions faster?\n\nThe stopping criterion of k-means appears to be determined by a threshold parameter \\epsilon. How critical is this parameter? How often is k-means iterated on average? Is there a discrepancy in the number of iterations of k-means for different layers? Wouldn't it be better to use a fixed number of iterations to ensure a similar behavior throughout the training and to ensure a fixed memory overhead during training?\n\nMinor:\n- The word \"regress\" was used several times which I confused with \"regression\" (i.e., predicting a real-valued target)\n- Why is DKM sometimes referred to as an \"activation layer\" (e.g., right above Section 3.2). This also confused me since there are weights involved, not activations.", "summary_of_the_paper": "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.", "summary_of_the_review": "Mostly well-written paper with good experimental results on benchmark data, but hardly any insights into the method itself.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LF_sYxu7C3u", "reviewer": "Reviewer_5RxN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes a novel differentiable k-means clustering layer (DKM) for deep neural network model compression. The DKM utilizes attention mechanism to align the weight-to-cluster assignment with the training loss function. Overall, the idea is novel but the paper is not prepared enough. \nThe comments are listed as following: \n\n1.\tIn Figure 2, there are some confusions in logic. To my understanding, the criterion of convergence should be put after the calculation of Attention Matrix   to decide whether to continue update the Centroid Matrix   or not. However, it was put after updating the Centroid Matrix and weight approximation. \n2.\tThe description of multi-dimensional weight-clustering is not so clear. It is better to help to understand by adding a figure with more details. \n3.\tIn Section 3.2, the procedure of DKM is elaborated. However, theoretical interpretation for each step of DKM is also required. \n4.\tSome details should be noticed, please check and correct them: \na)\tThe format of reference should be unified; \nb)\tSome format errors and spelling mistakes exist in the paper; \nc)\tThe Figure 1 and Figure 2 are not clear, I recommend to use vector graphics; \n5.\tIn Section 4, more related clustering-based quantization methods should be added for further improving the reliability of experiments. \n", "summary_of_the_paper": "This paper proposes a novel differentiable k-means clustering layer (DKM) for deep neural network model compression. The DKM utilizes attention mechanism to align the weight-to-cluster assignment with the training loss function. Overall, the idea is novel but the paper is not prepared enough. ", "summary_of_the_review": "This paper proposes a novel differentiable k-means clustering layer (DKM) for deep neural network model compression. The DKM utilizes attention mechanism to align the weight-to-cluster assignment with the training loss function. Overall, the idea is novel but the paper is not prepared enough. \nThe comments are listed as following: \n\n1.\tIn Figure 2, there are some confusions in logic. To my understanding, the criterion of convergence should be put after the calculation of Attention Matrix   to decide whether to continue update the Centroid Matrix   or not. However, it was put after updating the Centroid Matrix and weight approximation. \n2.\tThe description of multi-dimensional weight-clustering is not so clear. It is better to help to understand by adding a figure with more details. \n3.\tIn Section 3.2, the procedure of DKM is elaborated. However, theoretical interpretation for each step of DKM is also required. \n4.\tSome details should be noticed, please check and correct them: \na)\tThe format of reference should be unified; \nb)\tSome format errors and spelling mistakes exist in the paper; \nc)\tThe Figure 1 and Figure 2 are not clear, I recommend to use vector graphics; \n5.\tIn Section 4, more related clustering-based quantization methods should be added for further improving the reliability of experiments. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "H6OcImOkA30", "reviewer": "Reviewer_r68F", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Strength \n- Experimental results are done on important and representative datasets.\n- Easy to understand what the paper is doing. \n\n- Weakness\n\n- I think the paper didn't provide a scientific way of answering the research question. It reads to me that the main problem with existing method is that a discrete, 1 center clustering is the problem so they relax it. If this were true, then I believe the following 2 experiments should also be done:\n\n a) used a fixed ratio from 3 or 4 nearest clusters (each with ratio .33 , .25). This sort of fixed attention has been studied and used in accelerating transformer inference and reported good performance. If the hypothesis in this paper were right, by using this fixed attention should also give good results. \n\n  b) also try to learn it with other relaxation method such as gumbel softmax to see if an improvement can be observed.\n\nI think in general the method proposed in this paper is simple, which is ok to me. But honestly empirically I don't see a great improvement so I am more skeptical to the claim of the root cause. Only after other type of verification of this simple idea is provided I will be recommending this paper. \n\n- Misc\nAlso I don't see the baseline methods I know is included. This paper mainly focused on model compression so speed is not a concern. Thus, I don't think the reason why hessian computation is a problem. I am not peculiarly familiar with quantization based method, but I believe Hawq-v3 (v2 is cited in this paper) is a competitive baseline method to discuss. I'd like to see authors add the experiments of this method and have some meaningful discussions. I will categorize this paper as quantization method so I think this comparison is important to gain more insights.", "summary_of_the_paper": "The paper claims that a competitive branch of model compression method: Weight clustering, can be greatly improved by allowing a soft(or say differentiable) usage of clustering mechanism to achieve a better approximation result.", "summary_of_the_review": "In sum, a very simple idea is proposed. I think simple is good, but empirical results don't quite convince me. I am seeking more ways of validating this idea in order to recommend it.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a simple approach to quantizing neural network weights with encouraging empirical results. The authors did work hard to improve the paper and address reviewers' concerns during the discussion period. I believe the presentation of results can improve by adding a discussion of inference time. I am not sure if all of the baselines (e.g., in Figure 4) have the same inference cost.

PS1: The method does seem to unroll the iterative optimization process (ie. EM) of a Gaussian mixture model (GMM) and differentiates through the unrolled iterations. The paper makes the connection to attention, but does not seem to make a clear connection with GMM and EM. If this connection is correct, adding a discussion can be helpful.

PS2: I am not a big fan of using differentiable k-means as the method name. Differentiable k-means is confusing partly because k-means is differentiable, i.e., one can optimize k-means centers using gradient descent. The proposed approach seems more relevant to meta-learning, where one differentiate though one optimization process to optimize a secondary objective.
Decision: Accept (Poster)