Thank you for your valuable suggestions, in particular to clarify computational aspects of our work. This has promted us to conduct an extensive series of additional experiments, as outlined below. We hope that the additional information below speaks for the merits of RaT-BPTT (and might help to determine whether you could raise your score).

1. 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.

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.

2. It is hard to analyze why the proposed method helps improving performance compared to BPTT or truncated BPTT.

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).

3. 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.

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.

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

4. Secondly, I wonder the proposed method scales to ImageNet dataset which has 1000 classes.

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.

• I am not still convinced that the proposed method scales to large dataset because it still requires backpropagating through truncated inner trajectory. I am not sure the technique from [1] is applicable to the proposed method since [1] requires the inner optimization procedure "invertible". In general, Reverse Mode Differentiation requires large space complexity [2]. I am not familiar with [3], hard to say whether it does really save the memory. I highly recommend authors empirically show that some of those technique reduce the memory without performance degradation so that the proposed method scales to large datasets.

• I do not think FRePO is applicable. Although we sample a feature extractor from a model pool, we need to optimize a linear classifier for inner optimization step.

• Based on this, I keep my initial score as 5.

References

[1] Maclaurin, Dougal, David Duvenaud, and Ryan Adams. "Gradient-based hyperparameter optimization through reversible learning." International conference on machine learning. PMLR, 2015.

[2] Franceschi, Luca, et al. "Forward and reverse gradient-based hyperparameter optimization." International Conference on Machine Learning. PMLR, 2017.

[3] Sachdeva, Noveen, et al. "Farzi Data: Autoregressive Data Distillation." arXiv preprint arXiv:2310.09983 (2023).

We wish to re-emphasize that within the Dataset Distillation community, Tiny-ImageNet is recognized as one of the largest datasets commonly utilized. Our review of over 20 papers reveals that only two recent studies have employed ImageNet-1K. While we are open to including results on ImageNet, the time constraints of the rebuttal phase preclude this addition. However, we would like to reinforce our confidence in our method's scalability with the following arguments, substantiated by relevant references.

Scaling of our method

Contrary to the concerns raised, we disagree with the notion that 'requiring backpropagation through a truncated inner trajectory' inherently leads to scalability issues with larger datasets. It is unrealistic to expect a method to perform equally across all datasets with identical resource allocation. A critical factor is the rate at which resource requirements increase with dataset size. When comparing ImageNet to smaller datasets such as CIFAR10, CUB-200, and TinyImageNet, we note three main areas of increase: the number of classes, the data volume per class, and the network size.

• Our findings with CUB-200 and CIFAR-100 demonstrate that our method does not require additional unrolling for a larger number of classes.

• An increase in data volume linearly increases training time, affecting only the number of steps per epoch. This is a natural consequence of fixed Image Per Class number, akin to the increase in epochs for standard training.

• Our results with TinyImageNet (using ConvNet 4) and our generalization studies on architectures like ResNet-18, VGG-11, and AlexNet indicate that our method can adapt to these larger networks without extra unrolling steps. The memory increase is proportional to the network size, an expected and manageable outcome.

Based on these observations, we conclude that our method's scaling to larger datasets is feasible and does not inherently suffer from the fact of using backpropagation through certain trajectories.

Reducing the memory consumption

Further, we propose strategies to reduce memory consumption beyond natural scaling expectations.

We agree that reversing the inner optimization "requires the inner optimization procedure to be "invertible"". We note that both SGD and Adam satisfy this criterion. Given the final state and velocity (momentums), we could revert to a previous state's weights. Recalculating the gradient and then subtracting it from the velocity or momentum reverts the previous velocity (momentums). With infinite precision, such inversion is exact, as detailed in Algorithm 2 [1]. We refer the reviewer to check out a blog post at [danielewworrall.github.io/blog/2020/12/reversible-optimisers/]. [2] implements similar reversing methods for dataset distillation. While numerical issues, such as loss of precision during backward propagation, are a potential concern, the methodologies outlined in lines 9 to 13 of Algorithm [1] are designed to mitigate this. Furthermore, [3] observes significant potential for gradient reuse, suggesting the accumulation of the overall gradient during the forward pass. This approach results in an equivalent, exact method that maintains constant memory usage, and has been successfully applied to MTT, which similarly involves backpropagation through inner trajectories. Notably, this technique has led to enhanced performance.

Regarding the reviewer's query about [2] "hard to say whether it really saves memory", we direct attention to the fourth figure in [2], which clearly illustrates successful memory consumption reduction. Could the reviewer kindly clarify further concerns?

Using ideas from FRePO

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.

Let us clarify again on how we could use the idea from FRePO. Instead of using random initializaions for every inner unrolling, we could choose one model randomly from the model pool and perform unroll and backward on it. This adjustment could save the forward pass time. This modification affects only lines 5 and 6 of our algorithm and the subsequent optimization and backward stages remain unchanged. This change does not conflict with "Although we sample a feature extractor from a model pool, we need to optimize a linear classifier for inner optimization step".

References

[1] Maclaurin, Dougal, David Duvenaud, and Ryan Adams. "Gradient-based hyperparameter optimization through reversible learning." International conference on machine learning. PMLR, 2015.

[2] Sachdeva, Noveen, et al. "Farzi Data: Autoregressive Data Distillation." arXiv preprint arXiv:2310.09983 (2023).

[3] Cazenavette, George, et al. "Dataset distillation by matching training trajectories." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

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.

We thank the reviewer for a host of very valuable comments and suggestions and for the generally positive reception of our paper. We have tried to address most of your requests, and hope to incite you to raise your score if you like our responses.

1. Using meta-gradient norm as an indicator for stability (more in questions). (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.

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.

2. Boosting is optimization-agnostic procedure however it’s only tested with RaT-BPTT (more in questions). Do you expect the boosting idea to work with techniques other than RaT-BPTT, e.g., DSA [3], MTT [4]?

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.

We thank the reviewer for their time and valuable feedback. We hope the following responses address your questions.

1. 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.

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.

2. 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.

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.

[1] Fort, Stanislav, et al. "Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel." Advances in Neural Information Processing Systems 33 (2020): 5850-5861.

[2] Karp, Stefani, et al. "Local signal adaptivity: Provable feature learning in neural networks beyond kernels." Advances in Neural Information Processing Systems 34 (2021): 24883-24897.