Paper: Self-Supervised Dataset Distillation for Transfer Learning (ICLR-2024)

Reviewer comment (Experiments):
In Equation (2), both $X_s$ and $Y_s$ were trainable parameters and $\theta$ was minimized in the inner optimization. The process looked a little weird to me although I understand that this design was to make inner and outer optimization consistent. More detailed analysis of $Y_s$ can be provided. For example, we know that $X_s$ was initialized with real images, then how about $Y_s$? Would initialization have great impact on the final performance?

Author response:
**[Q1]** More detailed analysis of $Y\_s$ can be provided.

> The target representations $Y\_s=[\hat{\mathbf{y}}\_1 \cdots \hat{\mathbf{y}}\_m]^\top$ are  learnable parameters that are optimized along with the distilled images $X\_s=[\hat{\mathbf{x}}\_1 \cdots \hat{\mathbf{x}}\_m]^\top\$.  We aim to optimize the pair of $X_s$ and $Y_s$ such that a representation space of the model $\hat{g}\_\theta$, which is  trained to minimize $\frac{1}{2}\lVert \hat{g}\_\theta(X_s)- Y_s\rVert^2\_F = \frac{1}{2}\sum\_{i=1}^m\lVert \hat{g}\_\theta(\hat{\mathbf{x}}\_i) - \hat\{\mathbf{y}}\_i\rVert^2\_2$, is similar to that of the model  trained on the full dataset with SSL objective. 

> In other words, each target representation $\hat{\mathbf{y}}\_i$ serves a **response variable** of the corresponding $\hat{\mathbf{x}}\_i$ such that we want to minimize $\lVert \hat{g}\_\theta(\hat{\mathbf{x}}\_i) - \hat{\mathbf{y}}\_i\rVert\_2^2$ for inner optimization. Then we evaluate the mode $\hat{g}\_\theta$ resulting from the inner optimization on the outer objective, compute the gradient with respect to $X_s$ and $Y_s$, and update $X_s$ and $Y_s$. 

> After dataset distillation, we can expect that a model trained with the optimized  $X_s$ and $Y_s$ will obtain a representation space similar to that  of the model trained with the full dataset $X_t$.

---

**[Q2]** How do we initialize $Y_s=[\hat{\mathbf{y}}_1 \cdots \hat{\mathbf{y}}_m]^\top$? 

> Thank you for asking a clarification. We randomly choose a subset of the full dataset $X_t$ for initializing the distilled images $X_s=[\hat{\mathbf{x}}\_1\cdots \hat{\mathbf{x}}\_m]^\top$ and initialize each target representation $\hat{\mathbf{y}}\_i$ with the target model representation $g_\phi(\hat{\mathbf{x}}\_i)$ for $i=1,\ldots, m$. We have clarified how we initialize both $X_s$ and $Y_s$ in Algorithm 1.

---

**[Q3]** Does the initialization of $Y_s$ matter for the final performance?

> Thank you for your important question on our work. Following your suggestion, we conduct an ablation study by initializing the target representation $Y_s=[\hat{\mathbf{y}}\_1 \cdots \hat{\mathbf{y}}\_m]^\top$ with $\hat{\mathbf{y}}\_i$ drawn from standard normal distribution (i.e., $\hat{\mathbf{y}}\_i \overset{\mathrm{iid}}{\sim} N(0, I\_{d_y})$) for $i=1,\ldots, m$ instead of initializing it with the target model (i.e., $\hat{\mathbf{y}}_i = g\_\phi(\hat{\mathbf{x}}\_i)$ for $i=1,\ldots,m$). We use CIFAR100 as source dataset, and the results are as follows:
| Method            | CIFAR100       | CIFAR10        | Aircraft       | Cars           | CUB2011        | Dogs           | Flowers        |
|-------------------|----------------|----------------|----------------|----------------|----------------|----------------|----------------|
| KRR-ST (**w/o init**)   | $65.09\pm0.09$ | $87.56\pm0.19$ | $25.54\pm0.36$ | $17.75\pm0.07$ | $16.68\pm0.18$ | $20.82\pm0.30$ | $54.94\pm0.15$  |
| KRR-ST (**w init**) | $66.46\pm0.12$ | $88.87\pm0.05$ | $32.94\pm0.21$ | $23.92\pm0.09$ | $24.38\pm0.18$ | $23.39\pm0.24$ | $64.28\pm0.16$ |

> We observe that the initialization with the target model representation is crucial for the final performance. We have included these results in Appendix E. This mirrors the observation where the neural network pre-trained with self-supervised learning methods tends to converge to better solution than the model with random initialization. 

