K L<EMAIL_ADDRESS>Soroush Abbasi<EMAIL_ADDRESS>Ajinkya<EMAIL_ADDRESS>Hamed<EMAIL_ADDRESS>University of Maryland, Baltimore County
Maryland, USA University of California, Davis
California, USA Regressed SSL distillation

# SimReg: Regression as a Simple Yet Effective Tool for Self-supervised
Knowledge Distillation

###### Abstract

Feature regression is a simple way to distill large neural network models to
smaller ones. We show that with simple changes to the network architecture,
regression can outperform more complex state-of-the-art approaches for
knowledge distillation from self-supervised models. Surprisingly, the addition
of a multi-layer perceptron head to the CNN backbone is beneficial even if
used only during distillation and discarded in the downstream task. Deeper
non-linear projections can thus be used to accurately mimic the teacher
without changing inference architecture and time. Moreover, we utilize
independent projection heads to simultaneously distill multiple teacher
networks. We also find that using the same weakly augmented image as input for
both teacher and student networks aids distillation. Experiments on ImageNet
dataset demonstrate the efficacy of the proposed changes in various self-
supervised distillation settings. Code is available at
https://github.com/UCDvision/simreg

##  1 Introduction

There has been a tremendous improvement in deep learning methodologies and
architectures in the last few years. While this has lead to significant
improvements in performance on various computer vision tasks, it has also
resulted in complex and deep networks that require high compute during
inference [He et al.(2016)He, Zhang, Ren, and Sun, Huang et al.(2017)Huang,
Liu, Van Der Maaten, and Weinberger, Huang et al.(2019)Huang, Cheng, Bapna,
Firat, Chen, Chen, Lee, Ngiam, Le, Wu, et al., Zagoruyko and Komodakis(2016)].
Various specialized architectures [Tan and Le(2019), Iandola et
al.(2016)Iandola, Han, Moskewicz, Ashraf, Dally, and Keutzer, Howard et
al.(2017)Howard, Zhu, Chen, Kalenichenko, Wang, Weyand, Andreetto, and Adam,
Sandler et al.(2018)Sandler, Howard, Zhu, Zhmoginov, and Chen] have been
proposed to minimize the inference time and memory requirements of the model
to be deployed. Knowledge distillation [Buciluǎ et al.(2006)Buciluǎ, Caruana,
and Niculescu-Mizil, Hinton et al.(2014)Hinton, Vinyals, and Dean] has been
proposed as an effective technique to compress information from larger but
effective models (teachers) to lighter ones (students).

With availability of large scale unlabeled datasets, self-supervised learning
(SSL) has received great attention in recent times. Several SSL methods
achieve close to supervised performance on the benchmark ImageNet object
classification task [Caron et al.(2018)Caron, Bojanowski, Joulin, and Douze,
Grill et al.(2020)Grill, Strub, Altché, Tallec, Richemond, Buchatskaya,
Doersch, Pires, Guo, Azar, et al.]. Unlike supervised models, the outputs of a
self-supervised network are latent feature vectors and not class
probabilities. An additional module is generally trained atop the pretrained
SSL models using supervision to perform the downstream task. Conventional
knowledge distillation methods proposed for supervised classification are thus
not applicable for distillation from self-supervised networks. A simple way to
handle this is to directly regress the teacher latent features. Recent works
[Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash, Fang et
al.(2021)Fang, Wang, Wang, Zhang, Yang, and Liu] have proposed more complex
solutions that try to capture the structure of the teacher latent space and
are shown to outperform the regression baselines.

Figure 1: Proposed distillation pipeline: We propose a simple modification of
using a MLP prediction module during distillation. The module is discarded
during inference. Surprisingly, performance of backbone features $f_{s}$ is
better than those from MLP output, $f_{s}^{\prime}$, though $f_{s}^{\prime}$
more closely matches the teacher. A deeper MLP helps improve distillation
performance.

Use of a multi-layer perceptron (MLP) head atop CNN backbone model has been
shown to help self-supervised models prevent overfitting to the SSL task and
generalize better to downstream applications [Chen et al.(2020a)Chen,
Kornblith, Norouzi, and Hinton, Chen et al.(2020b)Chen, Fan, Girshick, and He,
Grill et al.(2020)Grill, Strub, Altché, Tallec, Richemond, Buchatskaya,
Doersch, Pires, Guo, Azar, et al., Chen and He(2021)]. Such modules are used
only during SSL pre-training and are not part of inference network. In this
work, we consider the task of distilling self-supervised models. We employ a
similar prediction head atop the student backbone network to effectively mimic
the teacher. As in SSL, the prediction module is discarded after distillation
and thus, there is no change in the time and memory required during inference
(refer Fig. 1). We empirically demonstrate that doing so does not hurt
classification performance. Counter-intuitively, we observe that the features
from the backbone network outperform those from the final layer of the
prediction head though the final layer best matches the teacher. Unlike in
SSL, overfitting to the training task (i.e, exactly mimicking the teacher)
benefits distillation [Beyer et al.(2021)Beyer, Zhai, Royer, Markeeva, Anil,
and Kolesnikov] and it is not clear why generalization could be better at
intermediate layers where the similarity with teacher is reduced. Our finding
suggests that we require a deeper analysis to understand how well the student
models mimic the teacher in general and how knowledge distillation works.
Crucially, it also enables us to use a deeper prediction head to achieve lower
train and test error leading to better downstream performance without
increasing the student capacity.

We empirically show that the above observation generalises to distillation
with different teacher and student settings and to other self-supervised
distillation technique. Our simple regression model with a MLP prediction head
outperforms complex state-of-the-art approaches that require the use of memory
banks and tuning of temperature parameter. Our work serves as an important
benchmark for future self-supervised distillation works. The use of MLP heads
also facilitates effective distillation from multiple SSL teacher networks.
Additionally, we demonstrate that using the same augmented image with weak
augmentation for both student and teacher networks results in better student
models. Since aggressive augmentation is necessary for effective self-
supervised learning [Chen et al.(2020a)Chen, Kornblith, Norouzi, and Hinton,
Chen et al.(2020b)Chen, Fan, Girshick, and He] but hurts their ability to
generalize [Purushwalkam Shiva Prakash and Gupta(2020)], our approach could be
used to learn better SSL models. To summarize, our contributions are simple
changes to architecture and augmentation strategy of distillation networks
that not only achieve state-of-the-art performance on SSL model distillation
but also question our current understanding of knowledge distillation.

##  2 Related Works

Supervised knowledge distillation:  Bucilua _et al_ [Buciluǎ et
al.(2006)Buciluǎ, Caruana, and Niculescu-Mizil] and Hinton _et al_ [Hinton et
al.(2014)Hinton, Vinyals, and Dean] pioneered the use of knowledge
distillation for compressing information. The methods used the teacher
prediction logits as soft-labels in addition to the supervised label to
regularize the student model. [Passalis and Tefas(2018)] minimizes the
divergence between the student and teacher probability distributions. Several
works ([Romero et al.(2015)Romero, Ballas, Kahou, Chassang, Gatta, and Bengio,
Komodakis and Zagoruyko(2017), Heo et al.(2019)Heo, Lee, Yun, and Choi])
utilize intermediate teacher outputs in distillation. FitNets [Romero et
al.(2015)Romero, Ballas, Kahou, Chassang, Gatta, and Bengio] match both the
final and intermediate teacher representations while [Komodakis and
Zagoruyko(2017)] transfers knowledge from the attention maps of the teacher.
RKD [Park et al.(2019)Park, Kim, Lu, and Cho] transfers mutual relations
instead of instance wise distillation. [Wang et al.(2020)Wang, Ge, and Wu]
proposes directly regressing the final teacher features with a modified loss
function that strictly matches the direction of the features but allows
flexibility in terms of feature magnitude.

