Paper: Efficient Unsupervised Knowledge Distillation with Space Similarity (ICLR-2024)

Reviewer comment (Reproducibility):
The reviewer believes that it is unfair to use supervised methods without supervised signals (human annotations) as baselines unless their hyperparameters are tuned without the annotations, but neither hyperparameter tuning nor choices is described in this paper.

Author response:
If the reviewer is referring to supervised methods in experiment 5.7, we would like to apologise for the oversight. As mentioned in the paper, we followed the official implementation of CRD for generating these results. For a fair comparison, for each baseline, we scaled the recommended loss multiplier $\beta$ by factors of **{0.5, 1, 2, 4, 8 ,16}** and selected the value which yielded the highest top-1 for the **Resnet56/20** teacher-student pair. Note that the contribution of supervision and logics based KD is set to 0. Keeping this $\beta$ fixed, we then performed distillation on all the other teacher-student pairs and reported the results. For CoSS, we performed the hyper-param search for **Resnet56/20** and fixed the obtained $\beta=140$  for all other teacher-student pair as well. We will be adding the above mentioned information along with shortlisted $\beta$ for each baseline to A.1 in the newer version. If there is any other additional that the reviewer would like us to cover, we will happily oblige.

Reviewer comment (Reproducibility):
The reviewer has a concern about reproducibility of this work. Appendix A.1 is not detailed enough (how teacher model is trained, architectures of additional trainable layers, baseline hyperparameters, etc) and does not cover all the experiments (at least) in the main body.

Author response:
We would like to clarify the reproducibility aspect of our work. As mentioned in A.1, we have adopted SEED’s training hyper-parameters (teacher, scheduling, epochs, optimizer, augmentations etc) for the ImageNet distillation experiments. The cited paper thoroughly covers the necessary information. However, we would like to apologise for omitting the overall loss scaling value which was learnt from the CIFAR experiments as mentioned in the previous response. We directly apply the scaling factor ($\beta$=140) learnt from CIFAR to the ImageNet experiments as well. Moreover, post decision deadline, we will be sharing the official code for our work.

Reviewer comment (Evaluation):
The reviewer needs more clarifications in the description of the additional evaluation with kNN as it's not convincing. While Section 5.1 explains it is because the approach allows them to evaluate the methods directly without the need of tuning parameters, it is still dependent on the choice of $k$, which is not justified but heuristic.

Author response:
Thank you for the insightful suggestion. kNN performance indeed depends on the choice of the neighbourhood. The choice of $k$, as reported in Section 5.1, is kept to be in-line with the evaluation setting adopted in SEED. The authors of SEED chose the value of $k=10$ in their evaluation. Though this justifies as to why we have used $k=10$, it does not explain as to why only $k=10$ should be preferred. To circumvent this issue, we have reported kNN with k={1,51,101} in A.2.2, where we see a similar trend across different k.

Reviewer comment (Evaluation):
In Section 1, the reduction of GPU memory requirement is emphasized, but the GPU memory requirement is not defined in this paper. The method of measurement is not explained. If the teacher's embeddings are pre-computed and cached, it may easily save GPU memory usage as much as the proposed method does.

Author response:
Thank you for pointing this out. We measured GPU memory as the maximum GPU space occupied by the method while processing an input batch. We reported the mean computed over the first 100 training batches. Training batch size for all methods is the same (as per implementation notes). Caching the representations will only allow us to not load the teacher model onto the GPU for computing embeddings. However, there are few downsides with this.

 (1) It will only save the GPU memory required by the teacher, since the teacher embeddings are required to be on GPU for loss (gradient) computation and hence does not free up GPU space required for embedding themselves.

(2) Teacher Resnet-50 utilises considerably less GPU space (~100MB) when loaded only for inference/embedding generation. Additionally, the cost of a forward pass for this computation is negligible compared to the overall time some methods take.

(3) Each method (including CoSS) loads the teacher onto the GPU, hence if an optimisation can help other methods, it should also help CoSS.

(4) Usage of heavy augmentations like CutMix as in BINGO, makes it hard to pre-process and store all the teacher embeddings and hence requires an online computation of embeddings.     

We will be adding these omitted details to the appendix.

