Multi-view Contrastive Learning for Online Knowledge Distillation

As depicted in Fig.1, $M$ peer networks $\{f_{m}(\bm{x})\}_{m=1}^{M}$ participate in the process of distillation during training, where each network includes a CNN feature extractor and a linear classifier. For contrastive learning, we aim to optimize the embeddings after Global Average Pooling (GAP) layer among peer networks. We linearly transform these embeddings with $l_{2}$-normalization into the same contrastive embedding space with embedding size of 128. Given a instance $\bm{x}_{i}$, the generated contrastive embeddings of $\{f_{m}(\bm{x}_{i})\}_{m=1}^{M}$ are denoted as $\{\bm{v}_{m}^{i}\}_{m=1}^{M}$ Similar with the previous branch-based OKD[9], low-level layers across the $M$ peer networks can also be shared to reduce the complexity and regularize the training networks.

Training and deployment. At the training stage, we jointly optimize $\{f_{m}(\bm{x})\}_{m=1}^{M}$. At the test stage, we discard $M-1$ auxiliary networks $\{f_{m}(\bm{x})\}_{m=1}^{M-1}$, only the last network $f_{M}(\bm{x})$ is kept, resulting in no additional inference cost.

Learning from labels. Each network is trained by Cross-Entropy (CE) loss between predictive probability distribution and hard labels. Given a instance $\bm{x}$ with label $y$, CE loss of the $m$-th network is:

$$\mathcal{L}_{ce}^{m}=-\sum_{c=1}^{C}\eta_{c,y}\log{p_{m}(y|\bm{x})}$$

(1)

Where $\eta_{c,y}$ is indicator that return $1$ if $c=y$ else $0$. $p_{m}(y|\bm{x})$ is the class posterior calculated from the logits distribution $[z_{m}^{1},z_{m}^{2},\cdots,z_{m}^{C}]$ by softmax normalization:

$$p_{m}(y|\bm{x})=\exp{(z_{m}^{y})}/[\sum_{c=1}^{C}\exp{(z_{m}^{c})}]$$

(2)

Overall, CE loss of $M$ networks is $\mathcal{L}_{ce}=\sum_{m=1}^{M}\mathcal{L}_{ce}^{m}$.

Distillation from an online teacher. We simply construct an online teacher by implementing the naive ensemble for the predictive probability distribution of all networks and softened by a temperature $T$ as:

$$\tilde{p}_{E}(c|\bm{x})=\frac{\exp{(\frac{1}{M}\sum_{m=1}^{M}z_{m}^{c}/T)}}{\sum_{d=1}^{C}\exp{(\frac{1}{M}\sum_{m=1}^{M}z_{m}^{d}/T)}}$$

(3)

Where $c\in\{1,2,\cdots,C\}$ and $\tilde{p}_{E}(c|\bm{x})$ denotes the soft probability of the $c$-th class. The soft probability distribution of the $M$-th network $f_{M}(\bm{x})$ is:

$$\tilde{p}_{M}(c|\bm{x})=\exp{(z_{M}^{c}/T})/[\sum_{d=1}^{C}\exp{(z_{M}^{d}/T)}]$$

(4)

We consider transferring probabilistic knowledge from online teacher to the final deployment network $f_{M}(\bm{x})$, KL divergence is thus used for aligning the soft predictions between the former and latter as:

$$\mathcal{L}_{kl}=\mathbf{KL}(\tilde{p}_{E}\parallel\tilde{p}_{M})=\sum_{c=1}^{C}\tilde{p}_{E}(c|\bm{x})\log{\frac{\tilde{p}_{E}(c|\bm{x})}{\tilde{p}_{M}(c|\bm{x})}}$$

(5)

Multi-view contrastive learning. Given a training set $\mathcal{D}=\{(\bm{x}_{i},y_{i})\}_{i=1}^{N}$ includes $N$ instances with $C$ classes. We consider learning the relationship derived from different networks among various data instances: feature embeddings of the same data instance are mutually closed, while that of two data instances with different classes are far away.