Self-supervised representation learning (SSL): Earlier works on SSL ([Doersch
et al.(2015)Doersch, Gupta, and Efros, Gidaris et al.(2018)Gidaris, Singh, and
Komodakis, Caron et al.(2018)Caron, Bojanowski, Joulin, and Douze, Noroozi and
Favaro(2016), Noroozi et al.(2017)Noroozi, Pirsiavash, and Favaro, Pathak et
al.(2016)Pathak, Krahenbuhl, Donahue, Darrell, and Efros]) learn effective
representations by solving pretext tasks that do not require supervised
labels. Recently, works based on contrastive learning ([Bachman et
al.(2019)Bachman, Hjelm, and Buchwalter, He et al.(2020)He, Fan, Wu, Xie, and
Girshick, Chen et al.(2020b)Chen, Fan, Girshick, and He, Tian et
al.(2019)Tian, Krishnan, and Isola, Chen et al.(2020a)Chen, Kornblith,
Norouzi, and Hinton, Misra and Maaten(2020), Caron et al.(2020)Caron, Misra,
Mairal, Goyal, Bojanowski, and Joulin]) have gained focus. In contrastive
learning, the distances between representations of positive pairs are
minimized while those between negative pairs are maximized. The positive and
negative pairs are generally constructed by utilizing multiple augmentations
of each image. BYOL [Grill et al.(2020)Grill, Strub, Altché, Tallec,
Richemond, Buchatskaya, Doersch, Pires, Guo, Azar, et al.] is closer to
knowledge distillation, where the distance between teacher and student
representations are minimized. The inputs to the two networks must be
different augmentations of the same image and the teacher network is obtained
as a moving average of the student. Similar to our work, [Grill et
al.(2020)Grill, Strub, Altché, Tallec, Richemond, Buchatskaya, Doersch, Pires,
Guo, Azar, et al.] employs MLP head atop the student network to predict the
teacher features.

Distillation of self-supervised models: In [Noroozi et al.(2018)Noroozi,
Vinjimoor, Favaro, and Pirsiavash], the student mimics the unsupervised
cluster labels predicted by the teacher. CRD [Tian et al.(2020)Tian, Krishnan,
and Isola] maximizes a lower bound of the mutual information between the
teacher and student networks. However, it additionally uses supervised loss
for optimization. CompRess [Koohpayegani et al.(2020)Koohpayegani, Tejankar,
and Pirsiavash] and SEED [Fang et al.(2021)Fang, Wang, Wang, Zhang, Yang, and
Liu] are specifically designed for compressing self-supervised models. In both
these works, student mimics the relative distances of teacher over a set of
anchor points. Thus, they require maintaining large memory banks of anchor
features and tuning temperature parameters. As in regression, proposed
prediction heads can also be used to improve CompRess and SEED.

##  3 Knowledge Distillation

We first consider the supervised model distillation formulation proposed in
[Hinton et al.(2014)Hinton, Vinyals, and Dean]. The teacher is trained on the
task of object classification from images. Let $X$ be the set of images, $Y$
the set of corresponding class labels and $c$ the total number of classes.
Consider a teacher network $\mathcal{T}$ with $f_{t}=\mathcal{T}(x)$,
$f_{t}\in\mathbb{R}^{c}$ as the output vector (logits) corresponding to input
image $x$. The predicted class probabilities can be obtained by applying
softmax operation $\sigma(.)$ atop the vector $f_{t}$.

$\hat{y}_{t}=\sigma(f_{t};\tau_{t})=\frac{e^{f_{t}/\tau_{t}}}{\sum_{i}e^{f_{t}^{i}/\tau_{t}}}$
(1)

where $f_{t}^{i}$ is the $i^{\textrm{th}}$ dimensional output of the feature
vector and $\tau_{t}$ is the temperature parameter. The teacher network
$\mathcal{T}$ is trained using the image-label pairs in a supervised fashion
with standard cross-entropy loss. The trained teacher network is to be
distilled to a student network. Once trained, the teacher network parameters
are frozen during the distillation process. Let $\mathcal{S}$ be the student
network, $f_{s}=\mathcal{S}(x)$, $f_{s}\in\mathbb{R}^{c}$ the feature vector
corresponding to input image $x$ and $\hat{y}_{s}=\sigma(f_{s};\tau_{s})$ the
predicted class probability vector. Knowledge distillation loss is given by

$\mathcal{L}_{\textrm{KD}}(\hat{y}_{t},\hat{y}_{s})=\sum_{j=1}^{c}\hat{y}_{t}^{j}\,\textrm{log}(\hat{y}_{s}^{j})$
(2)

The student is trained using a combined objective function involving
supervised cross entropy loss on student features $\mathcal{L}_{\textrm{CE}}$
and distillation loss $\mathcal{L}_{\textrm{KD}}$:

$\mathcal{L}=\lambda\mathcal{L}_{\textrm{CE}}+(1-\lambda)\tau_{s}^{2}\mathcal{L}_{\textrm{KD}}$
(3)

where $\lambda$ is a hyperparameter that determines the relative importance of
each loss term. Since $\tau_{t}$ is generally set to 1, KD loss is multiplied
by a factor of $\tau_{s}^{2}$ to match the scale of gradients from both loss
terms.

###  3.1 Distillation of Self-supervised Models

The value of $\lambda$ in Eq. 3 can be set to 0 if the class labels are not
available during student distillation. However, the formulation cannot be
directly employed to distill from self-supervised teacher networks since the
teacher outputs are latent representations and not logits or class probability
vectors. Thus, to distill from such teachers, we simply regress the final
feature vector of the teacher. Let $f_{t}=\mathcal{T}(x)$,
$f_{t}\in\mathbb{R}^{d}$ and $f_{s}=\mathcal{S}(x)$, $f_{s}\in\mathbb{R}^{m}$.
Since it is not necessary for the student and teacher representation
dimensions to be the same, we use a linear projection of the student feature
to match the dimensions.

$f_{s}^{\prime}=W^{T}f_{s}+b;\,\,W\in\mathbb{R}^{m}\times\mathbb{R}^{d},\,b\in\mathbb{R}^{d}$
(4)

The distillation objective is then given by
$\mathcal{L}=\mathcal{L}_{\textrm{reg}}=\textrm{d}(f_{t},f_{s}^{\prime})$
where $\textrm{d}(.)$ is a distance metric. Here, we consider squared
Euclidean distance of $l_{2}$ normalized features as the metric.

###  3.2 Prediction Heads for Regression Based Distillation

For a more effective matching of the teacher latent space, we propose a non-
linear prediction head $g(.)$ atop the student backbone network $\mathcal{S}$
in place of the linear projection in Eq. 4. During training, the student
feature is then obtained as $f_{s}^{\prime}=g(f_{s})$ where $g(.)$ is modeled
using a multi-layer perceptron (MLP). Each layer in $g(.)$ is given by a
linear layer with bias followed by batch-normalization and a non-linear
activation function (we use ReLU non-linearity in all our experiments). The
number of such layers is a hyperparameter to be optimized. The dimension of
the last layer output matches that of the teacher. There is no non-linearity
in the final layer to prevent constraining the output space of the student
network. During inference, the prediction head $g(.)$ is removed and the
output of the student network is obtained as $f_{s}=\mathcal{S}(x)$ (refer
Fig. 1). Thus, there is no change in the architecture or the number of
parameters of the model to be deployed. In our experiments, we demonstrate
that the use of such MLP heads plays a crucial role in improving downstream
performance. Surprisingly, we also observe that preserving the prediction
heads during inference is not necessarily beneficial and might result in
reduction in performance.

###  3.3 Multi-teacher Distillation

The prediction heads are particularly beneficial in distillation from multiple
teacher networks. Independent deep non-linear projections of the student
backbone features can be employed during distillation to match each of the
teachers. Let $f_{t}^{k}$ be the output vector of the $k^{th}$ teacher and
$f_{s}^{k}=g^{k}(f_{s})$ that of the corresponding student prediction head
$g^{k}(.)$. The multi-teacher distillation objective for $K$ teachers is given
by:

$\mathcal{L}=\frac{1}{K}\sum_{k}\textrm{d}(f_{t}^{k},f_{s}^{k})$ (5)

The prediction heads are trained by the loss term from corresponding teachers
while the backbone $S$ is trained using the summation in Eq. 5.

##  4 Experiments

We consider distillation of pretrained self-supervised models. We consider
four such methods for teacher networks - MoCo-v2 [Chen et al.(2020b)Chen, Fan,
Girshick, and He], BYOL [Grill et al.(2020)Grill, Strub, Altché, Tallec,
Richemond, Buchatskaya, Doersch, Pires, Guo, Azar, et al.], SwAV [Caron et
al.(2018)Caron, Bojanowski, Joulin, and Douze] and SimCLR [Chen et
al.(2020a)Chen, Kornblith, Norouzi, and Hinton]. We use the official publicly
released models for all the teacher networks (details in suppl.). We also use
a ResNet-50 model trained with supervised labels (provided by PyTorch in
[pyt(b)]) as a teacher. All teacher training and student distillation is
performed on the train set of ImageNet. We consider different teacher and
student backbone network architectures. For the prediction head, we experiment
with linear, 2 and 4 layer MLPs. Let the dimension of the student backbone
output be $m$ and that of teacher $d$. Similar to the prediction head in
[Grill et al.(2020)Grill, Strub, Altché, Tallec, Richemond, Buchatskaya,
Doersch, Pires, Guo, Azar, et al.], the MLP dimensions are $(m,2m,m,2m,d)$.

Implementation details: We use SGD optimizer with cosine scheduling of
learning rate and momentum of 0.9. Initial learning rate is set to 0.05. As in
[Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash] the networks
are trained for 130 epochs with batch size of 256. Cached teacher features are
utilized for faster distillation in experiments with SimCLR, BYOL and SwAV
teachers. We publicly release the code111Code is available at
https://github.com/UCDvision/simreg.

Datasets: We primarily evaluate the performance of distilled networks on
ImageNet [Russakovsky et al.(2015)Russakovsky, Deng, Su, Krause, Satheesh, Ma,
Huang, Karpathy, Khosla, Bernstein, et al.] classification task. Additionally,
for transfer performance evaluation, we consider the following datasets:
Food101 [Bossard et al.(2014)Bossard, Guillaumin, and Van Gool], CIFAR10
[Krizhevsky et al.(2009)Krizhevsky, Hinton, et al.], CIFAR100 [Krizhevsky et
al.(2009)Krizhevsky, Hinton, et al.], SUN397 [Xiao et al.(2010)Xiao, Hays,
Ehinger, Oliva, and Torralba], Cars [Krause et al.(2013)Krause, Stark, Deng,
and Fei-Fei], Aircraft [Maji et al.(2013)Maji, Rahtu, Kannala, Blaschko, and
Vedaldi], DTD [Cimpoi et al.(2014)Cimpoi, Maji, Kokkinos, Mohamed, and
Vedaldi], Pets [Parkhi et al.(2012)Parkhi, Vedaldi, Zisserman, and Jawahar],
Caltech-101 [Fei-Fei et al.(2004)Fei-Fei, Fergus, and Perona] and Flowers
[Nilsback and Zisserman(2008)]. We train a single linear layer atop the frozen
backbone network for transfer evaluation (refer suppl.)

Metrics: We use k-nearest neighbour (k-NN) and linear evaluation on all tasks.
We also report mean squared error (MSE) between the student and teacher
features over the test set. For k-NN evaluation, k=1 and 20 are considered and
cosine similarity is used to calculate NNs. We employ FAISS [fai()] GPU
library to perform fast k-NN evaluation. For linear evaluation, a single
linear layer is trained atop the features from the network to be evaluated. As
in [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash], the
inputs to the linear layer are normalized to unit $l_{2}$ norm and then each
dimension is shifted and scaled to have unit mean and zero variance. The layer
is trained for 40 epochs using SGD with learning rate of 0.01 and momentum of
0.9.

Train and Inference Arch | 1-NN | 20-NN | Linear | MSE
---|---|---|---|---
MobileNet-v2+4L-MLP | 54.5 | 58.7 | 68.5 | 0.090
MobileNet-v2+2L-MLP | 54.0 | 58.0 | 67.9 | 0.097
MobileNet-v2+Linear | 50.8 | 55.1 | 58.3 | 0.149
Table 1: Role of MLP Heads. We train three models with varying number of layers in prediction head and use the features from final MLP layer of each prediction model for evaluation. Deeper models more closely match the teacher (lower MSE) and achieve better classification performance (1 and 20 Nearest Neighbour and Linear evaluation). Train | Backbone(BB)+4L-MLP
---|---
Inference | Backbone(BB) | BB+2L-MLP | BB+4L-MLP
Metric | 1-NN | Linear | MSE | 1-NN | Linear | MSE | 1-NN | Linear | MSE
ResNet-18 | 55.3 | 65.7 | - | 56.0 | 66.4 | 1.99 | 53.4 | 65.2 | 0.1
Table 2: Effect of MLP Heads on inference. We train a single model with 4
layer MLP head and perform evaluation using features from different layers
(pre-MLP, intermediate MLP layer and MLP output). Since the final layer
outputs are trained to mimic the teacher, MSE with teacher features is lowest
at 4L-MLP while that at 2L-MLP is extremely high (dimension of BB and teacher
are different, hence MSE is not reported). However, features from backbone and
intermediate layer (+2L-MLP) outperform those from final layer (+4L-MLP) on
classification, contrary to the notion that features with lower MSE generalize
better.

###  4.1 Baseline Approaches

Regression: In addition to proposed MLP prediction head based regression
(termed SimReg-MLP), we consider two additional regression baseline methods
proposed in [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash]
termed ‘Regress’ and ‘Regress-BN’. While Regress distills from unnormalized
teacher features, Regress-BN uses batch-norm layer atop the final student and
teacher features during distillation. Unlike Regress-MLP, both these
approaches use a linear prediction head.

CompRess:  CompRess [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and
Pirsiavash] is designed to distill specifically from self-supervised models.
Given a set of anchor points, the student is encouraged to have the same
similarities with the anchors as that of the teacher. The anchor point
features can either be common features from a teacher memory bank
(CompRess-1q) or features from individual memory banks for teacher and student
(CompRess-2q). We additionally implement CompRess with our MLP prediction
head, termed CompRess-1q-MLP and CompRess-2q-MLP. SEED [Fang et al.(2021)Fang,
Wang, Wang, Zhang, Yang, and Liu] proposes similarity based distillation
similar to [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash]
but uses pre-trained teacher models with significantly lower number of
training epochs and performance. Further, it requires access to the projection
heads used atop teacher networks used only during SSL training and not
inference, These parameters are generally not publicly released, making the
setting less replicable. Thus we provide comparisons with only CompRess
[Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash].

Contrastive Representation Distillation (CRD):  CRD [Tian et al.(2020)Tian,
Krishnan, and Isola] uses a contrastive formulation to bring corresponding
teacher and student features closer while pushing apart those from unrelated
pairs. While the paper considered a supervised setup and loss term utilizing
labels, we use the formulation with just the contrastive loss as proposed in
[Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash].

Cluster Classification (CC):  In CC [Noroozi et al.(2018)Noroozi, Vinjimoor,
Favaro, and Pirsiavash], the student predicts unsupervised labels obtained by
clustering samples using teacher features. We report metrics for CC and CRD
from [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash].

##  5 Results