Reviewer comment (Writing):
There are many typos and grammatical errors:
- "distilled student" should be replaced with "trained student", as (knowledge of) teachers are distilled into students, and students are not distilled  
- and Attention transfer -> and attention transfer  
- "(i) feature queues (ii) contrastive objectives (iii) heavy augmentations (iv) and custom batch composition" -> "(i) feature queues, (ii) contrastive objectives, (iii) heavy augmentations, and (iv) custom batch composition"  
- Some notations are not defined where used, e.g., $\widehat{A_t^T}$ and $\widehat{A_s^T}$ in Eq. (4) and $\lambda$ in Eq. (5)  
- DisCo vs. DISCO  
- section vs. Section  
- Table vs table  
- Figure vs. figure  
- In the Appendix -> In Appendix  
- "We compute nearest neighbour in the size 10 neighbourhood of the sample" -> "We choose 10 nearest neighbours from the training samples"  
- Ericsson et. al. Ericsson et al. (2021) -> Ericsson et al. (2021)  
- ImageNet vs. imagenet vs. Imagenet  
- "two settings 1. traditional ... 2. Following" -> "two settings. 1) Traditional ... 2) Following"  
- "as reported in 4" -> "as reported in Table 4"  
- We'd like -> We would like (not an error, but suggested)

Author response:
$\textbf{Typos}$: We once again thank the reviewer for providing extremely thorough feedback on our work. We have addressed and fixed the grammatical errors and typos in our newer draft. We will be uploading it once we have incorporated suggestions from other reviewers as well.

Reviewer comment (Novelty):
The contribution of the method is limited. A similar idea has already been discussed in [1]. For each pair of prediction vectors from the student and teacher within a batch $A^s, A^t \in \mathbb{R}^{b\times d}$, [1] proposed inter-relation loss $\frac{1}{b}\sum^b_{i=1}g(A^s_{i,:}, A^t_{i,:})$, intra-relation loss $\frac{1}{d}\sum^d_{j=1}g(A^s_{:,j}, A^t_{:,j})$ where $g(\cdot,\cdot)$ is a distance function, $b$ is batch size and $d$ is the feature dimensions. [1] employ those loss functions on the logits vectors where $d$ is the number of classes. According to Figure 2, the proposed method is similar to [1]. The difference is the proposed method employs those loss functions on the feature vectors.

Author response:
Thank you for bringing this to our notice. Indeed, [1] has a similar objective where they consider intra-class responses for their supervised setting. One subtle yet important difference in implementations of the two is their treatment of intra-class responses prior to the similarity computation. In [1], **the authors normalise the features prior to computing their inter and intra class losses**. We do not perform this step so as to avoid the loss of original manifold information which a normalisation operation induces. Another important difference is in the motivation and subsequently the setting to which these two strategies are applied. [1] utilises the class-wise scores as means to capture (class) prior from the teacher which consequently is modelled as a correlation maximisation problem. For this purpose, it makes sense to pre-normalize the logits to probability over classes before computing the intra-class loss. On the other hand, we have focused on modelling the manifold space of a teacher for which the space-wise similarity allows us to rectify the loss of information caused by the normalised feature-wise similarity. For this, we do not normalise features prior to computing space similarity. Lastly, we hope to remind the reviewer that unsupervised vs supervised is also usually considered as a distinction. We have added this discussion with [1] in the related work of our new draft.

Reviewer comment (Experiments):
The proposed method is behind some existing methods in the large dataset (ImageNet 1K) in Table 1. There is also an absence of a comparison with state-of-the-art methods, such as SMD[2].

Author response:
Thank you for noticing this, we agree that the proposed method is behind other methods in only terms of top-1 in Table 1. A gentle reminder is that we can achieve this with significantly less computing efforts. For example, BINGO which performs 0.7% better on (Resnet-18) top-1 requires roughly **9x more time** to train. And, for DisCO in case of efficientnet-b0, the GPU compute and training time required increases **~5x**. More importantly, additional experiments in 5.2 and onwards showed that the model from  the proposed method can still serve as an alternative backbone for various tasks where it often performs better than other models. 

Thanks for the suggestion, we didn’t compare SMD because of its different settings. For example, BINGO, DisCo, and SEED are using the same evaluation protocol, while the one SMD uses is different. 

