Learning to Retain while Acquiring: Combating Distribution-Shift in
Adversarial Data-Free Knowledge Distillation

In the Adversarial DFKD framework, a generator ($\mathcal{G}$) is used to create pseudo-samples as surrogates to perform knowledge distillation/transfer, and the teacher-student ($\mathcal{T}$-$\mathcal{S}$) setup acts as a joint discriminator to penalize and update generator parameters ($\theta_{\mathcal{G}}$) in an adversarial manner. After updating $\theta_{\mathcal{G}}$, random samples are generated and used to minimize the $\mathcal{T}$-$\mathcal{S}$ discrepancy by updating the student parameters ($\theta_{\mathcal{S}}$). The generator and the student are optimized alternatively up until a fixed number of pre-defined iterations. In essence, the goal of DFKD is to craft a lightweight student model ($\mathcal{S}$) by harnessing valuable knowledge from the well-trained Teacher model ($\mathcal{T}$) in the absence of training data. A general overview of Adversarial DFKD framework is illustrated in Figure 1.

A student model ($\mathcal{S}$) is trained to match the teacher’s ($\mathcal{T}$) predictions on its unavailable target domain ($\mathcal{D}_{\mathcal{T}}$) as part of the distillation process. Let, $p_{\mathcal{S}}(x)=\operatorname{Softmax}(\mathcal{S}_{\theta_{\mathcal{S}}}(x))$ and $p_{\mathcal{T}}(x)=\operatorname{Softmax}(\mathcal{T}_{\mathcal{\theta_{T}}}(x))$, $\forall x\in\mathcal{D}_{\mathcal{T}}$, denote the predicted Student and Teacher probability mass across the classes, respectively. We seek to find the parameters $\theta_{\mathcal{S}}$ of the student model that will reduce the probability of errors ($P$) between the predictions $p_{\mathcal{S}}(x)$ and $p_{\mathcal{T}}(x)$:

where the superscript $i$ denotes the $i^{th}$ probability score of the predicted masses, $p_{\mathcal{S}}(x)$ and $p_{\mathcal{T}}(x)$.

However, DFKD suggests minimizing the student’s error on a pseudo dataset $\mathcal{D}_{\mathcal{P}}$, since the teacher’s original training data distribution $\mathcal{D}_{\mathcal{T}}$ is not available. Typically, a loss function, say $\mathcal{L}_{KD}$, that gauges disagreement between the teacher and the student is minimized $\forall\hat{x}\in\mathcal{D}_{\mathcal{P}}$ by optimizing the student parameters ($\theta_{\mathcal{S}}$):

We use the Mean Absolute Error (MAE) to define loss measure ($\mathcal{L}_{KD}$) as suggested in [10]. Suppose, given the predicted logits (pre-softmax predictions) $t(\hat{x})=\mathcal{T}_{\theta_{\mathcal{T}}}(\hat{x})$ from the teacher model and the predicted logits $s(\hat{x})=\mathcal{S}_{\theta_{\mathcal{S}}}(\hat{x})$ from the student model, we define $\mathcal{L}_{KD}$ as:

Our novelty resides in the student update strategy, as described earlier. Knowledge-Retention and Knowledge-Acquisition relate to two independent goals and can therefore be considered as two discrete tasks. In fact, by aligning these two goals, we empirically observe that, they can cooperate to retain the knowledge on previously encountered examples while acquiring knowledge from newly generated samples. The proposed method utilizes a meta-learning-inspired optimization strategy to effectively replay the samples from memory, say $\mathcal{M}$.

In the typical Adversarial DFKD setup, the student update objective with the generated pseudo samples ($\hat{x}$), i.e., the Knowledge-Acquisition task ($\mathcal{L}_{Acq}$) (Figure 2(a)), is formulated as:

where $\mathcal{L}$ is the MAE (2) between the teacher and the student logits, and $\hat{x}=\mathcal{G}(z),z\sim\mathcal{N}(0,I)$, denotes the batch of randomly generated samples.

Moreover, to alleviate the distribution drift during knowledge distillation in the adversarial setting, previous works have maintained a memory buffer, to store $\hat{x}$ and replay at later iterations to help the student recall the knowledge [3]. At a fixed frequency, batches of samples from current iterations are updated in the memory [3]. Also, recent works [2] have explored using the Generative Replay [28] strategy to simultaneously store the samples, as they are encountered, in a generative model (VAE [15]), and later replay by generating samples from the VAE decoder. Hence, the optimization objective for the Knowledge-Retention ($\mathcal{L}_{Ret}$) can be defined as follows:

where, with a slight abuse of notation, $\mathcal{M}$ denotes a Memory Buffer or Generative Replay or any other replay scheme. Thus, the overall optimization objective to update $\mathcal{\theta_{\mathcal{S}}}$, while considering both the new samples (generated from the latest $\theta_{\mathcal{G}}$) and the old samples (sampled from $\mathcal{M}$) (Figure 2(b)) is defined as:

However, the objective in (5) attempts to simultaneously optimizes $\mathcal{L}_{Ret}$ and $\mathcal{L}_{Acq}$ but does not seek to align the objectives, which leaves them to potentially interfere with one another. Say, we denote the gradients of Knowledge Acquisition and Knowledge Retention tasks with $\nabla\mathcal{L}_{Acq}(\theta)$ and $\nabla\mathcal{L}_{Ret}(\theta)$, respectively. If $\nabla\mathcal{L}_{Acq}(\theta)$ and $\nabla\mathcal{L}_{Ret}(\theta)$ point in a similar direction or are said to be aligned, i.e., $\nabla\mathcal{L}_{Acq}(\theta).\nabla\mathcal{L}_{Ret}(\theta)>0$, then taking a gradient step along $\nabla\mathcal{L}_{Acq}(\theta)$ or $\nabla\mathcal{L}_{Ret}(\theta)$ improves the model’s performance on both tasks. This is however not the case when $\nabla\mathcal{L}_{Acq}(\theta)$ and $\nabla\mathcal{L}_{Ret}(\theta)$ point in different directions, i.e., $\nabla\mathcal{L}_{Acq}(\theta).\nabla\mathcal{L}_{Ret}(\theta)\leq 0$, and the weight parameters obtained, may not be optimal for both the tasks simultaneously. Hence, the intended effect of having the gradient directions aligned is to obtain student parameters ($\theta_{\mathcal{S}}$) that have optimal performance on both $\mathcal{L}_{Acq}$ and $\mathcal{L}_{Ret}$.

The proposed meta-learning inspired approach, seeks to align the two tasks, $\mathcal{L}_{Acq}$ and $\mathcal{L}_{Ret}$. We take cues from Model-Agnostic Meta-learning (MAML)[12], and adapt it to Adversarial DFKD. At each iteration of the parameter update, we pose Knowledge-Acquisition and Knowledge-Acquisition as meta-train and meta-test, respectively. Which means, we perform a single gradient descent step using the current samples ($\hat{x}$) produced from the generator network ($\mathcal{G}$), on the parameters $\mathcal{\theta_{S}}$ and obtain $\mathcal{\theta_{S}}^{\prime}$ as $\theta_{\mathcal{S}}^{\prime}=\theta_{\mathcal{S}}-\alpha\nabla\mathcal{L}_{Acq}(\theta_{\mathcal{S}})$, where $\nabla\mathcal{L}_{Acq}(\theta_{\mathcal{S}})$ denotes gradient of $\mathcal{L}_{Acq}$ at $\theta_{\mathcal{S}}$. Then we optimize on the batch of samples ($\hat{x}_{m}$) obtained from the memory ($\mathcal{M}$), with the parameters $\theta_{\mathcal{S}}^{\prime}$, and then make a final update on $\theta_{\mathcal{S}}$. Thus, formally, the overall meta-optimization objective, with the task of Knowledge Acquisition serving as meta-train and the task of Knowledge Retention as meta-test (Figure 2(c)), can be defined as follows:

In this subsection, we analyze the proposed objective (3.3) to understand how it results in the desired alignment between the Knowledge-Acquisiton and Knowledge-Retention tasks. We assume that meta-train and meta-test tasks provide us with losses $\mathcal{L}_{Acq}$ and $\mathcal{L}_{Ret}$; in our case, $\mathcal{L}_{Acq}$ and $\mathcal{L}_{Ret}$ are the same function computed on the batches $\hat{x}$ and $\hat{x}_{m}$, respectively. We utilize Taylor’s expansion to elaborate the gradient of $\mathcal{L}_{Ret}$ at a point $\theta$ displaced by $\phi_{\theta}$, as described in Lemma 1,

In the absence of the training dataset ($\mathcal{D}_{\mathcal{T}}$), the generator ($\mathcal{G}$) is utilized to obtain pseudo-samples ($\hat{x}$) and perform knowledge-distillation, i.e., $\hat{x}=\mathcal{G}(z)$, $z\sim\mathcal{N}(0,I)$. The generator is learned to maximize the disagreement between Teacher network ($\mathcal{T}_{\theta_{\mathcal{T}}}$) and the Student network ($\mathcal{S}_{\theta_{\mathcal{S}}}$). Additionally, for the generated data $\hat{x}$ to mimic similar responses from the teacher as the real data, we include a prior loss $\mathcal{L}_{\mathcal{P}}$ [4] to be minimized alongside maximizing the discrepancy ($D$). Hence, we update the generator parameters ($\theta_{\mathcal{G}}$) by maximizing the following objective:

Typically the disagreement function ($D$) for the generator is identical to the teacher-student disagreement term [10, 5]. Instead, for teacher-student discrepancy maximization we use the Jensen-Shannon (JS) ($\mathcal{L}_{JS}$) divergence. Our motivation to use JS divergence is based on empirical study by Binici et al. [3]. Hence, $D$ is defined as:

Here $\mathcal{L}_{KL}$ stands for the Kullback–Leibler divergence and $p(a)$ and $p(b)$ denote the probability vectors obtained after the $\operatorname{Softmax}$ applied to the arguments $a$ and $b$, respectively.

Moreover, in Figure 4, we visualize the learning curves of the student networks trained on multiple datasets. We plot the test accuracy of the student at the end of each training epoch. The proposed method exhibits significant improvement in the learning evolution and the peak accuracies achieved, suggesting that the proposed approach can retain the knowledge from previously encountered samples as the learning progresses. However, on Tiny-ImageNet, with Generative replay, we did not observe substantial improvements; we conjecture that this may be due to the complexity of the dataset, and the inability to capture crucial samples as replay for the complex dataset, for a large number of epochs. Also, with Generative Replay we sometimes faced difficulty in training a VAE on a stream of synthetic samples (especially for complex dataset like Tiny-ImageNet) as it suffers due to the distribution drift of its own.

Additionally, we emphasize that we do not strive toward achieving state-of-the-art student classification accuracy (requiring extensive hyper-parameter tuning) in the DFKD setting, but verify the viability of our hypothesis of retaining the previously acquired knowledge while learning on new samples. Nonetheless, we observe that our method improves upon the student classification accuracy on CIFAR100 [16] compared to the contemporary works and the current state-of-the-art [2] with the ResNet-34 [14] ($\mathcal{T}$) and ResNet-18 [14] ($\mathcal{S}$) setup, as shown in Table 2. Additionally, since previous works use the same teacher network with different test accuracies, we also report the teacher accuracies of the respective methods used to distill the knowledge to the student. Nonetheless, we also compute the Teacher-Student accuracy difference ($\Delta_{\mathcal{T}_{Acc}-\mathcal{S}_{Acc}}$) to assess the distance of the student from its teacher in terms of classification accuracy.
Interpreting the Result: Because of the proposed student update strategy, we observe a global monotonicity in the student’s learning evolution which the existing approaches with naive replay [3, 2] claim, but do not warrant (described in Section 3.3). The global monotonicity in the learning evolution encompasses crucial advantages. For example, when the validation data is unavailable, the developer cannot assess the student’s accuracy and identify the best parameters for the student. In such a scenario, the final accuracy is dependent on the random termination epoch set by the developer. In other words, the ideal DFKD approach should sustain high accuracy via monotonically enhancing it during the course of distillation. Therefore, $\mu[\mathcal{S}_{Acc}]$ and $\sigma^{2}[\mathcal{S}_{Acc}]$ contribute as crucial metrics to asses the distillation method as opposed to the maximum accuracy (Acc_max), since the Acc_max value can be attained at any point of time prior to the termination, and can be misleading. The improvements in the monotonicity and the $\mu[S_{Acc}]$ and $\sigma^{2}[S_{Acc}]$ values of proposed method are evident from Table 1, Figure 4 and Figure 5.
Architecture Scalability: The proposed student update strategy is generic and is scalable across different neural network architecture families since the method is not constrained to a particular choice of architecture. From the ResNet-18 [14], WRN-16-1 [33] and WRN-16-2 [33] student learning curves (in Figure 4 and Figure 5), we observe our method’s advantage on both the network architectures. Moreover, for large student network architectures (Deeper or Wider) that include complex layers, the proposed method efficiently handles the intricacies with regard to computing the Hessian and the Hessian-product, which becomes highly essential for cumbersome models.
Improvement across Replay Schemes: Furthermore, the proposed method is agnostic to the memory scheme employed for replay, as demonstrated by the experiments (in Table 1, Figure 4 and Figure 5) using a Memory Buffer and Generative-Replay, thus, rendering our method generalized to the choice of replay. In Table 1 and Figure 5, we can observe that the proposed method enhances the student’s performance on both the replay schemes (Memory Bank and Generative Replay) used in the prior arts. Moreover, we experiment with different memory buffer sizes on WRN-16-1 [33] and WRN-16-2 [33] distillation (in Figure 5) and observe consistent and substantial improvements across different memory sizes. Here, the memory size is defined as the maximum number of pseudo-example batches that the bank can contain and each batch consists of randomly sampled 64 examples from $\hat{x}$.
GPU Memory Utilization: Moreover, our student update strategy brings in no practical memory overhead, compared to memory-based Adversarial DFKD methods. We observe only a minimal increase in the GPU memory usage of few MBs ($\approx$ 40 MB) due to the higher order gradients computed as a part of the update on $\theta_{\mathcal{S}}$ through $\theta_{\mathcal{S}}^{\prime}$. Moreover, we use a single gradient descent step to obtain $\theta_{\mathcal{S}}^{\prime}$, which does not incur a large memory overhead. Thus, we do not opt for a first order approximation [24] of our method, which is much prevalent in the meta-learning literature.