Teacher | MoCo-v2 ResNet-50 | MoCo-v2 ResNet-50 | SimCLR ResNet-50x4
---|---|---|---
Student | MobileNet-v2 | ResNet-18 | ResNet-50
Method | 1-NN | Linear | 1-NN | Linear | 1-NN | Linear
Teacher* | 57.3 | 70.8 | 57.3 | 70.8 | 64.5 | 75.6
Supervised* | 64.9 | 71.9 | 63.0 | 69.8 | 71.4 | 76.2
Regress* | 38.6 | 48.0 | 41.7 | 52.2 | - | -
Regress-BN* | 48.7 | 62.3 | 47.3 | 58.2 | - | -
CC [Noroozi et al.(2018)Noroozi, Vinjimoor, Favaro, and Pirsiavash]* | 50.2 | 59.2 | 51.1 | 61.1 | 55.6 | 68.9
CRD [Tian et al.(2020)Tian, Krishnan, and Isola]* | 36.0 | 54.1 | 43.7 | 58.4 | - | -
CompRess-2q [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash]* | 54.4 | 63.0 | 53.4 | 61.7 | 63.0 | 71.0
CompRess-1q [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash]* | 54.8 | 65.8 | 53.5 | 62.6 | 63.3 | 71.9
CompRess-2q-4L-MLP | 56.3 | 67.4 | 54.4 | 64.0 | 62.5 | 73.5
CompRess-1q-4L-MLP | 55.5 | 67.1 | 54.9 | 64.6 | 60.9 | 72.9
SimReg-4L-MLP | 55.5 | 69.1 | 54.8 | 65.1 | 60.3 | 74.2

Table 3: Comparison of SSL distillation methods on ImageNet classification.
Our regression method with MLP head (SimReg-4L-MLP) is comparable to or better
than the complex state-of-the-art approaches, especially on the linear
evaluation metric. We also observe that CompRess-1q and 2q are improved when
MLP heads are utilized. Interestingly, regression gets a significantly higher
boost compared to CompRess upon addition of MLP layers. Note that the MLPs are
used only during training and the inference network architecture remains the
same for all approaches making the comparison fair. * metrics from CompRess
[Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash].

Role of Prediction Head: A deeper prediction head results in a student with
higher representational capacity and thus a model that better matches the
teacher representations. Table 1 shows results for models with a common
MobileNet-v2 [Sandler et al.(2018)Sandler, Howard, Zhu, Zhmoginov, and Chen]
backbone and different prediction head architectures. The prediction head is
used during both student training and evaluation. We observe that a deeper
model has lower MSE with teacher features and better classification
performance. However, a deeper model also implies greater inference time and
memory requirements. The student architecture is fixed based on deployment
needs and thus requirement of larger model goes against the very essence of
distillation. To analyze performance at different layers of the prediction
head, we train a single ResNet-18 [He et al.(2016)He, Zhang, Ren, and Sun]
student with all intermediate dimensions of MLP equal to that of the output.
Surprisingly, a model trained with MLP prediction head performs well on
downstream task even when the prediction head is discarded during inference
(Table 2). The performance using features from backbone network is slightly
better than that from the final layer outputs whenever a MLP head is used
(more results in suppl.). More importantly, this observation enables us to use
deeper prediction heads for distillation in place of linear layers without any
concerns about altering the student architecture or increasing inference time.

Comparison with existing approaches: In all the remaining experiments, we use
SimReg-4L-MLP with the prediction head used only during distillation. We
compare the proposed regression method with other baselines and self-
supervised distillation methods in tables 3 and 4. Surprisingly, our simple
regression performs comparably or even outperforms the state-of-the-approaches
on all settings and metrics. On linear evaluation, we outperform previous
methods (without MLP) by 3.3, 2.5 and 2.3 points respectively on MobileNet-v2,
ResNet-18 and ResNet-50 students. Our observation generalizes to similarity
based distillation too. Use of MLP prediction head also consistently improves
the classification performance of both the CompRess variants (table 3). Note
that the linear metrics of our student model on MobileNet-v2 and ResNet-50 are
just 1.7 and 1.4 points below the corresponding teacher accuracies. Our
ResNet-50 model distilled from SimCLR teacher outperforms a ResNet-50 model
trained from scratch using SimCLR (69.3% [Chen et al.(2020a)Chen, Kornblith,
Norouzi, and Hinton]) by 4.9 points.

Teacher | BYOL ResNet-50 | SwAV ResNet-50
---|---|---
Method | 1-NN | 20-NN | Linear | 1-NN | 20-NN | Linear
Teacher | 62.8 | 66.8 | 74.3 | 60.7 | 64.8 | 75.6
CompRess-2q-4L-MLP | 56.0 | 60.6 | 65.2 | 53.2 | 58.1 | 63.9
CompRess-1q-4L-MLP | 55.4 | 60.0 | 65.2 | 52.4 | 57.1 | 63.4
SimReg-4L-MLP | 56.7 | 61.6 | 66.8 | 54.0 | 59.3 | 65.8
Table 4: ImageNet Evaluation with different teacher networks. We distill from
two pretrained ResNet-50 SSL models, BYOL and SwAV to ResNet-18 students. When
distilled from these stronger teacher networks, SimReg is significantly better
than both CompRess variants on all metrics. Both SimReg and CompRess contain
MLP head only during training.

Transfer Learning: Since an important goal of self-supervised learning is to
learn models that generalize well to new datasets and tasks, we evaluate the
transfer learning performance of our distilled networks. The results in table
5 suggest that the regression model transfers as well as or better than the
state-of-the-approaches on most datasets. Among CompRess variants,
CompRess-2q-MLP is generally better on ImageNet classification (table 3) but
transfers poorly (table 5) compared to CompRess-1q-MLP. However, the same
SimReg model performs comparably or outperforms them both in ImageNet and
transfer tasks.

Arch | ResNet-50 | MobileNet-v2 | ResNet-18
---|---|---|---
Method | Teacher | Comp-2q | Comp-1q | SimReg | Comp-2q | Comp-1q | SimReg
|  | -4L-MLP | -4L-MLP | -4L-MLP | -4L-MLP | -4L-MLP | -4L-MLP
Food | 72.3 | 71.4 | 72.5 | 73.1 | 61.7 | 65.9 | 65.4
CIFAR10 | 92.2 | 90.3 | 90.4 | 91.2 | 87.3 | 89.3 | 88.6
CIFAR100 | 75.1 | 73.9 | 74.5 | 76.1 | 68.4 | 71.9 | 70.2
SUN | 60.2 | 58.0 | 58.1 | 59.4 | 54.3 | 56.0 | 57.1
Cars | 50.8 | 60.3 | 63.1 | 62.4 | 37.2 | 44.1 | 42.3
Aircraft | 53.5 | 57.7 | 59.7 | 58.7 | 42.3 | 47.8 | 45.8
DTD | 75.1 | 71.7 | 71.3 | 74.5 | 69.3 | 71.2 | 70.9
Pets | 83.6 | 86.7 | 86.3 | 85.6 | 84.0 | 84.4 | 83.9
Caltech | 89.3 | 91.1 | 91.5 | 91.7 | 87.3 | 90.1 | 89.2
Flowers | 91.3 | 94.3 | 95.4 | 95.1 | 86.4 | 91.3 | 90.9
Mean | 74.3 | 75.5 | 76.3 | 76.8 | 67.8 | 71.2 | 70.4

Table 5: Transfer learning results on multiple classification tasks. Since the
teacher networks are self-supervised, generalization of learnt features to
other datasets is important. SimReg is significantly better than CompRess-2q
and comparable to CompRess-1q on most datasets. All methods employ 4 layer MLP
heads only during distillation.

In addition to transfer learning on different datasets on the classification
task, we consider the task of object detection. We distill a ResNet-50 teacher
to multiple ResNet-18 student networks with different MLP heads. The MLP heads
are not part of the model fine-tuned for detection. Following [He et
al.(2020)He, Fan, Wu, Xie, and Girshick], we use Faster-RCNN [Ren et
al.(2015)Ren, He, Girshick, and Sun] with R18-C4 backbone architecture for the
detection task. All methods are trained on the VOC trainval07+12 and tested on
test07 subsets of the PASCAL VOC [Everingham et al.(2010)Everingham, Van Gool,
Williams, Winn, and Zisserman] dataset using the code from [pyt(a)]. We report
the standard AP50, AP(COCO-style) and AP75 metrics in table 6. Unlike in
classification tasks, we find that the use of deeper MLP heads during
distillation does not aid detection performance. The performance of different
distillation architectures is nearly identical on the detection task.