We are currently running distillation on ImageNet using SMD and will report the results in the coming days. Below are the preliminary results on CIFAR-100 for SMD compared with CoSS.

| **Method** | **Resnet20/56** | **Resnet8x4/32x4** | **VGG8/13**   | **WRN16/40**  |
|--------|-------------|----------------|-----------|-----------|
| SMD[2]    | 70.33       | 71.41          | 74.10     | **74.93** |
| Ours   | **71.11**   | **73.90**      | **74.58** | 74.65     |

Reviewer comment (Experiments):
Can authors compare the proposed method with the method mentioned in reference [1]?

Author response:
Thank you for the suggestion. We have listed below the performance of [1] when adapted for an unsupervised distillation setting on CIFAR-100. We compute the losses of [1] over the embeddings features. To highlight the differences to our method, [1] applies the normalisation to the features **before** computing inter and intra category similarities. This pre-normalisation has the same effect of L2 normalisation which causes a loss in original manifold information. 

| **Method** | **Resnet20/56** | **Resnet8x4/32x4** | **VGG8/13**   | **WRN16/40**  |
|--------|-------------|----------------|-----------|-----------|
| DIST[1] | 67.13       |  67.67        | 73.40     | 72.85     |
| Ours   | **71.11**   | **73.90**      | **74.58** | **74.65**     |

Reviewer comment (Theory):
The proposed feature similarity and space similarity is neither well-explained nor intuitive. The authors discuss the reason why they apply normalization on the embedding matrix and use the cosine similiaity instead of the L2 distance in Section 4. The author *treat the embedded manifold as a topological manifold*, and then introduce an argument based on Homeomorphism. Such a conclusion **assumes** that the unsupervised learning methods learn a low dimensional manifold and the manifold is locally euclidean. However, there is a lack of references or theoretical analysis to support their point.

Author response:
> **W1.A**: The proposed feature similarity and space similarity is neither well-explained nor intuitive. The authors discuss the reason why they apply normalisation on the embedding matrix and use the cosine similarity instead of the L2 distance in Section 4 

In section 4, we highlight that normalising the features (L2) alone erases the information present in the pre-normalization embedding manifold. A simple argument for this we presented is that all points lying on a ray through the origin will coincide on the hypersphere. With these collapsed representations we cannot model the teacher's original (unnormalized) manifold which was our original goal. For this purpose, we introduce space similarity. For two points now coinciding on the d-hypersphere will have different contributions to the space similarity loss. The student will learn how data is distributed along different dimensions (up to some scalar). We provide a toy example to help visualise this aspect. In the figure 3 of the paper, we showed that the baseline SEED, though it learns well separated embeddings, is not able to capture the teacher’s embedding local shape as effectively as SEED. 

Lastly, we hope to clarify that for (unit) normalised features, L2 distance and cosine similarity capture similar information. We do not prefer one over the other in our formulation. This can be demonstrated by the following derivation. The Euclidean distance between two vectors $\mathbf{u}$ and $\mathbf{v}$ is given by:
$$D_{\text{Euclidean}}(\mathbf{u}, \mathbf{v}) = \|\mathbf{u} - \mathbf{v}\|_2 $$

For normalized vectors, we have:
$$ \|\mathbf{u} - \mathbf{v}\|_2 = \sqrt{(\mathbf{u} - \mathbf{v})^T (\mathbf{u} - \mathbf{v})} $$
$$ = \sqrt{\|\mathbf{u}\|_2^2 - 2 \mathbf{u}^T \mathbf{v} + \|\mathbf{v}\|_2^2} $$
$$= \sqrt{2 - 2 \cos(\theta)} $$

Here, $\theta$ is the angle between the two vectors which is also captured by cosine similarity. Hence, we can either aim to minimise the distance or maximise the cosine similarity.

> **W1.B**: The author treat the embedded manifold as a topological manifold, and then introduce an argument based on Homeomorphism. Such a conclusion assumes that the unsupervised learning methods learn a low dimensional manifold and the manifold is locally euclidean. However, there is a lack of references or theoretical analysis to support their point.

