Paper: Embarrassingly Simple Dataset Distillation (ICLR-2024)

Reviewer comment (Experiments):
Although the proposed method uses the same amount of computational cost as the truncated BPTT method, it still requires to store a trajectory of length $M$. If the dataset becomes large, we may need large $M$ for convergence at inner loop, which hinders scalability of the proposed algorithm.

Author response:
In our experiments, both CIFAR10 with IPC50 (total 500 images) and Tiny-ImageNet with IPC10 (total 2000 images) are tested using the same value of $M$. A moderate $M$ is already large enough for yielding good performance and we observe that the benefits from increasing $M$ significantly diminish after $M=50$. This is evidenced in Figure 10 in the appendix, where we compare the effects of varying  $M$ for CIFAR10 IPC10. The figure demonstrates that increasing $M$ from 40 to 100 only yields small improvements of around 0.6%.

Additionally, we truly appreciate Reviewer dNLL's reference [1] for efficient meta-gradient calculation and reference [2] for algorithmic design for Adam with constant memory consumption. Integrating such methods into our framework presents an exciting avenue for developing a scalable algorithm. This integration could potentially enhance the computational efficiency and memory management of our method, an aspect we are keen to explore in future work.

Reviewer comment (Evaluation):
It is hard to analyze why the proposed method helps improving performance compared to BPTT or truncated BPTT.

Author response:
Sorry for the confusion. We have corrected Figure 3 as in the General Response 1. Please also see our discussion on theoretical groundings for RaT-BPTT with Reviewer 1 (Qtcx).

Reviewer comment (Experiments):
The paper requires more comprehensive empirical experiments to verify the effectiveness of the proposed method. First and foremost, the authors of the paper did not conduct experiments for architecture generalization. While Table 1 simply demonstrates the method's ability to generalize from a shallow convnet to a wide convnet, I believe this is not adequate. I am curious about how the proposed dataset distillation method can be applied to other architectures, such as VGG, EfficientNet, and ResNet.

Author response:
In our original Appendix C.3 (now C.5), we already include direct training and transfer results for ResNet-18. We further assessed our method across various architectures to demonstrate its universality. As suggested by the reviewer, we conduct experiments for VGG-11, AlexNet, and ResNet-18. The EfficientNet is designed for ImageNet dataset with input size 224 $\times$ 224, and is only applied for finetuning on rescaled CIFAR datasets. Therefore, we choose to incorporate AlexNet instead. 

For all these architectures, we conduct training it from scratch and transferring from the distilled dataset with 3-layer ConvNet. To our knowledge, we are the pioneers in applying direct distillation to a standard-sized network like ResNet-18 and VGG-11. Prior works never train directly on VGG11 and they only use small / modified ResNets like ResNet-10 [3,5], ResNet-12 [4] and ResNet-AP10 [5,6] in these settings. We hope these direct training results demonstrate the generality of our method.

| Architecture | VGG-11 | AlexNet | ResNet-18 |
|----------|----------|----------|----------|
| Transfer from ConvNet | 46.6\% | 60.1\% | 49.2\% |
| Direct Training | 47.7\% | 63.7\% | 53.0\% |

Both transfer and direct training achieve good performance. 
We have added these results as Appendix C.5.

Reviewer comment (Experiments):
The paper does not examine whether the proposed method scales to the ImageNet dataset which has 1000 classes.

Author response:
Most existing works in the dataset distillation field do not directly train on the full ImageNet dataset. However, there are notable exceptions, with two recent works [7, 8] scaling dataset distillation to ImageNet 1K. Given our method's strong performance on the CUB dataset and Tiny-ImageNet, both encompassing 200 classes, we are confident in its scalability to the 1,000 classes in ImageNet-1K. To substantiate this claim, we are happy to include preliminary results on ImageNet 1K in the final camera-ready version of our paper, demonstrating the potential of our approach in handling larger, more complex datasets.

Reviewer comment (Evaluation):
There is no comparison of computational cost between the proposed method and the other baselines. I think the proposed method is way more expensive than other baselines (such as FrePo) which simplify the inner loop.

Author response:
We have conducted a comparative analysis of the total training time for several methods, utilizing a consistent computational environment on an RTX8000 with 48GB. It should be noted that we have excluded RCIG from this comparison, as our reproduced accuracy falls short of the reported number. The following are the recorded training times (in hours) for CIFAR10 with IPC10: KIP (over 150), LinBa (100), FrePO (8), RaT-BPTT (22), MTT (14), and IDC (35). Among these methods, our cost ranks as the third best.