Method | AP50 | AP | AP75
---|---|---|---
SimReg-4L-MLP | 74.0 | 45.4 | 47.8
SimReg-2L-MLP | 74.2 | 45.5 | 47.4
SimReg-Linear | 73.6 | 45.1 | 47.9
Table 6: Transfer learning for object detection on Pascal VOC dataset. Student
models with different MLP head architectures are used to perform distillation
on ImageNet dataset and the backbone with R18-C4 architecture is fine-tuned on
PASCAL VOC. Unlike in classification tasks, the performance of different
distillation architectures is nearly identical.

Effect of Data Augmentation: As shown in SEED [Fang et al.(2021)Fang, Wang,
Wang, Zhang, Yang, and Liu] and CompRess [Koohpayegani et
al.(2020)Koohpayegani, Tejankar, and Pirsiavash], for a given student
architecture, distillation from larger models trained using a particular
method is better than directly training the student using the same method. Use
of different and strong augmentations in contrastive SSL approaches has been
shown to hurt generalization performance [Purushwalkam Shiva Prakash and
Gupta(2020)]. Here, we show that when distilling models, the best performance
is obtained when the same augmented image with a weaker augmentation (details
in suppl.) is used as input to both teacher and student networks (table 7).
This suggests that our method can be used to improve generalizability of SSL
models.

Aug Type | Aug Strength | ImageNet | Transfer
---|---|---|---
| Teacher | Student | 1-NN | 20-NN | Linear | Linear
Same | Weak | Weak | 54.8 | 59.9 | 65.1 | 70.4
Same | Strong | Strong | 53.4 | 59.0 | 64.3 | 70.3
Different | Weak | Weak | 54.7 | 59.8 | 64.6 | 70.0
Different | Weak | Strong | 51.1 | 56.5 | 62.0 | 68.9
Different | Strong | Strong | 50.3 | 56.0 | 61.4 | 68.7
Table 7: Role of augmentation strength: During distillation either the ‘same’
augmented image or two ‘different’ augmentations of a single image are used as
inputs to the teacher and student networks. The augmentations strength is
varied for both the settings. We find that the performance is best when the
same image with weak augmentation is used. This is significant since using
different and stronger augmentations improve classification performance of SSL
models but decrease their generalizability [Purushwalkam Shiva Prakash and
Gupta(2020)].

Multi-teacher Distillation: We train a single student model from multiple
teacher networks trained with different SSL methods. Regression with a 4 layer
MLP head significantly outperforms one with linear prediction (table 8).

Training | ResNet-18+Linear | ResNet-18+4L-MLP
---|---|---
Inference | ResNet-18 | ResNet-18
MoCo-v2 | BYOL | SwAV | 1-NN | 20-NN | Linear | 1-NN | 20-NN | Linear
✓ | ✓ |  | 50.9 | 56.5 | 62.6 | 56.4 | 61.3 | 66.3
✓ |  | ✓ | 49.9 | 55.6 | 61.8 | 55.1 | 60.4 | 65.4
| ✓ | ✓ | 52.0 | 57.7 | 64.8 | 56.5 | 61.4 | 67.0
✓ | ✓ | ✓ | 51.1 | 56.7 | 63.2 | 56.5 | 61.6 | 66.9
Table 8: Multi-teacher distillation on ImageNet. We train a single student
model (ResNet-18) from multiple SSL teacher networks (ResNet-50) using a
common backbone network and a separate prediction head for each teacher.
Networks with 4 layer prediction heads can better match each of the teachers
and thus vastly outperform those with a linear head on both k-NN and linear
evaluation metrics.

Distillation of Supervised Teacher: All the previous teacher networks were
trained in a self-supervised manner. We additionally analyze the distillation
from a teacher trained with supervision (table 9). Note that the distillation
remains unsupervised and only the backbone CNN features of the teacher are
regressed. Similar to distillation of self-supervised teachers, we observe
that the use of a deep MLP head during training significantly improves
performance on the ImageNet classification task.

Teacher | Student Arch | Prediction | 1-NN | 20-NN | Linear
---|---|---|---|---|---
| (Inference) | Head (Train) |  |  |
Supervised ResNet-50 |  MobileNet-v2 Backbone | 4L-MLP | 63.77 | 67.87 | 73.5
2L-MLP | 64.7 | 69.3 | 73.5
Linear | 55.4 | 62.0 | 67.5
Table 9: Distillation of supervised teacher. Here, we analyze the role of MLP
heads when distilling the features of a teacher trained using supervision.
Note that the distillation remains unsupervised and only the backbone CNN
features are regressed. Similar to distillation of SSL teachers, we observe
that the use of a deep MLP head during training significantly improves
classification performance on ImageNet classification task.

##  6 Conclusion

Distilling knowledge with deeper student networks leads to better downstream
performance. We surprisingly find that intermediate layer outputs of a
distilled student model have better performance compared to the final layer,
though final layer is trained to mimic the teacher representations. Thus, we
use a prediction MLP head only for optimizing the distillation objective and
achieve boosts in performance with just the backbone network during inference.
We believe studying the reasoning for this effect is an interesting future
work. Our work also serves as an improved benchmark for future self-supervised
distillation works. Additionally, we show that using the same weakly augmented
image for both teacher and student aids distillation.

Acknowledgment: This material is based upon work partially supported by the
United States Air Force under Contract No. FA8750?19?C?0098, funding from SAP
SE, and also NSF grant numbers 1845216 and 1920079. Any opinions, findings,
and conclusions or recommendations expressed in this material are those of the
authors and do not necessarily reflect the views of the United States Air
Force, DARPA, or other funding agencies.

