Variational Information Distillation for Knowledge Transfer
Sungsoo Ahn, Shell Xu Hu, Andreas Damianou, Neil D. Lawrence, Zhenwen Dai
Introduction
Deep neural networks (DNNs) play important roles in various computer vision tasks, e.g., depth estimation , pose estimation , optical flow , object classification , detection , and segmentation . A typical DNN approach for a computer vision task is to train a sophisticated end-to-end neural network with a large amount of labeled data. Such an approach often delivers state-of-the-art performance if a sufficient amount of data is available. However, in many cases, it is impossible to gather sufficiently large data to train a DNN. For example, in many medical image applications , the amount of available data is constrained by the number of patients of a particular disease.
A popular approach for handling such lack of data is transfer learning , where the goal is to transfer knowledge from the source task to facilitate learning on the target task. Typically, one considers the source task to be generic with a larger amount of available data that contains useful knowledge for learning the target task, e.g., knowledge from natural image classification is likely to be useful for fine-grained bird classification . Hinton et al. proposed the teacher-student framework for transferring such knowledge between DNNs being trained on the source and target tasks respectively. The high-level idea is to introduce an additional regularization for the DNN being trained on the target task, i.e., the student network, which allows learning the knowledge existing in the DNN that was pre-trained on the source task, i.e., the teacher network. While the framework was originally designed for knowledge transfer between DNNs on the same dataset, recent works started exploiting its potential for more general transfer learning tasks, i.e., when the source data and the target data are different.
Many knowledge transfer methods have been proposed with various intuitions. Hinton et al. and Ba and Caruana propose to match the final layers of the teacher and the student network, as the outputs from the final layer of the teacher network provide more information than raw labels. Romero et al. proposes to match intermediate layers of the student network to the corresponding layers of the teacher network. Recent works relax the regularization of matching the entire layer by matching carefully designed features/statistics extracted from intermediate layers of the teacher and the student networks, e.g., attention maps and maximum mean discrepancy .
Evidently, there is no commonly agreed theory behind knowledge transfer. This causes difficulty in understanding empirical results and in developing new methods in a more principled way. In this paper, we propose variational information distillation (VID) as an attempt towards this direction in which we formulate the knowledge transfer as maximization of the mutual information between the teacher and the student networks. This framework proposes an actionable objective for knowledge transfer and allows us to quantify the amount of information that is transferred from a teacher network to a student network. Since the mutual information is computationally intractable, we employ a variational information maximization scheme to maximize the variational lower bound instead. See Figure 1 for the conceptual diagram of the proposed knowledge transfer method. We further show that several existing knowledge transfer methods can be derived as specific implementations of our framework by choosing different forms of the variational lower bound. We empirically validate the VID framework, which significantly outperforms existing methods. We observe the gap is especially large in the cases of small data and heterogeneous architectures.
In summary, the overall contributions of our paper are as follows:
We propose variational information distillation, a principled knowledge transfer framework through maximizing mutual information between two networks based on the variational information maximization technique.
We demonstrate that VID generalizes several existing knowledge transfer methods. In addition, our implementation of the framework empirically outperforms state-of-the-art knowledge transfer methods on various knowledge transfer experiments, including knowledge transfer between (heterogeneous) DNNs on the same dataset or on different datasets.
Finally, we demonstrate that heterogeneous knowledge transfer between a convolutional neural networks (CNN) and a multilayer perceptrons (MLP) is possible on CIFAR-10. Our method yields a student MLP that significantly outperforms the best-reported MLPs in the literature.
Variational information distillation (VID)
In this section, we describe VID as a general framework for knowledge transfer in the teacher-student framework. Specifically, consider training a student neural network on a target task, given another teacher neural network pre-trained on a similar (or related) source task. Note that the source task and the target task could be the same, e.g., for model compression or knowledge distillation. The underlying assumption is that the layers in the teacher network have been trained to represent certain attributes of given inputs that exist in both the source task and the target task. For a successful knowledge transfer, the student network must learn how to incorporate the knowledge of such attributes to its own learning.
From a perspective of information theory, knowledge transfer can be expressed as retaining high mutual information between the layers of the teacher and the student networks. More specifically, consider an input random variable drawn from the target data distribution and pairs of layers , where each pair is selected from the teacher network and the student network respectively. Feedforwarding the input through the networks induces pairs of random variables which indicate activations of the selected layers, e.g., . The mutual information between the pair of random variables is defined by:
where the entropy and the conditional entropy are derived from the joint distribution . Empirically, the joint distribution is a result of aggregation over the layers with input sampled from the input distribution . Intuitively, the definition of can be understood as a reduction in uncertainty in the knowledge of the teacher encoded in its layer when the the student layer is known.
We now define the following loss function which aims to learn a student network for the target task while encouraging high mutual information with the teacher network:
where is the task-specific loss function for the target task and is a hyper-parameter introduced for regularization of the mutual information in each layer. Equation (2) needs to be minimized with respect to the parameters of the student network. However, the minimization is hard since exact computation of the mutual information is intractable. We instead propose a variational lower bound for each mutual information term , in which we define a variational distribution that approximates :
where the expectations are over the distribution and the last inequality is due to the non-negativity of the Kullback-Leiber divergence . This technique is known as the variational information maximization . Finally, we obtain VID by applying the variational information maximization to each mutual information term in (2), leading to a minimization of the following loss function:
The objective is jointly minimized over the parameters of the student network and the variational distribution . Note that the entropy term has been removed from the equation (2) since it is constant with respect to the parameters to be optimized. Alternatively, one could interpret the objective (4) as jointly training the student network for the target task and maximization of the conditional likelihood to fit the activations of the selected layers from the teacher network. By doing so, the student network obtains the “compressed” knowledge required for recovering activations of the selected layers in the teacher network.
where indicates the -th entry of the vector , represents the output of a single unit of neural network and is, again, parameterized by softplus function to enforce positivity. For this case, the corresponding layer in the student network is the penultimate layer instead of the logit layer to match the hierarchy of two layers without being too restrictive on the output of the student network. Furthermore, we found that using a simple linear transformation for the parameterization of the mean function was sufficient in practice, i.e., for some weight matrix .
The aforementioned implementations turned out to perform satisfactorily during the experiments. We also considered using heteroscedastic variance , but it gave unstable training with ignorable improvements. Other types of parameterizations such as a heavy-tailed distribution or the mixture density network could be used to gain additional performance. We leave these ideas for future exploration.
See Figure 2 for an illustration of the training VID using the implementation based on equation (5). Here, we display the change in the evaluated log-likelihood of the variational distribution aggregated over channels, i.e., , given input (Figure 2(a)) throughout the VID training process. One observes that the student network is trained gradually for the variational distribution to estimate the density of the intermediate layer from the teacher network (Figure 2(b), 2(c) and 2(d)). As a comparison, we also optimize the variational distribution for the student network trained without knowledge transfer, (Figure 2(e)). For this case, we observe that this particular instance of the variational distribution fails to obtain high log-likelihoods, indicating low mutual information between the teacher and the student networks. Interestingly, the parts that correspond to the background achieve higher magnitudes compared to that of the foreground in general. Our explanation is that the output of layers corresponding to the background that mostly corresponds to zero activations (Figure 2(f)) and contains less information, being a relatively easier target for maximizing the log-likelihood of the variational distribution.
2 Connections to existing works
We first describe the relationship between our framework and the celebrated infomax principle applied to representation learning , stating that “good representation” is likely to contain much information in the corresponding input. Especially, such a principle has been successfully applied to semi-supervised learning for neural networks by maximizing the mutual information between the input and output of the intermediate layer as a regularization to learning the target task, e.g., learning to reconstruct input based on autoencoders . Our framework can be viewed similarly as an instance of semi-supervised learning with modification of the infomax principle: layers of the teacher network contain important information for the target task, and a good representation of the student network is likely to retain much of their information. One recovers the traditional semi-supervised learning infomax principle when we set in the equation (2).
Generalizing mean squared error matching.
Next, we explain how existing knowledge transfer methods based on mean squared error matching can be seen as a specific instance of the proposed framework. In general, the methods will be induced from the equation (4) by making a specific choice of the layers for knowledge transfer and parameterization of heteroscedastic mean in the variational distribution:
Note that Equation (7) corresponds to a Gaussian distribution with unit variance over every dimension of the layer in the teacher network. Ba and Caruana showed that knowledge can be transferred between the teacher and the student networks that were designed for the same task, by matching the output of logit layers from the teacher and the student networks with respect to mean squared error. Such a formulation is induced from the equation (7) by letting , and in the equation (7). This was later extended for knowledge transfer between the teacher and the student networks designed for different tasks by Li and Hoiem , through adding an additional linear layer on top of the penultimate layer in the student network to matching with logit layer in the teacher network. This is induced similarly from the equation (7) by letting , and being a linear transformation, i.e., . Next, Romero et al. proposed a knowledge transfer loss for minimizing the mean squared error between intermediate layers from the teacher and the student networks, with additional convolutional layer introduced for adapting different dimension size between each pair of matched layers. This is recovered from the regularization term in the equation (7) by choosing layers for the knowledge transfer to be intermediate layers of the teacher and the student networks, and being a linear transformation corresponding to a single convolutional layer.
These methods are all similar to our implementation of the framework in that they all use Gaussian distribution as the variational distribution. However, our method differs in two key ways: (a) allowing the use of a more flexible nonlinear functions for heteroscedastic mean and (b) modeling different variances for each dimension in the variational distribution. This allows transferring mutual information in a more flexible manner without wasting model capacity. Especially, modeling unit variance for all dimensions of the layer in the teacher network could be highly restrictive for the student network. To illustrate, the layer of the teacher network might include an activation that contains information irrelevant to the task of the student network, yet requires much capacity for regression of to . This would raise over-regularization issues, i.e., wasting the majority of the student network’s capacity on trying to fit such a unit. Instead, modeling high homoscedastic variance for such dimension make its contribution ignorable to the overall loss, allowing one to “filter” out such unit in an efficient way.
Comparison with feature matching.
Besides the knowledge transfer methods based on mean squared error matching, several works have proposed indirectly matching the handcrafted features extracted from intermediate layers. More specifically, Zagoruyko and Komodakis proposed matching the “attention maps” generated from activations from the layers. Huang and Wang later generalized the attention map to matching the maximum mean discrepancy of the activations. Yim et al. proposed matching the feature called the Flow of Solution Procedure (FSP) defined by the Gram matrix of layers adjacent in the same network. Chen et al. considered matching the reconstructed input image from the intermediate layers of the teacher and the student networks. These methods could be seen as smartly avoiding the aforementioned over-regularization issue by filtering out information in the teacher network using expert knowledge. However, such methods potentially lead to suboptimal results when the feature extraction method is not apt for the particular knowledge transfer task and may discard important information from the layer of the teacher network in an irreversible way.
Experiments
We demonstrate the performance of the proposed knowledge transfer framework by comparing VID to state-of-the-art knowledge transfer methods on image classification. We apply VID to two different locations: (a) VID between intermediate layers of the teacher and the student network (VID-I) and (b) VID between the logit layer of the teacher network and the penultimate layer of the student network (VID-LP). For comparison, we consider the following knowledge transfer methods: the original knowledge distillation (KD) , learning without forgetting (LwF) , hint based transfer (FitNet) , activation-based attention transfer (AT) and polynomial kernel-based neural selectivity transfer (NST) . Note that we consider FitNet as a regularization for training the student network instead of a stage-wise training procedure as first proposed in . We compare knowledge transfer methods for knowledge transfer between same and different datasets, which is commonly referred to as the knowledge distillation and transfer learning tasks respectively.
In all the experiments, we select the same pairs of intermediate layers for knowledge transfer based on VID-I, FitNet, AT and NST. Similarly, the same pairs of layers for knowledge transfer are used for LwF and VID-LP. All the hyper-parameters of all the methods are chosen according to the performance on a validation set, which is 20% of the training set. We carefully pick the set of candidate values of hyper-parameters such that all the values proposed in the original works are included. The presented performances are the average of three repeated runs. More details about experiments are included in the supplementary material. The implementation of the algorithm will be made publicly available shortly.
We first compare knowledge transfer methods on the traditional knowledge distillation task, where a student network is trained on the same task as the teacher network. By distilling the knowledge from a large teacher network into a small student network, we can speed up the computation for prediction. We further investigate two problems for this task: whether we can benefit from knowledge transfer in the small data regime and how much performance we lose by reducing the size of the student network? Note that we do not evaluate the performance of VID-LP and LwF as they are designed for transfer learning. When applied, KD, VID-LP and LwF delivered similar performance.
Knowledge transfer can be a computationally expensive task. Given a pre-trained teacher network on the whole training data set, we explore the possibility of using a small portion of the training set for knowledge transfer. We demonstrate the effect of a reduced training set by applying knowledge distillation on CIFAR-10 with four different sizes of training data. We employ wide residual networks (WRN) for the teacher network (WRN-40-2) and the student network (WRN-16-1), where the teacher network is pre-trained on the whole training set of CIFAR-10. Knowledge distillation is applied to four different sizes of training set: 5000 (the full size), 1000, 500, 100 data points per class.
We compare VID-I with KD, FitNet, AT and NST. We also provide performances of the teacher network (Teacher) and the student network trained without any knowledge transfer (Student) as baselines. We choose four pairs of intermediate layers similarly to , each of which is located at the end of a group of residual blocks. We implemented VID-I using three convolutional layers with hidden channel size as twice of the output channel size. The results are shown in Table 1. Our method, VID-I, outperforms other knowledge transfer methods consistently across all regimes. The performance gap increases as the size of dataset get smaller, e.g., VID-I only drops of accuracy even when data points per each class are provided to the student network. There is a drop without knowledge transfer and a drop for the best baseline, i.e., KD + AT.
Varying the size of the student network.
The size of the student network gives a trade-off between the speed and the performance in knowledge transfer. We evaluate the performance of knowledge transfer methods on different sizes of the student network. The teacher network (WRN-40-2) is pre-trained on the whole training set of CIFAR-100. A student network with four choices of size, i.e., WRN-40-2, WRN-16-2, WRN-40-1, WRN-16-1, is trained on the whole training set of CIFAR-100. We compare our VID-I with KD, FitNet, AT and NST along with the Teacher and Student baselines. The choices of intermediate layers are the same as the previous experiment.
The results are shown in in Table 1. As also noticed by Furlanello et al. , the student network with the same size as the teacher network outperforms the teacher network with all the knowledge transfer methods. One observes that VID-I consistently outperforms FitNet, AT and NST, which correspond to the same choice of layers for knowledge transfer. It also outperforms KD except for the case when the structure of the student network is identical to that of the teacher network, i.e., WRN-40-2, where two methods can be combined to yield the best performance.
2 Transfer learning
We evaluate knowledge transfer methods on transfer learning. The teacher network is a residual network (ResNet-34) pre-trained on the ImageNet dataset . We apply transfer learning to improve the performance of two separate image classification tasks. The first task is a fine-grained bird species classification based on the CUB-200-2011 dataset , which contains 11,788 images in total for 200 bird species. The second task is an indoor scene classification based on the MIT-67 dataset , which contains 15,620 images for 67 classes of indoor scenes. For both tasks, there are a relatively few images per class, which can significantly benefit from knowledge transfer from the ImageNet classification task. To evaluate the performance at various levels of data scarcity, we subsample both datasets into three different sizes (50, 25, 10 per class for MIT-67 and 20, 10, 5 per class for CUB-200-2011) and compare the knowledge transfer methods.
We evaluate the knowledge transfer methods in two scenarios: a smaller student network of the same architecture (ResNet-18) and different architecture (VGG-9) . We compare our VID-I and VID-LP with LwF, FitNet, AT and NST. We evaluate the performance of the student network without transfer learning (Student) as a baseline. For the teacher and the student network with ResNet architecture, we choose the outputs of the third and fourth groups of residual blocks (from the input) as the intermediate layers for knowledge transfer. In the case of the VGG-9 student network, we choose the fourth and fifth max-pooling layers as the intermediate layers for knowledge transfer, which corresponds to the same spatial dimension as the intermediate layers selected from the teacher network. For applying VID-I to the ResNet-18 student network, we use two convolutional layers with the size of intermediate channels as half of the output channel size. When the student network is VGG-9, a single convolutional layer without non-linearity is used.
The results are shown in Table 3. The knowledge transfer from ResNet-34 to VGG-9 gives very similar performance to the transfer from ResNet-34 to ResNet-18 for all the knowledge transfer methods. This shows that knowledge transfer methods are robust against small architecture changes. Our methods outperform other knowledge transfer methods in all regions of comparison. Both VID-I and VID-LP outperforms baselines that correspond to the same choice of layers for knowledge transfer. For the MIT-67 dataset, we observe that our algorithm outperforms even the finetuning method, which requires pre-training of the student network on the source task.
3 Knowledge transfer from CNN to MLP
The transfer learning experiments show the robustness of the knowledge transfer method against small architecture changes. This leads to an interesting question: whether a knowledge transfer method can work between two completely different network architectures. A solution to this question can open a new direction of knowledge transfer and potentially offer solutions to many problems, e.g., speeding up prediction of recurrent neural networks (RNNs) by transferring knowledge from a RNN to a CNN, speeding up prediction of CNN on CPU or low-energy device by transferring knowledge from a CNN to a multi-layer perceptron (MLP).
In this paper, we evaluate the performance of knowledge transfer from CNN to MLP on CIFAR-10. There is a well-known performance gap between CNN and MLP on CIFAR-10 . The state-of-the-art performance on CIFAR-10 with MLP is 78.62% with initialization from auto-encoders and 74.32% using knowledge distillation . Urban et al. also trained a single convolutional layer achieving the performance of 84.6% using knowledge distillation.
We apply the knowledge transfer methods in the knowledge distillation setting as mentioned in Section 3.1. We use a teacher network with convolutional layers (WRN-40-2) pre-trained on CIFAR-10. We use a MLP with five fully connected hidden layers as the student network, constructed by stacking one linear layer, three bottleneck linear layers and one linear layer in sequence. Each is followed by a non-linearity activation in between. Here, the bottleneck layer indicates a composition of two linear layers without non-linearity that is introduced to speed up learning by reducing the number of parameters. All the hidden layers have the same number of units and the bottleneck linear layer is composed of two linear layers with a size of and .
The knowledge transfer between intermediate layers is defined between the outputs of four residual groups of the teacher network and the outputs of the first four fully connected layers of the student network. We compare VID-I with KD and FitNet since these knowledge transfer methods do not rely on spatial structures. For the same reason, AT and NST are not applicable to multilayer perceptrons. VID-I is implemented with multiple transposed convolutional layers without non-linearities. Specifically, the inputs for the variational distributions, i.e., the hidden layers of the MLP are treated as a tensor with spatial dimensions. Single transposed convolutional layer with a kernel, unit stride and zero padding is followed by multiple transposed convolutional layers with a kernel, two strides, and single padding to match the spatial dimension of the corresponding layer of the teacher network for knowledge transfer. More details on implementations of the student network and the auxiliary distribution are in the supplementary material.
The results are shown in Table 4. Both FitNet and VID-I improve the performance comparing the baseline of directly training the intermediate layers of the student network. VID-I significantly outperforms FitNet on MLPs with different sizes. Furthermore, MLP-4096 outperforms the the state-of-the-art performance with MLP reported by Lin et al. (78.62%) and Ba et al. (74.32%) significantly. More importantly, our method bridges the performance gap between CNN (84.6% using one convolutional layer ) and MLP shown in previous works.
Conclusion
In this work, we proposed the VID framework for effective knowledge transfer by maximizing the variational lower bound of the mutual information between two neural networks. The implementation of our algorithm is based on Gaussian observation models and is empirically shown to outperform other benchmarks in the distillation and transfer learning tasks. Using more flexible recognition models, e.g., , for accurate maximization of mutual information and alternative estimation of mutual information, e.g., , are both ideas of future interest.
References
Appendix A Implementation details
For the WRNs and ResNets used throughout the experiments, we use the same architectures as originally described by Zagoruyko et al., and He et al., respectively. For the VGG-9 network used in transfer learning, i.e., Section 3.2, we make a slight modification from the VGG-11 network without deviating from the VGG design philosophy. It is conducted by first stacking eight convolutional layers with channels in order with batch normalization and rectified linear unit (ReLU) after every convolutional layers. Furthermore, additional max-pooling layers are inserted after the -th ReLUs. Then the final max-pooling layer is followed by global average pooling layer and a linear layer leading up to the prediction of the labels. For the MLP used in knowledge transfer from CNN to MLP, i.e., Section 3.3, we sequentially stack one linear layer, three bottleneck linear layers and one linear layer leading to the prediction of labels, where dropout with drop rate of , batch normalization and ReLU was inserted between each layers. Here, the bottleneck layer indicates a composition of two linear layers without non-linearity that is introduced to speed up learning by reducing the number of parameters. All of the hidden layers have the same number of units and the bottleneck linear layer is composed of two linear layers with a size of and .
A.2 Parameterization of VID
A.3 Loss function and training scheme
In the experiments, the loss function for VID takes the following form:
Finally, we describe the training scheme used for the experiments. Due to unstable gradients in some cases, we clipped the norm of gradients by throughout the experiments. Additionally, the homoscedastic variance for the variational distribution in equation (5) and (6) was initialized with value of . In the knowledge distillation experiments, i.e., Section 3.1, when training on the full dataset, we used stochastic gradient descent (SGD) for epochs with batch size of and weight decay of . Initial learning rate of is decayed times at -th epoch. When training on subset of the dataset, the numbers are appropriately scaled to have similar number of updates for parameters. In the transfer learning experiments, i.e., Section 3.2, when training on the full dataset for ResNet-34 and ResNet-18, we use SGD for epochs with batch size and weight decay of . Initial learning rate of is decayed by times at -th epoch. For the case of VGG-9, we use SGD for epochs with batch size without weight decay. Initial learning rate of is decayed by times at and -th epoch. Again, the numbers are appropriately scaled to match the number of updates for parameters when training on subset of the full dataset. In the knowledge transfer from CNN to MLP experiments, i.e., Section 3.3, we used SGD for epochs with batch size of and weight decay of . Initial learning rate of was decayed by times at and -th epoch.
Appendix B Additional Experiments
In order to examine the behavior of the learned variance parameters in VID, we plot its channel-wise value for different layers in Figure 3. Here, one can observe that the learned variance parameters are diverse, especially accross different layers. Hence, modeling of homoscedastic variance is necessary for obtaining a tighter lower bound of the mutual information in the equation (2).
B.2 Transfer learning from SVHN to MNIST
We also provide additional experimental results for transfer learning from SVHN to MNIST in Table 5. To this end, the teacher network is trained on the full SVHN dataset that was converted to grayscale, then the student network is trained on MNIST with data points per class. We employ LeNet-like architectures for both networks. Again, one observes that VID outperforms over other methods.
B.3 Comparison with adversarial network compression
We additionally compare with the recently proposed adversarial network compression . by repeating the knowledge distillation experiment on CIFAR-10 between ResNets presented by . The corresponding results are reported in Table 6. One observes that our methods outperforms the ANC with a small margin.
B.4 Experimental results with standard deviation
In Table 7, 8, 9, 10, 11 and 12, we provide full experimental results corresponding to the Table 1, 2, 3(a), 3(b), 3(d) and 4 respectively with additional standard deviations for the three repeated runs.
B.5 Additional heat maps for VID training
In Figure 4, we provide additional visualization results of the knowledge transfer based on VID that was plotted in the same way as in Figure 2.