Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
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:

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

-  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)? 

- 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? 

[1] Choi, Yoojin, Mostafa El-Khamy, and Jungwon Lee. "Towards the limit of network quantization." arXiv preprint arXiv:1612.01543 (2016).

[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_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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
main_review: - Strength 
- Experimental results are done on important and representative datasets.
- Easy to understand what the paper is doing. 

- Weakness

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

 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. 

  b) also try to learn it with other relaxation method such as gumbel softmax to see if an improvement can be observed.

I 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. 

- Misc
Also 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_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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
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.

My 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.

Since 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.

How 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).

How important is the initialization of the weights using a pre-trained model? How well does the proposed method work when trained from scratch?

I 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?

The 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?

Minor:
- The word "regress" was used several times which I confused with "regression" (i.e., predicting a real-valued target)
- 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_review: Mostly well-written paper with good experimental results on benchmark data, but hardly any insights into the method itself.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
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. 
The comments are listed as following: 

1.	In 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. 
2.	The description of multi-dimensional weight-clustering is not so clear. It is better to help to understand by adding a figure with more details. 
3.	In Section 3.2, the procedure of DKM is elaborated. However, theoretical interpretation for each step of DKM is also required. 
4.	Some details should be noticed, please check and correct them: 
a)	The format of reference should be unified; 
b)	Some format errors and spelling mistakes exist in the paper; 
c)	The Figure 1 and Figure 2 are not clear, I recommend to use vector graphics; 
5.	In Section 4, more related clustering-based quantization methods should be added for further improving the reliability of experiments. 

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. 
The comments are listed as following: 

1.	In 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. 
2.	The description of multi-dimensional weight-clustering is not so clear. It is better to help to understand by adding a figure with more details. 
3.	In Section 3.2, the procedure of DKM is elaborated. However, theoretical interpretation for each step of DKM is also required. 
4.	Some details should be noticed, please check and correct them: 
a)	The format of reference should be unified; 
b)	Some format errors and spelling mistakes exist in the paper; 
c)	The Figure 1 and Figure 2 are not clear, I recommend to use vector graphics; 
5.	In Section 4, more related clustering-based quantization methods should be added for further improving the reliability of experiments.
review_title: null

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)