Thank you for highlighting this issue. Many manifold learning techniques compute local neighbourhood distance with the assumption that the space is locally euclidean[1,2,3]. Explicit statements about the assumption of a locally Euclidean manifold in the context of unsupervised learning can be relatively rare, as this assumption is often fundamental to the underlying methods without being explicitly articulated in the papers. For example, many unsupervised learning methods employ manifold learning based data visualisations which implies that the learnt manifold is locally euclidean [4,5]. We have now added this discussion to Section 4 to strengthen our assumption.

Reviewer comment (Theory):
The designed loss is analogous to contrastive learning, which computes the cosine similarity between two normalized features. However, the discussion does not explain why they only consider the positive samples while neglect the negative samples. I think this may be the key difference from other methods.

Author response:
> **W2**: The designed loss is analogous to contrastive learning, which computes the cosine similarity between two normalised features. However, the discussion does not explain why they only consider the positive samples while neglecting the negative samples. I think this may be the key difference from other methods.

Thank you for providing valuable insights. As noted by the reviewer, our methodology markedly deviates from baselines employing a contrastive learning framework for distillation. While we appreciate the acknowledgment of this distinction, we respectfully disagree with the assertion that our alignment objectives closely resemble conventional contrastive learning techniques. Rather, our approach shares greater affinity with methods which in the reviewer's categorization, exclusively relies on only positive samples [6,7]. An advantage of focusing on positive samples lies in the observation that comparisons with negative samples (as done by the baselines) lead to increase in training time and GPU compute. In summary, our objective focusing only on positive samples is not only more effective, but also more efficient, particularly in scenarios where resources may be constrained.

Reviewer comment (Theory):
It is unclear why the introduced losses discard the negative samples.

Author response:
> **Q2**: It is unclear why the introduced losses discard the negative samples.

The response is provided in **W2** above.

Reviewer comment (Experiments):
In the comparison to CRD, I guess the authors implement the negative contrastive learning. What is the performance of CRD when negative samples are removed?

Author response:
> **Q3**: In the comparison to CRD, I guess the authors implement the negative contrastive learning. What is the performance of CRD when negative samples are removed?

Yes, we utilised the complete objective of CRD which incorporates positives and negatives in the loss computation. Without the contribution of the negatives, CRD loss failed to optimise in our recent experiments achieving classification score of 1\% equivalent to that of random classification.

Reviewer comment (Experiments):
For other comparison methods, what is the performance if the negative samples are removed?

Author response:
> **Q4**: Also, for other comparison methods, what is the performance if the negative samples are removed?

Given that the other methods are also contrastive learning based and removing negatives had a drastic impact on the CRD training. We are currently running a training with SEED’s objective with negative comparisons removed and will provide an update shortly.

Reviewer comment (Experiments):
What is the performance of the proposed model if adding negative samples?

Author response:
> **Q5**: What is the performance of the proposed model if adding negative samples?

As alluded to earlier in our responses, our solution follows an alternative route to that of contrastive learning. We cannot add negatives (samples to increase dissimilarity from)  in the space similarity as this would change the entire paradigm under which we operate. Our intention is not to increase dissimilarity between positive-negative dimensions which is the core idea of contrastive learning.

Reviewer comment (Experiments):
The paper claims efficiency as their advantage. Is it possible to improve the performance using a larger batch size?

Author response:
> **Q6**: The paper claims efficiency as their advantage. Is it possible to improve the performance using a larger batch size?

The ablation over the batch size does not indicate that the performance can be increased by increasing the batch size directly. Please note that we did not perform any hyperparameter optimisation for this study. 

| **Method** | $b=32$ | $b=64$ | $b=128$   | $b=256$  | $b=512$ | $b=1024$ |
|--------|-------------|----------------|-----------|-----------|-----------|-----------|
| Resnet8x4/32x4   |   73.21    |  73.62        |  73.90     | 73.35     | 72.57 |  72.09  |
| Resnet20/56   |   70.34    |  71.32        |  71.11     | 71.05     | 70.59 | 70.11

Reviewer comment (Writing):
Please add more discussion regarding Section 4.

Author response:
> **Q7**: Please add more discussion regarding Section 4.

Thank you for the suggestion. Based on the discussions highlighted in the review, we have added more discussions in section 4 to support our assumptions and intuition. We have also included details about the setup used for the visualisations in Figure 3.

Reviewer comment (Experiments):
Comparison over baseline is lacked.

