Dataset Distillation via Knowledge Distillation: Towards Efficient Self-Supervised Pre-training of Deep Networks

We thank the reviewer for acknowledging the strengths of our work, namely: 1) the innovative solution to transform dataset distillation (DD) for SSL into a supervised approach using the knowledge distillation loss 2) tackling the important and under-explored problem of DD for SSL 3) simple and clear writing 4) theoretical and empirical motivations for high variance of gradients of SSL.

W1: Corrected in revision.

W2: In Table 5, we have already included results showing our method can be used to train models as large as ResNet-18.

W3: Discussed in Q2

W4: Changed in revision to "as illustrated theoretically in the simplified setting above".

W5: In Table 6, we have already included results for MKDT applied to SimCLR (another SSL approach).

W6: In Table 7, we have already included results evaluating the encoders trained using MKDT distilled sets with larger fractions of labeled data (10%,50%). We mainly focused on very small subsets in MKDT, since dataset distillation has additional benefits over SAS (subset selection for contrastive learning). SAS mainly focuses on reducing training time, while preserving accuracy; whereas dataset distillation is useful to generate extremely small subsets to enable very fast training/adaptation on memory-limited edge devices. Moreover, dataset distillation has the additional orthogonal goal of preserving privacy.

Q1: In supervised learning, very recent work of [1] has shown promise in generalizing to larger networks and larger datasets (e.g. training a ResNet on a distilled version of ImageNet to 34% accuracy using only 50 images per class). This, however, relies heavily on labels and cannot be applied to dataset distillation for SSL. Currently, dataset distillation for SSL is still underexplored. Our contribution sets the groundwork for how to perform dataset distillation for SSL. As happened with SL dataset distillation, we hope that subsequent work extend this to larger scale models and dataset sizes.

Q2: Recent work in distribution matching [40 from paper, 47 from paper, 1] that has achieved remarkable results has relied on data-free knowledge distillation. We did mention these works in our related work for distribution matching methods. The key idea from data-free distillation works of optimizing the synthetic data to match data statistics such as mean, variance etc. [2] has proved extremely helpful in dataset distillation for SL. However, as shown in [40 from paper, 47 from paper, 1], this still requires labels to prevent a collapse of synthetic data from different classes. Nonetheless, the connection to data-free knowledge distillation is indeed interesting and we believe future work could adapt DFKD for DD for SSL, using our idea of converting SSL to SL using knowledge distillation.
We have also added additional references to DKFD literature to our revision.

We are eager to engage in further discussion to resolve any other comments / concerns.

References:

[1] Shao, Shitong, et al. "Generalized large-scale data condensation via various backbone and statistical matching." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[2] Lopes, Raphael Gontijo, Stefano Fenu, and Thad Starner. "Data-free knowledge distillation for deep neural networks." arXiv preprint arXiv:1710.07535 (2017).

We thank the reviewer SPPk for their appreciation of 1) the innovative approach we provide for the important problem of dataset distillation for SSL 2) the theoretical motivation for using knowledge distillation to generate lower variance trajectories to enable trajectory matching for DD for SSL 3) the well-written related works section, situating the contribution of our work in the current research.

We now address the weaknesses and questions raised by the reviewer:

W1: As the reviewer observed, in our related work we did identify the label dependency of the SL method as preventing them from being applied to SSL. Gradient matching methods necessarily need labels, as it is crucial to match gradients per-class to learn class-discriminative features as shown in [46 from paper]. For distribution matching, a similar concern exists where distilling without labels would lead to a collapse in representations i.e. all classes would have nearly identical representations. We confirm this empirically in the table below, showing that DM without labels is not able to distill anything meaningful, performing no better than random subsets. Since MTT is the only method that can indeed distill without labels and still learn class-discriminative features, we do believe MTT is the best candidate from SL distillation techniques to be adapted to SSL.

Table: Distribution Matching (DM) without Labels (Distilled Data Size = 5%)