> Note that many **supervised dataset distillation methods [7, 9, 10] also employ specific techniques for initializing distilled class labels** that needs to be optimized along with distilled images. For example, FRePo [7] initializes the distilled class labels with mean-centered one-hot vectors scaled by $\frac{1}{\sqrt{C/10}}$, where $C$ is the number of classes (please see the implementation of FRePo in this **[link](https://github.com/yongchao97/FRePo/blob/43e028a5839a5de367701b9e9544a08ffccb3166/lib/datadistillation/frepo.py#L76)**).

Reviewer comment (Evaluation):
From the visualization, it seemed that the distilled images were not very different from original ones and looked exactly the same. Is it possible to report other metrics to measure the difference between synthetic examples and real ones apart from visualization?

Author response:
**[Q5]** Is it possible to report other metrics to measure the difference between synthetic examples and real ones?

> Yes, it is. We report the FID [1] between a synthetic dataset and a real dataset below:
| Method | CIFAR100 | TinyImageNet |
|--------|----------|--------------|
| Random | 34.49    | 20.31        |
| DSA    | 113.62   | 90.88        |
| DM     | 206.09   | 130.06       |
| MTT    | 105.11   | 114.33       |
| KRR-ST | 76.28    | 47.54        |

> Here, we exclude the FID score of FRePo [7] which uses zca normalization since the Inception network [2] used for computing FID score can not process zca normalized images.  Other than Random which is a subset of the original dataset, our distilled dataset is the most similar to the original dataset among the baselines. However, it is hard to say that FID is highly correlated with performance of the distilled dataset. Although Random achieves the best FID, it shows the lowest performance among dataset distillation methods. 

> Despite our distilled images appearing natural, their higher FID than Random indicates their distinctiveness from the original dataset. In most realistic scenarios, either choosing  a representative subset of the full dataset may not be optimal or the presence of such a subset is not guaranteed [8]. Thus, a distilled dataset that is directly optimized to compress the full dataset diverges from the randomly chosen subset. We have included these results in Appendix F.

Reviewer comment (Reproducibility):
It was not clear how dataset distillation methods such as DSA and DM were adapted to the self-supervised settings and it was a non-trivial process.

Author response:
**[Q6]** It is not clear how supervised dataset distillation methods are adapted to the self-supervised settings, which is a non-trivial process.

> Similar to the supervised learning method as a baseline against  self-supervised learning methods in transfer learning scenario [2, 3, 4, 5], we run supervised distillation methods (Kmeans, DSA, DM, MTT, KIP, FRePo)  to distill a source dataset with **its labels, which we do not use for our method**, into a small synthetic dataset. Subsequently, a model is pre-trained on the distilled dataset and fine-tuned on downstream datasets. Further details are provided in the “Implementation Details” section.

Reviewer comment (Writing):
Clarity Issue: The concept of "target representation" requires further elucidation for clarity and better understanding.

Author response:
The target representations $Y_s=[\hat{\mathbf{y}}\_1 \cdots \hat{\mathbf{y}}\_m]^\top$ are  learnable parameters that are optimized along with the  distilled images $X_s=[\hat{\mathbf{x}}\_1 \cdots \hat{\mathbf{x}}\_m]^\top\$.    We aim to optimize  the pair of $(X_s, Y_s)$  such that a representation space of the model $\hat{g}\_\theta$, which are  trained to minimize $\frac{1}{2}\lVert \hat{g}_{\theta} (X_s)- Y_s\rVert^2_F= \frac{1}{2}\sum\_{i=1}^m \lVert \hat{g}\_\theta (\hat{\mathbf{x}}_i) - \hat\{\mathbf{y}}_i\rVert^2_2 $,  is similar to that of the model, $g\_\phi$,  trained on the full dataset with SSL objective. 

> In other words, each target representation $\hat{\mathbf{y}}\_i$ serves a **response variable** of the corresponding $\hat{\mathbf{x}}\_i$ such that we want to minimize  $\lVert \hat{g}_\theta(\mathbf{x}_i) - \hat{\mathbf{y}}_i\rVert_2^2 $ for inner optimization.  Then we evaluate the mode $\hat{g}\_\theta$ resulting from the inner optimization on the outer objective, compute the gradient with respect to $X_s$ and $Y_s$, and update $X_s$ and $Y_s$. 

> After dataset distillation, we can expect that a model trained with the optimized  $X_s$ and $Y_s$ will obtain a representation space similar to that  of the model trained with the full dataset $X_t$.

Reviewer comment (Theory):
Theoretical Analysis of Instability: It is good to conduct a theoretical analysis of the instability inherent in the bilevel formulation when optimizing a condensed dataset with a self-supervised learning (SSL) objective. My question is, does this instability observed in the bilevel formulation similarly apply to the supervised dataset condensation formulation?

Author response:
Yes it does. The issue of biased gradient is independent of whether using a label for dataset distillation. If there is any randomness such as data augmentation for inner optimization, the (meta-)gradient of the outer objective is a biased estimator of the true gradient. 

 > For supervised objective functions such as cross-entropy loss, we can train a model without data augmentation to eliminate randomness for inner optimization [1]. However this approach is not feasible for self-supervised learning (SSL) objectives. Many of existing SSL objectives [1,2,3] require data augmentation or random masking, introducing the problem of biased gradient for bilevel optimization. Thus, we propose minimizing MSE between representation of distilled images and target representation without any randomness for inner optimization and thus we get an unbiased gradient estimator.

Reviewer comment (Theory):
Stability Issues with SSL Data Augmentation: Furthermore, our empirical observations indicate that the process of back-propagating through data augmentations utilized in self-supervised learning (SSL) introduces instability and poses challenges. However, further explanation is needed to fully comprehend why data augmentation in SSL leads to these instabilities.

Author response:
**Theorem 1** explains the empirically observed training instability when employing data augmentation for inner optimization. If we sample data augmentation or input mask for inner optimization process, the gradient of the outer objective with respect to distilled images $X_s$ and target representation $Y_s$ becomes a biased estimator. This motivates our proposed inner objective which eliminates any randomness and minimizes MSE between feature representation of distilled images $X_s$ and target representation $Y_s$.

Reviewer comment (Experiments):
In the experiments the authors distill the data sets into 1000-2000 examples, for self-supervised learning, without augmentation. The authors do not comment on augmentations when training on the distilled data. This approach might work for the small models and low resolution used in the experiments, but I’m not convinced that it generalizes to larger models, more complex data sets and higher resolution. Data augmentation is a central component in many SSL methods including Barlow Twins, which the authors use.

Author response:
First of all, as shown in Theorem 1, the reason why we do not use data augmentation for inner optimization is that the gradient of the outer objective becomes a **biased estimator when employing data augmentation** for the inner optimization. This theoretical analysis aligns with the empirically observed instability associated with the use of data augmentation for inner optimization.

Secondly, we use the exact same data augmentation as the original paper [1] for training a target model $g_\phi$ on the full dataset with Barlow Twins objective (line 2 in Algorithm 1). We have revised the draft to include the detail in Algorithm 1.

Lastly, in Table 1,2, and 3,  we have already shown that our proposed method **outperforms all the baselines with models (AlexNet, VGG11, MobileNet, ResNet10)  larger than ConvNet, on complex datasets with high resolution images** (Tiny ImageNet and full ImageNet with 64x64 resolution).

Reviewer comment (Experiments):
Some baselines might be weak; for example MobileNet and ResNet10 from scratch get < 4% accuracy on Cars.

Author response:
We more extensively tune hyperparameters of MobileNet and ResNet10 from scratch on Cars, resulting in notable improvement of test accuracy: MobileNet ($3.94\pm0.31 \rightarrow 7.06\pm0.04$) and ResNet10 ($2.43\pm0.08 \rightarrow 3.50\pm0.05$). 

Note that the Cars dataset is a particularly challenging task for the models without any pre-training due to the small number of training instances (8,144) with 196 classes.

Reviewer comment (Writing):
The abstract might be hard to follow for readers unfamiliar with prior works on dataset distillation.

Author response:
We have revised the draft to offer a more detailed background of dataset distillation.

Reviewer comment (Experiments):
This paper dose NOT choose the state-of-the-art baselines in dataset distillation for comparison, such as IDC, IDM, etc.

Author response:
Please recall that we have already included a strong baseline, **FRePo [3], which is a more recent model than IDC and outperforms IDM (IDM: 21.9 vs FRePo: 25.4 on TinyImageNet dataset) for supervised dataset distillation task**.

Furthermore, both IDC [1] and IDM [2] proposed efficient parameterization methods for dataset distillation, which are **orthogonal to our framework**. Specifically, given a storage budget (e.g., $1000\times3\times32\times32$ for CIFAR100), they parameterize low-resolution images ($4000\times3\times16\times16$). These low-resolution images are decoded into higher resolution images  (e.g., $4000\times3\times32\times32$ for CIFAR100) with upsampling functions, **resulting in a larger number of distilled images** compared to other dataset distillation methods. Note that these parameterization methods can also be plugged into our framework for improved performance. We leave this research direction as our future work.

Reviewer comment (Theory):
What is the motivation for minimizing MSE between the original data representation of the model from inner loop and that of the model pre-trained on the original dataset?

Author response:
Please recall that we have already described the motivation behind the outer objective function in Section 3.2: 

"The intuition behind the objective function is as follows. Assuming that a model trained with a SSL objective on a large-scale dataset generalizes to various downstream tasks (Chen et al., 2020b), we aim to ensure that the representation space of the model $\hat{g}_{\theta^{*}(X_s, Y_s)}$ ,  trained on the condensed data, is similar to that of the self-supervised target model $g\_\phi$".

Moreover, we have empirically validated our motivation by visualizing our distilled dataset in Figure 2. The visualization shows that the distilled data points cover most of the feature space induced by the original dataset.