Author response:
Thank you for the suggestion. We would like to take this opportunity to highlight that we provide an ablation study over $\lambda$ in section A2.1. $\lambda=0$ represents the case of only cosine similarity of features. New experiments on ImageNet ResNet-18 & Eff-b0 with ($\lambda=0$) reveal that we gain roughly $0.9\%$ and $0.8\%$ top-1 by employing space similarity over the baseline. We have now added these results to the appendix.

Reviewer comment (Evaluation):
Performance on CNN and ViT needs more reasonable analysis.

Author response:
Thanks for raising this interesting perspective. Underperforming due to differences in architecture is a widely observed phenomenon in the context of supervised distillation [1]. A reason for the drop is owed to different inductive biases present in the different teacher-student architectures. Contrative objectives encourage the model to learn discriminative features. W.r.t distillation, contrastive objectives encourage the student to learn discriminative features as learnt by their teachers. Given different network capacities and inductive biases, we suspect that these are stronger conditions than our alignment based objective.

Reviewer comment (Experiments):
Sensitivity of hyperparameter.

Author response:
Thank you for the insightful suggestion. Indeed, ideal value for $\lambda$ varies across architectures. We suspect that this is due to differing inductive biases of these architectures. In order to choose the best $\lambda$ for each student-teacher pair, unfortunately, one has to evaluate multiple configurations. However, based on our experiments, $\lambda=0.5$ provides strong results which also extends readily to ImageNet like datasets and for different architectures (Resnet, EfficientNets, ViTs). In addition, we hope to highlight that the additional cost of hyperparameter search is pertinent to all methods. But, given the speed advantage that our method possesses, this search will be relatively faster than baselines’.

Reviewer comment (Evaluation):
More details on the computational efficiency are required.

Author response:
Thank you for the feedback. We have now also added the table reporting the absolute values for GPU utilisation and training time per batch to the appendix. They reflect the values presented in Fig 1 of our paper. Moreover, we have also added relative speed and GPU utilisation (compared with CoSS) to Table 1 as well.

**EDIT:** We will also be reporting total training resources, which is (how much GPU memory is consumed by a method) * (how long the method takes to complete the training). Alongside we will also report the trend of memory and time consumption over the course of training in our paper.

Reviewer comment (Experiments):
The space similarity is like the traditional cosine similarity loss. But the comparison over the cosine similarity is lacked. It is difficult to evaluate the effects of so-called space similarity.

Author response:
Thank you for the suggestion. We would like to take this opportunity to highlight that we provide an ablation study over $\lambda$ in section A2.1. $\lambda=0$ represents the case of only cosine similarity of features. New experiments on ImageNet ResNet-18 & Eff-b0 with ($\lambda=0$) reveal that we gain roughly $0.9\%$ and $0.8\%$ top-1 by employing space similarity over the baseline. We have now added these results to the appendix.

Reviewer comment (Evaluation):
CoSS performs worse on CNN->CNN distillation (Table 1). Could the authors provide more analysis on why the method works well on ViT, but not so good on CNNs?

Author response:
Thanks for raising this interesting perspective. Underperforming due to differences in architecture is a widely observed phenomenon in the context of supervised distillation [1]. A reason for the drop is owed to different inductive biases present in the different teacher-student architectures. Contrative objectives encourage the model to learn discriminative features. W.r.t distillation, contrastive objectives encourage the student to learn discriminative features as learnt by their teachers. Given different network capacities and inductive biases, we suspect that these are stronger conditions than our alignment based objective.

Reviewer comment (Experiments):
The hyperparameter of lambda is somewhat sensitive to different architectures and datasets. How to choose the appropriate lambda needs further discussion. If hyperparameter search is required, additional training cost is required.

Author response:
Thank you for the insightful suggestion. Indeed, ideal value for $\lambda$ varies across architectures. We suspect that this is due to differing inductive biases of these architectures. In order to choose the best $\lambda$ for each student-teacher pair, unfortunately, one has to evaluate multiple configurations. However, based on our experiments, $\lambda=0.5$ provides strong results which also extends readily to ImageNet like datasets and for different architectures (Resnet, EfficientNets, ViTs). In addition, we hope to highlight that the additional cost of hyperparameter search is pertinent to all methods. But, given the speed advantage that our method possesses, this search will be relatively faster than baselines’.