Given two networks $f_{a}$ and $f_{b}$ for illustration, generated embeddings across the training set $\mathcal{D}$ are $\{\bm{v}_{a}^{i}\}_{i=1}^{N}$ and $\{\bm{v}_{b}^{i}\}_{i=1}^{N}$, respectively. We define the positive pair as $(\bm{v}_{a}^{i},\bm{v}_{b}^{i})$, and the negative pair as $(\bm{v}_{a}^{i},\bm{v}_{b}^{j}),y_{i}\neq y_{j}$. We expect to make positive and negative pairs achieve high and low cos similarities respectively by contrastive learning. Given the embedding $\bm{v}_{a}^{i}$ of instance $\bm{x}_{i}$ from the fixing view $f_{a}$, we enumerate the corresponding positive embedding $\bm{v}_{b}^{i}$ and $K$ negative embeddings $\{\bm{v}_{b}^{i,k}\}_{k=1}^{K}$ from $f_{b}$. For ease of notation, $\bm{v}_{b}^{i,k}$ is the $k$-th negative embedding relative to $\bm{v}_{a}^{i}$. We regard the optimization as correctly classifying the positive $\bm{v}_{b}^{i}$ against $K$ negative $\{\bm{v}_{b}^{i,k}\}_{k=1}^{K}$ to approximate the full distribution against all negative embeddings, which is inspired by Noise-Contrastive Estimation (NCE) [20].

The idea behind NCE-based approximation is to transform the instance-level multi-classification into binary classification for discriminating positive pairs and negative pairs. Given the anchor $\bm{v}_{a}^{i}$, the probability of $\bm{v}_{b}^{j},j\in\{1,2,\cdots,N\}$ matching $\bm{v}_{a}^{i}$ as the positive pair is:

$$p_{p}(\bm{v}_{b}^{j}|\bm{v}_{a}^{i})=\frac{\exp{(\bm{v}_{a}^{i}\cdot\bm{v}_{b}^{j}/\tau)}}{Z_{i}},Z_{i}=\sum_{n=1}^{N}\exp{(\bm{v}_{a}^{i}\cdot\bm{v}_{b}^{n}/\tau)}$$

(6)

Where $\tau$ is the temperature and $Z_{i}$ is a normalizing constant. Moreover, we define the uniform distribution for the probability of $\bm{v}_{b}^{j},j\in\{1,2,\cdots,N\}$ matching $\bm{v}_{a}^{i}$ as the negative pair, i.e. $p_{n}(\bm{v}_{b}^{j}|\bm{v}_{a}^{i})=1/N$. Assumed that frequency of sampling lies in every $K(K<<N)$ negative instances and 1 positive instance concurrently. Then the posterior probability of $\bm{v}_{b}^{j}$ drawn from the actual distribution of positive instance (denoted as $D=1$) is:

$\displaystyle p(D=1|\bm{v}_{a}^{i},\bm{v}_{b}^{j})=\frac{p_{p}(\bm{v}_{b}^{j}|\bm{v}_{a}^{i})}{p_{p}(\bm{v}_{b}^{j}|\bm{v}_{a}^{i})+K\cdot p_{n}(\bm{v}_{b}^{j}|\bm{v}_{a}^{i})}$

(7)

We minimize negative log-likelihood derived from positive pair $(\bm{v}_{a}^{i},\bm{v}_{b}^{i})$ and $K$ negative pairs $(\bm{v}_{a}^{i},\bm{v}_{b}^{j}),\bm{v}_{b}^{j}\sim p_{n}(\cdot|\bm{v}_{a}^{i})$, which approximates the contrastive loss from $f_{a}$ to $f_{b}$ as:

	$\displaystyle\mathcal{L}_{contrast}^{f_{a}\rightarrow f_{b}}=-$	$\displaystyle\log p(D=1|\bm{v}_{a}^{i},\bm{v}_{b}^{i})$
	$\displaystyle+K\cdot$	$\displaystyle\mathop{\mathbb{E}}_{\bm{v}_{b}^{j}\sim p_{n}(\cdot|\bm{v}_{a}^{i})}\log[1-p(D=1|\bm{v}_{a}^{i},\bm{v}_{b}^{j})]$		(8)

Where $\bm{v}_{b}^{j}\sim p_{n}(\cdot|\bm{v}_{a}^{i})$ denotes the randomly sampled negative embedding retrieved from an online memory bank $V$ [15] instead of real-time computing along with each mini-batch, which allows us efficiently obtaining abundant negative instances for generating contrastive knowledge. Note that only $\bm{v}_{a}^{i}$ and $\bm{v}_{b}^{i}$ are computed by $f_{a}$ and $f_{b}$ in real time from the input instance $\bm{x}_{i}$ in mini-batch.