## References

  * [fai()] A library for efficient similarity search and clustering of dense vectors.  https://github.com/facebookresearch/faiss.
  * [pyt(a)] Official pytorch code base for moco.  https://github.com/facebookresearch/moco, a.
  * [pyt(b)] Official pytorch resnet-50 pretrained model.  https://pytorch.org/vision/stable/models.html, b.
  * [Bachman et al.(2019)Bachman, Hjelm, and Buchwalter] Philip Bachman, R Devon Hjelm, and William Buchwalter.  Learning representations by maximizing mutual information across views.  _NeurIPS_ , 2019.
  * [Beyer et al.(2021)Beyer, Zhai, Royer, Markeeva, Anil, and Kolesnikov] Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov.  Knowledge distillation: A good teacher is patient and consistent.  _arXiv preprint arXiv:2106.05237_ , 2021.
  * [Bossard et al.(2014)Bossard, Guillaumin, and Van Gool] Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool.  Food-101–mining discriminative components with random forests.  In _European conference on computer vision_ , pages 446–461. Springer, 2014.
  * [Buciluǎ et al.(2006)Buciluǎ, Caruana, and Niculescu-Mizil] Cristian Buciluǎ, Rich Caruana, and Alexandru Niculescu-Mizil.  Model compression.  In _Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pages 535–541, 2006.
  * [Caron et al.(2018)Caron, Bojanowski, Joulin, and Douze] Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze.  Deep clustering for unsupervised learning of visual features.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , pages 132–149, 2018.
  * [Caron et al.(2020)Caron, Misra, Mairal, Goyal, Bojanowski, and Joulin] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  _NeurIPS_ , 2020.
  * [Chen et al.(2020a)Chen, Kornblith, Norouzi, and Hinton] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pages 1597–1607. PMLR, 2020a.
  * [Chen and He(2021)] Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 15750–15758, 2021.
  * [Chen et al.(2020b)Chen, Fan, Girshick, and He] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He.  Improved baselines with momentum contrastive learning.  _arXiv preprint arXiv:2003.04297_ , 2020b.
  * [Cimpoi et al.(2014)Cimpoi, Maji, Kokkinos, Mohamed, and Vedaldi] Mircea Cimpoi, Subhransu Maji, Iasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi.  Describing textures in the wild.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 3606–3613, 2014.
  * [Doersch et al.(2015)Doersch, Gupta, and Efros] Carl Doersch, Abhinav Gupta, and Alexei A Efros.  Unsupervised visual representation learning by context prediction.  In _Proceedings of the IEEE international conference on computer vision_ , pages 1422–1430, 2015.
  * [Everingham et al.(2010)Everingham, Van Gool, Williams, Winn, and Zisserman] Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman.  The pascal visual object classes (voc) challenge.  _International journal of computer vision_ , 88(2):303–338, 2010.
  * [Fang et al.(2021)Fang, Wang, Wang, Zhang, Yang, and Liu] Zhiyuan Fang, Jianfeng Wang, Lijuan Wang, Lei Zhang, Yezhou Yang, and Zicheng Liu.  Seed: Self-supervised distillation for visual representation.  _ICLR_ , 2021.
  * [Fei-Fei et al.(2004)Fei-Fei, Fergus, and Perona] Li Fei-Fei, Rob Fergus, and Pietro Perona.  Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories.  In _2004 conference on computer vision and pattern recognition workshop_ , pages 178–178. IEEE, 2004.
  * [Gidaris et al.(2018)Gidaris, Singh, and Komodakis] Spyros Gidaris, Praveer Singh, and Nikos Komodakis.  Unsupervised representation learning by predicting image rotations.  _ICLR_ , 2018.
  * [Grill et al.(2020)Grill, Strub, Altché, Tallec, Richemond, Buchatskaya, Doersch, Pires, Guo, Azar, et al.] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  _NeurIPS_ , 2020.
  * [He et al.(2016)He, Zhang, Ren, and Sun] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * [He et al.(2020)He, Fan, Wu, Xie, and Girshick] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9729–9738, 2020.
  * [Heo et al.(2019)Heo, Lee, Yun, and Choi] Byeongho Heo, Minsik Lee, Sangdoo Yun, and Jin Young Choi.  Knowledge transfer via distillation of activation boundaries formed by hidden neurons.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 3779–3787, 2019.
  * [Hinton et al.(2014)Hinton, Vinyals, and Dean] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  _NeurIPS Workshop_ , 2014.
  * [Howard et al.(2017)Howard, Zhu, Chen, Kalenichenko, Wang, Weyand, Andreetto, and Adam] 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:1704.04861_ , 2017.
  * [Huang et al.(2017)Huang, Liu, Van Der Maaten, and Weinberger] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4700–4708, 2017.
  * [Huang et al.(2019)Huang, Cheng, Bapna, Firat, Chen, Chen, Lee, Ngiam, Le, Wu, et al.] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al.  Gpipe: Efficient training of giant neural networks using pipeline parallelism.  _Advances in neural information processing systems_ , 32:103–112, 2019.
  * [Iandola et al.(2016)Iandola, Han, Moskewicz, Ashraf, Dally, and Keutzer] Forrest N Iandola, Song Han, Matthew W Moskewicz, Khalid Ashraf, William J Dally, and Kurt Keutzer.  Squeezenet: Alexnet-level accuracy with 50x fewer parameters and< 0.5 mb model size.  _arXiv preprint arXiv:1602.07360_ , 2016.
  * [Komodakis and Zagoruyko(2017)] Nikos Komodakis and Sergey Zagoruyko.  Paying more attention to attention: improving the performance of convolutional neural networks via attention transfer.  In _ICLR_ , 2017.
  * [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and Pirsiavash] Soroush Abbasi Koohpayegani, Ajinkya Tejankar, and Hamed Pirsiavash.  Compress: Self-supervised learning by compressing representations.  _NeurIPS_ , 2020.
  * [Koohpayegani et al.(2021)Koohpayegani, Tejankar, and Pirsiavash] Soroush Abbasi Koohpayegani, Ajinkya Tejankar, and Hamed Pirsiavash.  Mean shift for self-supervised learning.  _arXiv preprint arXiv:2105.07269_ , 2021.
  * [Krause et al.(2013)Krause, Stark, Deng, and Fei-Fei] Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei.  3d object representations for fine-grained categorization.  In _Proceedings of the IEEE international conference on computer vision workshops_ , pages 554–561, 2013.
  * [Krizhevsky et al.(2009)Krizhevsky, Hinton, et al.] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [Maji et al.(2013)Maji, Rahtu, Kannala, Blaschko, and Vedaldi] Subhransu Maji, Esa Rahtu, Juho Kannala, Matthew Blaschko, and Andrea Vedaldi.  Fine-grained visual classification of aircraft.  _arXiv preprint arXiv:1306.5151_ , 2013.
  * [Misra and Maaten(2020)] Ishan Misra and Laurens van der Maaten.  Self-supervised learning of pretext-invariant representations.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 6707–6717, 2020.
  * [Nilsback and Zisserman(2008)] Maria-Elena Nilsback and Andrew Zisserman.  Automated flower classification over a large number of classes.  In _2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing_, pages 722–729. IEEE, 2008.
  * [Noroozi and Favaro(2016)] Mehdi Noroozi and Paolo Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In _European conference on computer vision_ , pages 69–84. Springer, 2016.
  * [Noroozi et al.(2017)Noroozi, Pirsiavash, and Favaro] Mehdi Noroozi, Hamed Pirsiavash, and Paolo Favaro.  Representation learning by learning to count.  In _Proceedings of the IEEE International Conference on Computer Vision_ , pages 5898–5906, 2017.
  * [Noroozi et al.(2018)Noroozi, Vinjimoor, Favaro, and Pirsiavash] Mehdi Noroozi, Ananth Vinjimoor, Paolo Favaro, and Hamed Pirsiavash.  Boosting self-supervised learning via knowledge transfer.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 9359–9367, 2018.
  * [Park et al.(2019)Park, Kim, Lu, and Cho] Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho.  Relational knowledge distillation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 3967–3976, 2019.
  * [Parkhi et al.(2012)Parkhi, Vedaldi, Zisserman, and Jawahar] Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar.  Cats and dogs.  In _2012 IEEE conference on computer vision and pattern recognition_ , pages 3498–3505. IEEE, 2012.
  * [Passalis and Tefas(2018)] Nikolaos Passalis and Anastasios Tefas.  Learning deep representations with probabilistic knowledge transfer.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , pages 268–284, 2018.
  * [Paszke et al.(2019)Paszke, Gross, Massa, Lerer, Bradbury, Chanan, Killeen, Lin, Gimelshein, Antiga, Desmaison, Kopf, Yang, DeVito, Raison, Tejani, Chilamkurthy, Steiner, Fang, Bai, and Chintala] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. 2019.
  * [Pathak et al.(2016)Pathak, Krahenbuhl, Donahue, Darrell, and Efros] Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros.  Context encoders: Feature learning by inpainting.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 2536–2544, 2016.
  * [Purushwalkam Shiva Prakash and Gupta(2020)] Senthil Purushwalkam Shiva Prakash and Abhinav Gupta.  Demystifying contrastive self-supervised learning: Invariances, augmentations and dataset biases.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * [Ren et al.(2015)Ren, He, Girshick, and Sun] Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun.  Faster r-cnn: Towards real-time object detection with region proposal networks.  _Advances in neural information processing systems_ , 28:91–99, 2015.
  * [Romero et al.(2015)Romero, Ballas, Kahou, Chassang, Gatta, and Bengio] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets.  _ICLR_ , 2015.
  * [Russakovsky et al.(2015)Russakovsky, Deng, Su, Krause, Satheesh, Ma, Huang, Karpathy, Khosla, Bernstein, et al.] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  _International journal of computer vision_ , 115(3):211–252, 2015.
  * [Sandler et al.(2018)Sandler, Howard, Zhu, Zhmoginov, and Chen] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4510–4520, 2018.
  * [Tan and Le(2019)] Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _International Conference on Machine Learning_ , pages 6105–6114. PMLR, 2019.
  * [Tian et al.(2019)Tian, Krishnan, and Isola] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive multiview coding.  _arXiv preprint arXiv:1906.05849_ , 2019.
  * [Tian et al.(2020)Tian, Krishnan, and Isola] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive representation distillation.  _ICLR_ , 2020.
  * [Wang et al.(2020)Wang, Ge, and Wu] Guo-Hua Wang, Yifan Ge, and Jianxin Wu.  In defense of feature mimicking for knowledge distillation.  _arXiv preprint arXiv:2011.01424_ , 2020.
  * [Xiao et al.(2010)Xiao, Hays, Ehinger, Oliva, and Torralba] Jianxiong Xiao, James Hays, Krista A. Ehinger, Aude Oliva, and Antonio Torralba.  Sun database: Large-scale scene recognition from abbey to zoo.  In _2010 IEEE Computer Society Conference on Computer Vision and Pattern Recognition_ , pages 3485–3492, 2010.  10.1109/CVPR.2010.5539970.
  * [Zagoruyko and Komodakis(2016)] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  In _British Machine Vision Conference 2016_. British Machine Vision Association, 2016.