There are ways to further improve the efficiency. 1) The current package we utilize for meta-gradient calculation, the higher package, as noted in [2], lacks efficiency compared to other methods. We could lower the time cost by altering our implementation to more efficient methodologies. 2) The references [1, 2] contain efficient designs for the meta-gradient calculation.  As reported in [2], it could lead to up to 2x speedups compared with the higher package. This improvement would not only enhance the performance of our method but also bring it in line with the efficiency benchmarks set by methodologies like FrePO. 3) Similar to FrePO, we may keep a pool of parameter checkpoints to further optimize our method. This strategy would reduce the need for inner training from new random initializations.

Reviewer comment (Experiments):
What happens if we scale the meta-gradient of BPTT with the norm of the gradient during training? The authors argue training instability of BPTT based on the larger gradient norm. Then the most straightforward way to mitigate the issue would be gradient clipping or scaling the gradient.

Author response:
We would re-iterate that in Figure 3, we demonstrate the notable instability of meta-gradients through the gradient norm. As suggested by the reviewer, we conduct an ablation study on controlling the gradient norm of BPTT. Results in Appendix C.3 shows only marginal improvements. 

Our foundational approach has already incorporated gradient clipping to manage extreme gradient norm values, employing a standard exponential moving average (EMA) with a 0.9 decay rate and capping the gradient norm at twice the adaptive norm.

To further stabilize the gradient norm, we explored two additional methods: 1) BPTT-Gradient Clipping, limiting the gradient norm to no more than 1.05 times the adaptive norm, and 2) BPTT-Normalized Gradient, ensuring a consistent gradient norm of 1 throughout training. However, as Figure 12 illustrates, these methods achieved only marginal enhancements over the basic BPTT approach. Their performance trails behind RaT-BPTT, with a threefold increase in optimization time due to extended backpropagation.

These findings highlight challenges such as deviation from the kernel regime, variance from extensive unrolling and pronounced non-convexity, contributing to gradient instability, as evidenced by fluctuating gradient norms. Addressing these issues solely by adjusting gradient norms proves insufficient.

Please also see our discussion with Reviewer dNLL where we use the normalized variance metric to demonstrate the notable instability. The instability is not just in the overall norm, but also across different samples and different steps.

Reviewer comment (Evaluation):
Using meta-gradient norm as an indicator for stability (more in questions).

Author response:
Thank you for giving this suggestion! In Figure 3, we investigated the stability of meta-gradients using gradient norms as a metric, predicated on the notion that stable and efficient learning should manifest as consistent and decreasing gradient norms throughout training. Following your suggestion, we now introduce another metric for evaluating gradient stability: the normalized gradient variance, in line with the methodology proposed in [1]. Each variance value reflects the instability across the batch samples, and the values across time steps reflects the instability across training steps.

To calculate this metric, we compute the average variance of all gradient entries using a set of 100 samples from the evaluation batch. Given the different scales in gradient norms across different methods, we normalize this variance against the square of the norm. This normalization yields a more consistent metric, termed the normalized variance. Employing the same experimental setup as in Figure 3, we present the results in Figure 11. It shows that RaT-BPTT not only maintains lower variance at each training step but also demonstrates more consistent variance trajectories over the course of training. These findings, in conjunction with the earlier results, collectively offer a comprehensive view of the argued training instability. We have added the discussion in a new section (C.3) in the Appendix, as we feel it further strengthens our analysis and thank this reviewer again for the very valuable suggestion!

We do not employ the Hessian analysis due to the prevailing uncertainty in the field regarding the learning dynamics at the edge of stability, particularly when using the Adam optimizer with moderate and small batch sizes. This aspect of optimization remains a topic of ongoing research, as current results, like those presented in [2, 7], are predominantly confined to scenarios involving full batch training. Given this limitation, it becomes challenging to establish a definitive standard for assessing gradient stability from the Hessian perspective. The variance of gradients instead stands as a more interpretable metric and our results again validate the improvements of RaT-BPTT.

Reviewer comment (Experiments):
Boosting is optimization-agnostic procedure however it’s only tested with RaT-BPTT (more in questions).