Symmetrically, we can also fix $f_{b}$ and enumerate over $f_{a}$ by contrasting $\bm{v}_{b}^{i}$ with positive $\bm{v}_{a}^{i}$ and negative $\{\bm{v}_{a}^{i,k}\}_{k=1}^{K}$, resulting in the contrastive loss $\mathcal{L}_{contrast}^{f_{b}\rightarrow f_{a}}$. We thus deduce the overall objective of contrastive learning between $f_{a}$ and $f_{b}$ as $\mathcal{L}(f_{a},f_{b})=\mathcal{L}_{contrast}^{f_{a}\rightarrow f_{b}}+\mathcal{L}_{contrast}^{f_{b}\rightarrow f_{a}}$. We further extend contrastive loss from two views to multiple views, which captures more robust evidences for classification as well as discards noise information. Specifically, we model fully-connected interactions among each view pair that thus lead to $\binom{M}{2}$ relationships. Contrastive loss among $M$ networks is summarized as $\mathcal{L}_{contrast}=\sum_{1\leq a<b\leq M}\mathcal{L}(f_{a},f_{b})$.

Overall learning objective. We combine above three objectives to construct our final objective:

$$\mathcal{L}_{MCL-OKD}=\mathcal{L}_{ce}+T^{2}\mathcal{L}_{kl}+\beta\mathcal{L}_{contrast}$$

(9)

Where $T^{2}$ is used for balancing contributions between hard and soft labels, $\beta$ is a constant factor for rescaling the magnitude of contrastive loss.

Image Classification. We use CIFAR-100 [12] and ImageNet [13] benchmark datasets for evaluations. CIFAR-100 contains 50K training images and 10K test images with the input resolution 32$\times$32 from 100 classes, ImageNet contains 1.2 million images and 50K test images with the input resolution 224$\times$224 from 1000 classes. We use $T=3$ and $\beta=0.025$ in equ.(9). Following [11], we use $\tau=0.1$ and $\tau=0.07$ in equ.(6) for CIFAR-100 and ImageNet respectively, and architecture-aware $K\in[256,16384]$ in equ.(3.2). We use 4 networks in all OKD methods, i.e. $M=4$. Given a model, training graph shares the first several stages and separates from the last two stages. Detailed experimental settings can be found in our released codes. We report the mean error rate over 3 runs.

Few-shot learning. We use miniImageNet [21] benchmark for few-shot classification. Prototypical network [22] is used as the backbone, which also plays the role of a peer network in OKD. We use the standard data split following Snell et al. [22]. At the test stage, we report average accuracy over 600 randomly sampled episodes with 95% confidence intervals for miniImageNet. For logit-based OKD methods, we add an auxiliary global classifier to the original prototypical network over the class space of training set, and perform learning of probabilistic outputs among 4 peer networks.

Image Classification. Table 1 shows the comparisons of the performance among SOTA OKD methods across the various networks of VGG [24], ResNet [25], DenseNet [26] and HCGNet [2]. It can be obviously observed that our MCL-OKD consistently outperforms all other alternative methods of DML [8], CL-ILR [10], ONE [9] and OKDDip [11] by large margins, which indicates that performing MCL on the representations among various peer networks is more effective than transparent learning from probability distributions that previous OKD methods always focus on. Compared to the previous SOTA OKDDip, MCL-OKD achieves absolute 1.84% reduction of error rate on average on CIFAR-100. Extensive experiments on the more challenge ImageNet validate that MCL-OKD significantly outperforms the baseline by 0.79% margin. As shown in Table 3, MCL-OKD can also achieve the best ensemble error rate among compared OKD methods with retaining all peer networks.

Few-shot learning. Table 4 compares accuracy among SOTA KD and OKD methods on few-shot learning. We can observe that MCL-OKD significantly outperforms other OKD methods, which verifies that MCL on representations is more effective than learning on probabilistic outputs especially for metric learning tasks, due to the superiority of generating discriminative feature embeddings for instances from newly unseen classes. Moreover, MCL-OKD achieves better results than RKD [23], which is the SOTA KD method for few-shot learning but needs a pre-trained teacher as the prerequisite.

Training complexity. We take ResNet-34 on ImageNet as an example. MCL-OKD uses extra 1.56 GFLOPS for contrastive computing, which is about 16% of the original 9.60 GFLOPS by vanilla training. In practice, we observed no distinct addition of training time (1.23 hours/epoch v.s. 1.52 hours/epoch on a single NVIDIA Tesla V100 GPU). The memory bank of each peer network needs about 600MB memory for storing all 128-d features, resulting in the total 2.4GB memory for 4 peer networks.