## Supplementary Material

In this supplementary material, we present additional experimental results
(Sec. A) and details on experiment settings and implementation (Sec. B).
Additional results include those on the role of MLP head during training (Sec.
A.1) and self-distillation (Sec. A.2). We publicly release the code222Code is
available at https://github.com/UCDvision/simreg.

##  Appendix A Additional Experimental Results

###  A.1 Role of MLP Head

In tables 1 and 2 of main we analyze how the depth of MLP head during training
and inference affects classification performance. We present additional
results here in table A1 with different teacher and student network settings.
The student networks are trained with different prediction head
configurations. The evaluation is always performed using features from
backbone network for a fair comparison. In addition to the self-supervised
(SSL) teacher models used in the main paper, we consider a supervised teacher
network. The teacher is trained with cross-entropy loss using ground truth
labels on the ImageNet dataset. As in SSL teachers, we use only the backbone
network for distillation from a supervised teacher. Note that the supervised
labels are absent during student training. In both the supervised and self-
supervised settings, the student with 4 layer MLP head consistently
outperforms others on all metrics. Compared to Linear head, 4L-MLP achieves 5
(MoCo-v2, ResNet-18), 11.2 (MoCo-v2, MobileNet-v2) and 6 (Supervised,
MobileNet-v2) percentage points improvement on linear evaluation.

Teacher | Student Arch | Prediction | 1-NN | 20-NN | Linear
---|---|---|---|---|---
| (Inference) | Head (Train) |  |  |
MoCo-v2 ResNet-50 |  ResNet-18 Backbone | 4L-MLP | 54.8 | 59.9 | 65.1
2L-MLP | 52.7 | 58.5 | 63.6
Linear | 48.8 | 54.3 | 60.1
MoCo-v2 ResNet-50 |  MobileNet-v2 Backbone | 4L-MLP | 55.46 | 59.73 | 69.1
2L-MLP | 54.4 | 59.6 | 68.5
Linear | 48.7 | 54.2 | 57.9
Supervised ResNet-50 |  MobileNet-v2 Backbone | 4L-MLP | 63.77 | 67.87 | 73.5
2L-MLP | 64.7 | 69.3 | 73.5
Linear | 55.4 | 62.0 | 67.5
Table A1: Effect of MLP Heads on ImageNet classification performance. As in
table 1 of main paper, we analyze the role of the prediction head used during
training by varying the number of MLP layers. However, the evaluation here is
performed using the features from the backbone network and the prediction head
plays no role during inference. A linear prediction head corresponds to the
architecture used in earlier works [Koohpayegani et al.(2020)Koohpayegani,
Tejankar, and Pirsiavash]. We observe that a deeper prediction module during
training results in substantial boosts in performance. This observation is
consistent across different teacher networks (both SSl and supervised) and
student architectures. Compared to Linear head, 4L-MLP achieves 5(MoCo-v2,
ResNet-18), 11.2(MoCo-v2, MobileNet-v2) and 6(Supervised, MobileNet-v2)
percentage points improvement on linear evaluation.

###  A.2 Self-distillation

In all the previous experiments, a larger teacher network is distilled to a
shallower student. In self-distillation, we consider the same backbone
architecture for both teacher and student. Similar to other experiments, we
use a prediction head (linear or MLP) atop student backbone during
distillation and remove it during evaluation. As we observe in table A2, the
student with a 4 layer MLP head outperforms the teacher in both ImageNet
classification and transfer tasks. The improvement in transfer performance is
particularly significant (+4 percentage points) and might be attributed to the
use of prediction head and weaker augmentations during distillation.

Student Arch | Prediction Head | ImageNet | Transfer
---|---|---|---
(Inference) | (Train) | 1-NN | 20-NN | Linear | Linear
MoCo-v2 Teacher | - | 57.3 | 60.9 | 70.8 | 74.3
ResNet-50 | 4L-MLP | 58.2 | 62.2 | 72.0 | 78.3
ResNet-50 | Linear | 56.4 | 60.6 | 69.7 | 71.8
Table A2: ImageNet Classification and transfer results for self-distillation
with prediction head. In self-distillation, the teacher and student backbone
architectures are the same (ResNet-50). We use a MoCo-v2 pretrained teacher
and train student networks with linear and 4 layer MLP heads. All evaluations
are performed using backbone features. The student with MLP prediction head
outperforms the teacher on both ImageNet classification and transfer tasks. A
boost of 4 percentage points on average transfer accuracy suggests that the
use of prediction head and weaker augmentations during distillation are
beneficial in learning a good generalizable model.

###  A.3 Comparison with CompRess without MLP

In table 1 of the main paper, we observed that the use of MLP head during
distillation benefits both the CompRess variants on the ImageNet
classification task. Here, we show that similar boost in CompRess performance
can be achieved on transfer tasks when distilled with MLP head. We use the
officially provided pretrained models for vanilla CompRess-2q ResNet-18 and
MobileNet-v2 architectures for our comparison and perform transfer analysis
similar to that in table 5 of the main paper. Results in table A3 demonstrate
that performance of vanilla CompRess models are significantly worse compared
to both CompRess with MLP and proposed regression based distillation. Note
that the MLP heads are not used during inference for fair comparison.

Arch | MobileNet-v2 | ResNet-18
---|---|---
Method | CompRess-2q | CompRess-2q | SimReg | CompRess-2q | CompRess-2q | SimReg
| plain | -4L-MLP | -4L-MLP | plain | -4L-MLP | -4L-MLP
Food | 61.4 | 71.4 | 73.1 | 57.6 | 61.7 | 65.4
CIFAR10 | 85.3 | 90.3 | 91.2 | 82.5 | 87.3 | 88.6
CIFAR100 | 65.1 | 73.9 | 76.1 | 62.5 | 68.4 | 70.2
SUN | 53.9 | 58.0 | 59.4 | 52.2 | 54.3 | 57.1
Cars | 35.0 | 60.3 | 62.4 | 30.0 | 37.2 | 42.3
Aircraft | 42.1 | 57.7 | 58.7 | 38.0 | 42.3 | 45.8
DTD | 70.4 | 71.7 | 74.5 | 67.4 | 69.3 | 70.9
Pets | 82.9 | 86.7 | 85.6 | 81.6 | 84.0 | 83.9
Caltech | 85.6 | 91.1 | 91.7 | 85.3 | 87.3 | 89.2
Flowers | 87.6 | 94.3 | 95.1 | 83.0 | 86.4 | 90.9