The fraction % next to the dataset name denotes the fraction of downstream labels used. (Compare this to #s in Table 4 from paper)

W2, W3: We combined the responses for these 2 weaknesses since a clear explanation of Figure 1 helps us explain why other variance reduction techniques are unlikely to be sufficiently effective. Firstly, as shown in [1], SAM is a regularization for finding more-generalizable minima by penalizing the sharpness. This is distinct from the problem of reducing variance of gradients. In fact [5] shows that SAM’s optimization leads to higher variance of gradients and can benefit significantly itself from variance reduction. Secondly, classical variance reduction techniques have been shown to be unsuccessful for deep neural networks [4]. Thirdly, we provide additional results (see Table ) showing that MTT applied directly to higher variance SSL trajectories leads to distilled sets that perform worse than even trivial baselines (random subsets / no-pretraining). Additionally, in Figure 1, we do indeed explore a simple alternative to reduce variance for SSL, i.e., increasing the batch size, and show how this is still insufficient: In Figure 1a, we provide evidence for the higher variance in gradients using empirical estimates. Since weights at the end of each iteration are determined exactly by the gradient, we use the variance of the weights at end of each iteration for models starting from the same initialization to estimate this. As we see, the SL (KD) loss has far lower variance than SSL. Moreover, while increasing the batch size by 4x does reduce the variance slightly, it is not nearly as much as SL. This indicates that reducing the variance of SSL through other regularization techniques may also likewise be insufficient. Figure 1b shows that as a result of the higher variance, when the distillation process optimizes the distilled set to match the training trajectories, it is not able to successfully optimize the distilled set to match the trajectories i.e. it cannot minimize the distillation loss. This is due to how challenging this optimization is due to the high variance of SSL. Finally, Figure 1c confirms that the insufficient minimization of distillation loss is problematic. In particular, it shows that for the cases where distillation loss was insufficiently minimized, there is little to no change between distilled images and their initializations indicating that the distillation has been unsuccessful.

Q1: Answered with W2/W3.

Q2: Answered with W1.

Q3: Answered with W2/W3.

Q4: We empirically showed that, on most tasks, initializing with the high-loss subset performs slightly better than initializing with a random subset. Examples with a high loss correspond to more ambiguous data points likely to lie on the boundary of the latent classes of the pretraining data (high loss images from CIFAR100, shown in https://anonymous.4open.science/r/iclr_mkdt_rebuttal-2DE1/iclr_rebuttal_fig.png). As a result, initializing the distilled set with these high-loss examples allows the distilled set to better learn representations of boundary examples, leading the encoder to more closely preserve the teacher model’s representations on them (see average MSE on the top 1% of high-loss examples in https://anonymous.4open.science/r/iclr_mkdt_rebuttal-2DE1/iclr_rebuttal_fig.png). Consequently, since the boundary points are represented more accurately, the linear classifier trained on these representations with downstream data achieves higher accuracy. This is not central to our method but rather an additional component that provides further performance improvements at no extra cost.

Q5: MKDT does indeed enable dataset distillation for various SSL methods (including masked reconstruction) since the only requirement is representations from a teacher model trained with SSL. As we show in Table 6, we have results extending MKDT to SimCLR as well, showing significant gain over baselines. We believe that the effectiveness of the distilled set will be determined by how effective the given SSL algorithm is training the teacher model.

Q6: Largely we borrowed hyperparameters from MTT [1 from paper]. In particular, we use nearly identical number of synthetic steps (40 for us v/s 30-50 for [1 from paper] on CIFAR10/100 and 10 for both on TinyImageNet) and expert epochs (2 epochs, across all datasets, for both us and [1 from paper]). As seen in [1 from paper], when distillation is more challenging, e.g., on larger datasets such as TinyImageNet, the max start epoch is reduced in order to only distill early training dynamics and make optimization more tractable. Even with knowledge distillation trajectories, we observed that optimizing the distillation loss for SSL is harder (since we deal with the more difficult problem of supervised regression with unique representations as labels), thus we reduced the max start epoch by ~10x (2 for us v/s 20 for [1 from paper] on CIFAR10/100 and 2 for us v/s 10 [1 from paper] on TinyImageNet).

We are eager to engage in further discussion to resolve any other concerns or comments!

References:

[1] Foret, Pierre, et al. "Sharpness-aware minimization for efficiently improving generalization." arXiv preprint arXiv:2010.01412 (2020).

[2] Chen, Xuxi, et al. "Data distillation can be like vodka: Distilling more times for better quality." Proceedings of the International Conference on Learning Representations (ICLR), 2024.

[3] Yang, Yu, Hao Kang, and Baharan Mirzasoleiman. "Towards sustainable learning: Coresets for data-efficient deep learning." International Conference on Machine Learning. PMLR, 2023.

[4] Defazio, Aaron, and Léon Bottou. "On the ineffectiveness of variance reduced optimization for deep learning." Advances in Neural Information Processing Systems 32 (2019).

[5] Li, Bingcong, and Georgios Giannakis. "Enhancing sharpness-aware optimization through variance suppression." Advances in Neural Information Processing Systems 36 (2024).

We’d like to thank the reviewer YkVh for appreciating 1) the importance of the problem we tackle i.e. dataset distillation for self-supervised learning and 2) the empirical and theoretical evidence we provide for high variance of SSL trajectories that prevent successful application of trajectory matching (MTT).