Author response:
Yes, thank you for asking this question, prompting us to add what we hope is a valuable proof of concept. Indeed, we believe that the Boosting framework will work with other methods. As a proof of principle, and choosing a representative method, we implement strong boosting (with $\beta = 0$) for MTT from IPC 10 to IPC 50 in steps of IPC 10. The final generalization accuracy is 71.4%. The accuracy improves rapidly throughout the boosting and the final performance is marginally lower (by 0.2%) than the direct distillation of IPC50 with MTT, which stands at 71.6%. We have added the discussion in a new section (C.6) in the Appendix.

Reviewer comment (Evaluation):
(Figure 3) The change in gradient norm with more steps is not a good indicator of training stability. I would suggest looking at either the variance of these gradients [1], or the eigenvalues of the hessian matrix (a.k.a. sharpness) [2] as better aligned indicators.

Author response:
Thank you for giving this suggestion! In Figure 3, we investigated the stability of meta-gradients using gradient norms as a metric, predicated on the notion that stable and efficient learning should manifest as consistent and decreasing gradient norms throughout training. Following your suggestion, we now introduce another metric for evaluating gradient stability: the normalized gradient variance, in line with the methodology proposed in [1]. Each variance value reflects the instability across the batch samples, and the values across time steps reflects the instability across training steps.

To calculate this metric, we compute the average variance of all gradient entries using a set of 100 samples from the evaluation batch. Given the different scales in gradient norms across different methods, we normalize this variance against the square of the norm. This normalization yields a more consistent metric, termed the normalized variance. Employing the same experimental setup as in Figure 3, we present the results in Figure 11. It shows that RaT-BPTT not only maintains lower variance at each training step but also demonstrates more consistent variance trajectories over the course of training. These findings, in conjunction with the earlier results, collectively offer a comprehensive view of the argued training instability. We have added the discussion in a new section (C.3) in the Appendix, as we feel it further strengthens our analysis and thank this reviewer again for the very valuable suggestion!

We do not employ the Hessian analysis due to the prevailing uncertainty in the field regarding the learning dynamics at the edge of stability, particularly when using the Adam optimizer with moderate and small batch sizes. This aspect of optimization remains a topic of ongoing research, as current results, like those presented in [2, 7], are predominantly confined to scenarios involving full batch training. Given this limitation, it becomes challenging to establish a definitive standard for assessing gradient stability from the Hessian perspective. The variance of gradients instead stands as a more interpretable metric and our results again validate the improvements of RaT-BPTT.

Reviewer comment (Experiments):
One subtle change in RaT-BPTT is the usage of Adam optimizer in the inner-loop. Notably, all existing distillation techniques use SGD. Is this the main reason for improvement compared to other techniques? How does RaT-BPTT (SGD) compare with RaT-BPTT (Adam)? I would suggest referring to [5] for a better understanding of using Adam in data distillation.

Author response:
We have opted for Adam instead of SGD to simplify the tuning process for the inner loop. This decision was based on the ability to use a common learning rate without requiring decay in the inner loop. 

As the reviewer suggested, we have implemented RaT-BPTT (SGD) using SGD with learning rate 0.01 and learning rate decays at [120, 200] by 0.2. For IPC10 on CIFAR10, RaT-BPTT (SGD) achieves a 69.0% accuracy (std 0.3%), while RaT-BPTT (Adam) results in a slightly higher accuracy of 69.4% (std 0.4%). Thus, RaT-BPTT (SGD) also outperforms previous methods in this setting by a large margin. It is crucial to note that our improvement is attributed to factors beyond merely employing Adam in the inner loop. We are grateful to the reviewer for prompting this ablation study which, we believe, gives further credence to the RaT-BPP method. Given that RaT-BPP outperforms the benchmarks even when using SGD, we hope that it is not necessary to run some of the benchmark methods with Adam. In fact, we tried to combine MTT with Adam. However, either using Adam in both the teacher and student models or only in the student yields subpar performance, different from the observations in [5]. This discrepancy may be attributed to inherent differences between language and image data processing – with Adam being the predominant optimizer in NLP, while SGD maintains its prevalence in image classification tasks. As [5] unfortunately did not provide detailed configurations for their experiments, we are open to exploring these aspects in greater depth after communicating with their authors.

It is also noteworthy to point out that we are not the first to use Adam for the inner loop during training. [8] also uses Adam for their linearized inner loop. Some other papers [8, 9] have also adopted Adam for the linear loop during evaluations. We suspect that whenever Adam was an option, the benchmarking papers probably tried it without significant improvements.