Table A3: Transfer learning performance of CompRess with and without MLP.
Since the teacher networks are self-supervised, generalization of learnt
features to other datasets is important. Similar to ImageNet classification,
CompRess with MLP significantly outperforms vanilla CompRess (CompRess-2q
plain) on all datasets and metrics. MLP heads, if present, are only used
during distillation and are not part of inference networks.

###  A.4 Results with Intermediate Layers of CNN

In our results in table 2 of main paper, we analyze how the classification
performance changes as we consider features from the earlier layers of the
prediction head. Here, we analyize results from various intermediate layers
including those from the CNN backbone. We train a single ResNet-18 student
from a MoCo-v2 ResNet-50 teacher and perform k-NN evaluation using features
from different layers. In table A4, Conv-1 refers to the output of the first
convolutional layer while ResBlk-j refers to the output from the $j^{th}$
residual block. The CNN features for evaluation are obtained by reducing their
spatial dimension and then vectorizing. The spatial dimensions are reduced so
that the feature lengths are roughly the same throughout the backbone for fair
comparison. We observe that the performance increases as we go deeper into the
backbone. The best performance is achieved at the intermediate layer of
prediction head and there is a small drop in accuracy at the final prediction
layer.

Eval Layer | Conv-1 | ResBlk-1 | ResBlk-2 | ResBlk-3 | ResBlk-4 | 2L-MLP | 4L-MLP
---|---|---|---|---|---|---|---
Student | 1-NN | 20-NN | 1-NN | 20-NN | 1-NN | 20-NN | 1-NN | 20-NN | 1-NN | 20-NN | 1-NN | 20-NN | 1-NN | 20-NN
ResNet-18 | 6.2 | 7.2 | 16.0 | 18.0 | 21.8 | 24.3 | 33.4 | 37.4 | 55.3 | 60.2 | 56.0 | 60.9 | 53.4 | 57.6

Table A4: ImageNet classification using intermediate features. We consider a
single student network with 4 layer MLP head and perform k-NN evaluation using
features from various intermediate layer features from the network. For fair
comparison, we match the dimensions from the intermediate convolutional
features (Conv-1 and residual block features) to that of the final backbone
feature (ResBlk-4) by reducing their spatial dimensions. As expected,
performance improves as we use features from deeper layers of the CNN. This
changes in the MLP head where a drop in accuracy at the very last layer of the
prediction head is observed.

##  Appendix B Implementation Details

###  B.1 Teacher Networks

We use teacher networks trained using four different self-supervised
representation learning approaches - MoCo-v2 [Chen et al.(2020b)Chen, Fan,
Girshick, and He], BYOL [Grill et al.(2020)Grill, Strub, Altché, Tallec,
Richemond, Buchatskaya, Doersch, Pires, Guo, Azar, et al.], SwAV [Caron et
al.(2018)Caron, Bojanowski, Joulin, and Douze] and SimCLR [Chen et
al.(2020a)Chen, Kornblith, Norouzi, and Hinton]. We use the official and
publicly available pre-trained weights for these networks with ResNet-50x4
architecture pretrained model for SimCLR teacher and ResNet-50 models for the
remaining methods. MoCo-v2 and SwAV have been trained for 800 epochs and BYOL
and SimCLR for 1000 epochs. For distillation with BYOL, SwAV and SimCLR
teachers we use cached features from the teacher. The cached features are
obtained by passing the entire training data through the teacher network once
and storing the features. Random image augmentation as would be used in non-
cached version is employed to generate the inputs for caching.

###  B.2 Image Augmentations

We use two strategies for augmenting the input image during distillation -
‘weak’ and ‘strong’. ‘Strong’ augmentation refers to the setting used in
MoCo-v2 [Chen et al.(2020b)Chen, Fan, Girshick, and He]. In both augmentation
settings, we apply a series of stochastic transformations on the input image.
A random resized crop (scale is in range [0.2, 1.]), random horizontal flip
with probability 0.5 and normalization to channel-wise zero mean and unit
variance are common for both augmentation methods. In addition to these
transformations, ‘strong’ augmentations use random color jittering (strength
of 0.4 for brightness, contrast and saturation and 0.1 for hue) with
probability 0.8, random grayscaling with probability 0.2 and Gaussian blur
(standard deviation chosen uniformly from [0, 1]).

###  B.3 Optimizer

In all our distillation experiments, we use SGD optimizer with cosine
scheduling of learning rate, momentum of 0.9 and weight decay of 0.0001.
Initial learning rate is set to 0.05. The networks are trained for 130 epochs
with a batch size of 256 using PyTorch [Paszke et al.(2019)Paszke, Gross,
Massa, Lerer, Bradbury, Chanan, Killeen, Lin, Gimelshein, Antiga, Desmaison,
Kopf, Yang, DeVito, Raison, Tejani, Chilamkurthy, Steiner, Fang, Bai, and
Chintala] framework.

###  B.4 Evaluation Metrics

We utilize k-NN and linear evaluation to evaluate classification performance
on ImageNet and linear evaluation to evaluate transfer performance. For
ImageNet linear evaluation, the inputs to the linear layer are normalized to
unit $l_{2}$ norm and then each dimension is shifted and scaled to have unit
mean and zero variance [Koohpayegani et al.(2020)Koohpayegani, Tejankar, and
Pirsiavash]. The layer is trained for 40 epochs using SGD with initial
learning rate of 0.01 and momentum of 0.9. The learning rate is scaled by 0.1
at epochs 15 and 30. For evaluation of transfer performance, we use the
optimizer settings from [Koohpayegani et al.(2021)Koohpayegani, Tejankar, and
Pirsiavash]. The shorter side of the input image is resized to 256 and centre
crop with length 224 is used. The input is channel-wise normalized using the
statistics from ImageNet dataset. We use LBFGS optimizer with parameters
max_iter=20 and history_size=10. Learning rate and weight decay are optimized
by performing a grid search using validation set. The best model is obtained
by retraining with optimal parameters on the combined train and validation
set. 10 different log spaced values in [-3, 0] are used for learning rate
while 9 log values in [-10, -2] are used for weight decay.

###  B.5 MLP Architecture

For the proposed prediction head, we experiment with linear, 2 and 4 layer
MLPs. Each MLP layer is composed of a linear projection followed by 1D batch
normalization and ReLU activation. Let the dimension of the student backbone
output be $m$ and that of teacher $d$. For linear evaluation, a single layer
with input and output dimensions of $(m,d)$ is used. For a 2 layer MLP,
following [Grill et al.(2020)Grill, Strub, Altché, Tallec, Richemond,
Buchatskaya, Doersch, Pires, Guo, Azar, et al.], we use the dimensions
$(m,2m,d)$. We extend this to a 4 layer MLP with the following intermediate
feature dimensions: $(m,2m,m,2m,d)$. Batch normalization and ReLU activation
are not employed at the end of layer 2 for 4 layer MLP head (equivalent to
stacking two 2-layer MLP heads). For our ablation on the role of MLP head
during inference (table 2 in main paper), we compare the performance of our
method at different layers of the MLP head from a single trained network. For
fair comparison, we require all the intermediate dimensions to be same as that
of the output. Thus, for this experiment alone, we use an MLP such that the
feature dimensions are $(m,d,d,d,d)$. The output dimension ($m$) for
ResNet-18, ResNet-50 and MobileNet-v2 are 512, 2048 and 1280 respectively. The
teacher output dimensions are 2048 and 8192 respectively for ResNet-50 and
ResNet-50x4 architectures. From table 2 (network with MLP feature dimensions
(512, 2048, 2048, 2048, 2048)) and table 4 (network with MLP feature
dimensions (512, 1024, 512, 1024, 2048)) results, we observe that higher MLP
feature dimensions might help further boost performance (65.7 vs 65.1 on
ImageNet linear). More ablations on this are necessary to optimize the MLP
architecture.