We now address the weaknesses and questions raised by the reviewer:

W1: Our evaluation measures performance across several downstream datasets and We have significant improvement when pre-training on CIFAR10/CIFAR100, up to 13%. We acknowledge that, for TinyImageNet, on some datasets, the improvement over subsets is smaller. However, this is not a limitation of our method which proposes using knowledge distillation (KD) trajectories to enable data distillation for SSL. Instead, this is a limitation inherited from trajectory matching (MTT) for higher-resolution datasets such as TinyImageNet. As seen in the MTT paper [1 from paper], the improvement over random subsets is far smaller for higher resolution datasets such as TinyImageNet as compared to CIFAR10/CIFAR100. Remedying this problem for both trajectory matching for SL and SSL is an interesting direction for future work.

W2: We have added these results in our revision (to Tables, 4,5,6 and 7). We also include these results below. As we see, across all these settings, our proposal, MKDT, continues to outperform KRR-ST significantly.

For all the tables below, % in brackets indicates the % of labeled downstream data available.

Table 4: KRR-ST 5%

Table 5: ResNet10

Table 5: ResNet18

Table 6: KRR-ST SimCLR

Table 7: KRR-ST 2% Distilled Data with Larger Label Fractions

W3: We provide additional results (https://openreview.net/forum?id=c61unr33XA&noteId=am03RRoWCi) where we directly apply trajectory matching (MTT) to the higher variance SSL trajectories and show that the corresponding distilled sets perform worse than even trivial baselines (random subsets / no-pretraining). The key insight of our method is that using the student model as intermediary, we are able to convert the problem from SSL distillation, where the gradient has high variance, to distillation of a supervised regression problem where the labels are representations learned by the teacher model. As discussed in Sections 4.1 and 4.2, the supervised regression problem has lower gradient variance as the gradient of each example is independent of other examples in the batch, whereas for SSL, the gradient of each example depends on all other examples in the mini-batch (hence is very sensitive to the choice of mini-batches). The evidence we provide in Thm 4.1 and Figure 1 supports our claims: 1) the variance of the SL regression (knowledge distillation) is less than that of SSL and 2) the lower variance makes distillation more effective.

Q1: We have fixed this in our revision.

Q2: ResNet-18 trained with BarlowTwins (same as MMKDT) (we follow the exact methodology & hyperparameters specified by KRR-ST and have added this to our Appendix)

Q3: We have added examples of distilled images for CIFAR10, CIFAR100 and TinyImageNet for both MKDT and KRR-ST in Appendix E. From these, we can clearly observe that MKDT distilled images seem to sharpen the salient class-features of the image, while blurring out background and other extraneous details; in contrast, the images distilled by KRR-ST appear noisy and blurred.

We are eager to engage in further discussion to resolve any other concerns or comments!

We’d like to thank reviewer 1SyT for appreciating our work’s contribution as 1) the first effective dataset distillation method for SSL pre-training, as well as 2) our approach’s outstanding experimental performance on CIFAR-10, CIFAR100 and TinyImageNet.

We now address the weaknesses and questions raised by the reviewer:

W1: We included a simplified theoretical analysis to highlight how the interaction between examples in a batch by SSL loss leads to higher gradient variance. Thus, we unrolled all the terms to illustrate how the variance for SSL gradient would be larger. The works cited by the reviewer are all for supervised learning; we are not aware of existing analysis of the variance of gradients with non-linear models for contrastive learning. Due to the more complicated nature of the SSL/CL loss, prior works in CL theory [1,2,3] have also relied on analyses of linear models and have demonstrated how such analysis corresponds with the behavior of deep nonlinear networks. Providing an expression for variance of SSL gradient in a more general setting is a challenging problem in its own right. This is because gradient of each example depends on every other example in the batch, hence introducing additional complications. Such an analysis is beyond the scope of this work. Nonetheless, we have revised our claims from “confirmed theoretically” to “illustrate theoretically, in a simplified setting” to reflect the gap between theory and implementation.

W2: Although we use 4-5 layer ConvNets for distillation, in Table 5, we do demonstrate that datasets distilled with these smaller ConvNets can be used to pre-train models as large as ResNet-18. This pre-training is effective as it shows generalization to a variety of downstream tasks with only 1-5% of data, when evaluated by training only a linear classifier on the representations of the pretrained ResNets. We note that using 4-5 layer ConvNets is also common practice in dataset distillation for SL [1, 23, 25, 42, 46, 48 etc. from paper]. This is because gradient, weight, or distribution matching with networks larger than ConvNets becomes prohibitive due to the very high cost and difficulty of optimization. Enabling further scaling of this method to larger networks is indeed an interesting direction for future research. Recent work [6, 9 from paper] on improving the scalability and optimization of trajectory matching, orthogonal to our contribution, can be useful to this end. With regards to the citation for sparse-coding model, we have replaced them with citations from prior literature on unimodal contrastive learning.

W3: There might be a misunderstanding of the goal of our experiment here. We are not trying to demonstrate generalization to architectures larger than the “teacher network”. Instead, our goal is to demonstrate that our methodology can apply to networks larger than the 4-5 layer ConvNets used for distillation. The results on ResNet-10 and ResNet-18 indeed prove that this methodology can generalize to larger networks.

W4: Thank you for catching these! We've addressed them in our revision.

Q1: Since weights = lr * gradient, our Figure 1 shows the high variance of gradients using the variance of weights, after each iteration. The first step shows exactly the high variance of gradients and subsequent steps shows how this leads to even further increase in variance of weights over iterations. Moreover, we focus on showing the impact of the high variance of gradients on variance of weights since trajectory matching matches weights, at different iterations, rather than gradients. In Figure 1a, we compute the variance, in weights, after each iteration (step), across 5-models starting from the same initialization, for both SL (in particular, the knowledge distillation loss which is a supervised regression loss) and SSL. The figure confirms the smaller variance of our method (denoted as MTT SL) vs MTT SSL.

Q2: First, we provide additional results (see https://openreview.net/forum?id=c61unr33XA&noteId=am03RRoWCi) showing that MTT applied directly to higher variance SSL trajectories leads to distilled sets that perform worse than even trivial baselines (random subsets / no-pretraining). Second, as discussed earlier, Fig 1 confirms the connection between variance and distillation performance. Figure 1a established that the variance of the gradient is large for SSL as compared to SL. Figure 1b shows that as a result, distillation loss cannot be effectively minimized by MTT on higher variance trajectories. Figure 1c confirms that the inability to minimize distillation loss leads to no change in the images from their initialization - indicating no distillation has occurred. To further confirm this connection, we considered reducing variance for SSL by increasing the batch size (denoted by 4x batch size SSL). We confirm in Figure 1a that this does indeed reduce the variance. Figure 1b shows that this leads to a slightly better minimization in distillation loss and Figure 1c shows that this leads to a slightly more effective distillation with images changing from their initialization. Overall, the distilled images using the high variance SSL trajectories are nearly identical to the real images used for initialization. Thus, they offer no benefits to pre-training, privacy etc. We use these results together to confirm empirically the connection between gradient variance and effectiveness of trajectory matching distillation.

We are eager to engage in further discussion to resolve any other concerns or comments!

References:

[1] Ji, Wenlong, et al. "The power of contrast for feature learning: A theoretical analysis." Journal of Machine Learning Research 24.330 (2023): 1-78.

[2] Xue, Yihao, et al. "Investigating the Benefits of Projection Head for Representation Learning." Proceedings of the International Conference on Learning Representations (ICLR), 2024.

[3] Xue, Yihao, et al. "Which features are learnt by contrastive learning? On the role of simplicity bias in class collapse and feature suppression." International Conference on Machine Learning. PMLR, 2023.