Reviewer comment (Evaluation):
The authors claims that the CoSS is faster and more efficient. Yet the comparative details on computational efficiency is not provided. The analysis on how the method is of high efficiency is also lacked.

Author response:
Thank you for the feedback. We have now also added the table reporting the absolute values for GPU utilisation and training time per batch to the appendix. They reflect the values presented in Fig 1 of our paper. Moreover, we have also added relative speed and GPU utilisation (compared with CoSS) to Table 1 as well.

**EDIT:** We will also be reporting total training resources, which is (how much GPU memory is consumed by a method) * (how long the method takes to complete the training). Alongside we will also report the trend of memory and time consumption over the course of training in our paper.

Reviewer comment (Experiments):
The discussion on topological spaces in Section 4 is commendable. However, it should be noted that L2 normalization in previous approaches is designed to conform the semantic manifold to a hyperspherical space, thereby constraining the metric within the bounds of the cosine similarity. That means the manifold of teacher feature space is already a cosine space and the similarity is determined by the inner product between two hyper-sphere spaces. In light of this, the paper would benefit from a rigorous comparison demonstrating the superiority, if any, of a d-dimensional Euclidean manifold over a d-dimensional hypersphere for the learning tasks at hand.

Author response:
Thank you for pointing this out. As highlighted by the reviewer, similarity matching using only the feature happens in the d-hypersphere with projected embeddings from the student and teachers. We address the issue of information loss from the teacher's unnormalized manifold to a d-hypersphere. This matching on the d-hypersphere does not account for the structure of the original unnormalized manifolds of teacher and student. Through our empirical evaluations presented in the paper, we aimed to highlight the benefit of capturing the original teacher’s manifold information. CoSS is able to yield competitive results, moreover, lagging in top-1 does not hinder its performance on robustness and various transfer learning tasks.

Reviewer comment (Experiments):
This paper could be strengthened by investigating the connection between batch normalization and the proposed CoSS. Can the authors provide insights or very simple results about how BN might influence or interact with CoSS?

Author response:
We thank the reviewer for providing an interesting perspective. Indeed, we can draw a connection between Batch Normalisation (BN) [1] and our final objective.

To recall, during training, BN operates on a batch of input data, $X \in R^{b\times d}$, where batch size is $b$ and $d$ is the feature dimension. It first performs standardisation along each spatial (channel) dimension. $\hat{X_{:,I}} = \frac{X_{i} - \mu_i}{ \sigma_i}$ 

where, $\mu_i$, $\sigma_i$ are the mean and variances respectively for the $i^{th}$ spatial dimension. The normalized values are then scaled by trainable parameters $\gamma$ and $\beta$ as: $Z_{:,i} = \gamma_i \hat{X}_{:,i} + \beta_i$

Here, $\gamma \in R^d$ and $\beta \in R^d$ operate on each spatial dimension. These can be interpreted as affine transformations working on each dimension. We can potentially utilise it to map the standardised student embeddings to the teacher's un-normalized embedding space. The corresponding loss can be defined as follows: $\mathcal{L}= \frac{1}{b}\sum_{i=0}^{i<b}\mathcal{D}(Z^s_i, X^t_i) $

where, $X^t_i$ is the teacher's embedding for the $i^{th}$ sample and $Z^s_i$ corresponds to the BatchNormalized student's embeddings for the same sample. Below we share empirical findings with the BatchNorm formulation on CIFAR-100.
 
| **Method** | **VGG8/13** | **Resnet8x/32x** | **WRN16/40**  |
|--------|-------------|----------------|-----------|
| BN         |  74.01      | 72.22            |  73.42         |
| Ours       | 74.58       | 73.90            | 74.65         |

Reviewer comment (Theory):
In my view, to some degree, the joint constraint in CoSS appears to share conceptual relations with the optimal transport (Sinkhorn function). Could the authors elaborate on any theoretical underpinnings or empirical evidence that supports this connection?

Author response:
We thank the reviewer for the valuable suggestion. Unfortunately, due to limited familiarity with the topic, we are unable to provide any meaningful insights on the question raised by the reviewer. However, we are willing to explore this direction in the future.