We hope this response on Adam versus SGD is satisfactory to the reviewer and thank you again for this valuable suggestion.

Reviewer comment (Experiments):
Please include full-data performance in Table 1.

Author response:
We have updated Table 1 in the revised paper as per your suggestion, and hope these numbers are helpful.

Reviewer comment (Reproducibility):
The meta-gradient of SGD in the inner-loop can be efficiently computed [6]. Please mention it in the full-text.

Author response:
Thank you for highlighting the significance of references [5] and [6]. These works indeed open new avenues for accelerating RaT-BPTT. We have now included this discussion to acknowledge their contributions and point out future directions in Section 6.

Reviewer comment (Experiments):
It would be great to include an analysis of the effect of varying $\beta$ in Boost-DD on downstream generalization.

Author response:
Thank you for your valuable suggestions. We are currently in the process of conducting the experiments as recommended. 

In the meantime, could you please provide further clarification on your suggestion '*It would be great to include an analysis of the effect of varying $\beta$ in Boost-DD on downstream generalization*'? What does '*downstream generalization*' mean here? We want to ensure we fully understand your feedback to effectively address it.

Reviewer comment (Theory):
No theoretical results supporting the empirical results. Analyzing the bi-level optimization problem might be difficult in general, but it is better to find some simple setting where theory can explain when/why the proposed method outperforms existing methods.

Author response:
In numerous analyses of bi-level optimization, a common assumption is the convexity of the inner problem, with a focus on convergence to stationary points. However, in scenarios where the inner problem exhibits local strong convexity, along with other favorable conditions, it is unreasonable to expect that the convergence rate of RaT-BPTT would be faster in orders than traditional T-BPTT. On the other hand, in non-convex settings, which are typical of deep learning problems like ours and where RaT-BPTT demonstrates notable empirical performance, providing theoretical proofs of effectiveness becomes exceedingly challenging.

Nonetheless, in Section 3 we have aimed to provide a heuristic to justify the increased performance of RaT-BPTT, particularly from a gradient perspective. The inherent non-convexity in the inner loop necessitates long unrolling to adequately capture long term dependencies. This requirement, however, introduces significant computational challenges and leads to gradient instability due to the compounding effects of non-convex Hessian matrix products. RaT-BPTT is a natural solution to all three problems, by limiting the window to control stability and reduce computational burden, and by moving the window to cover the entire unrolling. Both gradient inspections and empirical performance support such heuristics.

To provide a different angle for theoretical justification, we lay out another hypothesis related to kernel approximations of neural nets (like the NTK). This hypothesis attributes the success of RaT-BPTT, in comparison to both BPTT and T-BPTT, to the separation between feature learning and kernel learning. Our inner loop is training a neural network from scratch. According to [1], there is a significant rotation in the empirical kernel at the onset of training, which then gradually stabilizes. This observation suggests a dominant role of feature learning initially, followed by kernel learning.

From this perspective, for large unrolling $T$, T-BPTT is akin to optimizing kernel learning from $\theta_{T-M}$. The initialization distribution $p_{\theta_{T-M}}$ is only loosely related to the data. Therefore, T-BPTT only optimizes data that favors good kernel learning, without any assurance of effective feature learning optimization. In contrast, RaT-BPTT, with its sliding window, explicitly focuses on also optimizing the feature learning aspect, which is crucial, especially in the initial stages of training. 

Previous works [2] have established provable benefits of feature learning with specifically tailored noises. We find it promising to adopt such constructions into bi-level optimziation problems to show different features captured by feature learning and kernel learning. However, we believe that such a theoretical guarantee is beyond the scope of our empirical study and likely requires an in-depth study of kernel evolution during unrolling, which likely will have empirical components as well.

Reviewer comment (Experiments):
Is the proposed algorithm applicable for general bi-level optimization problems, not only for the dataset distillation application? If so, it would be great to add some discussions on what other applications using bi-level optimization (which currently use BPTT) one can use RaT-BPTT instead.

Author response:
Yes! Thanks for pointing it out. As suggested, we have added a new paragraph in Section 6. This addition discusses the broader implications of our work in bilevel optimizations, specifically highlighting its potential in meta-learning for hyperparameters, MAML-like problems, and Poisoning Attacks.