Farzi Data: Autoregressive Data Distillation

First of all, we really appreciate your review — we believe this is a good quality review where you ask very valid questions; irrespective of your low score recommendation. However, we do hope our responses shed some light on your key confusions and convince you to improve your score. It would really help us push this important line of work to one of the most technically capable audiences, at ICLR. Please let us know if we can add / expand on something:

Q: Important notation is not defined in the main-text.

A: Thank you for raising these super-valid points! We will make sure to address the following changes as you suggested to improve the readability of our paper:

• what the terms $d\mathbf{m}, d\mathbf{x}, d\mathbf{w}$: We apologize for missing this detail. $d\mathbf{m}, d\mathbf{x}, d\mathbf{w}$ are the meta-gradients of $f(w_T)$ w.r.t $m_0, x, w_0$ respectively. Most of this notation is re-used from [1], but we will make sure to add this to the main-text.
• what will we actually update using the meta-gradient in Algorithm 1: Using simple, gradient based optimization in the outer loop (we use Adam), once we have the meta-gradient from Algorithm 1, i.e., $d\mathbf{x} \triangleq df(w_T) / d\mathbf{x}$; we use this gradient to update $\mathbf{x} \leftarrow Adam(\mathbf{x}, d\mathbf{x})$ in the standard way. Note that $\mathbf{x}$ in our case is parameterized as $\tilde{\mathcal{D}}_{syn}$ and $\mathbf{M}$.
• non-standard $\mathcal{O}(100)$: Absolutely, we’ll make sure to write 100x instead of $\mathcal{O}(100)$. Thanks!
• Statement of Proposition 3.2: We’ll make sure to write the complete statement in the main-text instead of the current way. Thank you for pointing this out!

Q: Issue in Theorem 3.1

A: Thank you for catching our honest mistake at the end of Theorem 3.1 — we are highly appreciative! We tried taking the lower-bound approach as you suggested but to no-end: there are no commonly available lower-bounds for Radamacher complexities because of the underlying $\sup()$ in its definition. Even trying to develop novel lower bounds isn’t a trivial process and/or contribution. Hence, given the short timeline of this rebuttal, we have decided to remove Theorem 3.1 from our paper and totally admit to our mistake.

We do hope that you believe that removing this Theorem doesn’t hurt our technical and practical contributions too much, as you have yourself suggested in your review.

Q: Notation questions in Proposition 3.2

A: Thanks for pointing this out - we can definitely improve upon the clarity here. You’re correct by mentioning that this is a recursive formula, starting from $t=T$ going all the way to $t=1$. As outlined in Algorithm 1, $d\mathbf{m}, d\mathbf{x}, d\mathbf{w}$---which are the meta-gradients of $f(w_T)$ w.r.t $m_0, x, w_0$ respectively---are variables which are carried over, with the equations in Algorithm 1 defining the update rules for our recursive meta-gradient calculation. Please note that $d\mathbf{x}$ which represents $df(w_T) / d\mathbf{x}$ is the final meta-gradient used for updating our synthetic data ($\tilde{\mathcal{D}}_{syn}$ and $\mathbf{M}$) that we’re interested in. Similarly, for other meta-learning applications like learning suitable weight initializations (e.g., MAML), we can use the same Algorithm 1 to estimate other meta-gradients like $d\mathbf{w}$ which represents $df(w_T) / dw_0$, if we use Adam in the inner-loop.

Please note that our Algorithm 1 shares the same spirit as the popular reverse-mode SGD algorithm with its recursive formulation [1], albeit much simpler due to being a first-order optimizer.

Q: Are the results in Fig. 6(b) with the total number of meta-gradient steps constant?

A: Great question again! We would like to mention here that all of the results in Fig. 6(b) are performed with the outer loop running till convergence of a maximum of 4000 steps. Please note that realistically, convergence for all settings in our paper happens between 1000-2000 steps, at most. We’ll make sure to add a clarification in the main text.

We would also like to provide some added intuition behind using the pre-trained trajectories as is currently used in Farzi. In the model weight space, $\Omega$ contains a set of points with varying levels of quality (models are trained starting from random till convergence, with intermediate checkpoints being stored). Farzi’s inner loop, using these points as starting points, now has a better exploration of the entire model weight space, and the final data will be better optimized to train models of varying quality. This makes the final data summary more robust and enables it to train models at different stages of training.

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

Thanks to the authors for their detailed response. I believe your proposed changes will improve the paper. However, as it stands, I still cannot recommend acceptance for the following reasons:

1. I don't see an updated pdf, so I cannot verify whether or not the updates to the notation, definitions, etc. were sufficient for the paper to be acceptable for publication.
2. Regarding Theorem 3.1, I agree with the authors that simply removing it from the paper would be acceptable. However, the result of this theorem was used as motivation for the novelty of your procedure when responding to other reviewer comments (last question in the response to Reviewer THWQ), so there is still a problem here as well. It also forms a significant part of the paper as presented, and if it is removed, this may have a large impact on other reviewers' scores.

In short, I don't believe it's possible for the proposed changes to be fairly assessed for the paper to be accepted to ICLR. I recommend the authors to revise the paper and resubmit to a future conference.

We really thank you for your time and effort in reviewing the paper and providing us with your valuable feedback! Below are the responses to your well-motivated questions, which we hope will shed more clarity on Farzi’s design. Please let us know if we can add / expand on something:

Q: Are there reasonable paths forward that would allow this to scale to longer sequence lengths/larger models?

A: Thanks for the question! Regarding longer sequence lengths, even in most LLM pre-training setups, typical input sequence length is in the order of a few thousands. This is certainly feasible with the setup proposed by Farzi’s latent parameterization (but would struggle if the sequence length, e.g., exceed hundreds of thousands). Further, performing the Farzi optimization for larger models is also possible out-of-the-box with hardware/compute requirement no-more-than training the LLM itself. The only reason we weren’t able to perform data distillation for LLMs is because we don’t have the compute to train LLMs by itself.

Q: What's the overall cost saving versus just training a model on the original dataset for more time?

A: Great question! Since this question was also raised by other reviewers, we have put the results in a separate comment up top for the sake of not repeating, and are happy to present our impressive results!

Q: Have the authors thought about cases where there is significant noise in the training corpus?

A: Wonderful question, and highly relevant in, e.g., pre-training scenarios as you mentioned. Even though we don’t have direct evidence in this paper, there is indeed very strong hope for automatic denoising while performing data distillation. Looking at Figure 4 in [1] which also performs data distillation, training on data summarized by data distillation is (1) definitely much better than training on random samples of noisy data; and (2) better than training on the full dataset when there is high-noise in the datasets. The intuition behind this is that data distillation implicitly promotes de-noising, as the overall optimization objective is to maximize model performance within a small data budget, implicitly minimizing the noise in such low-data scenarios.

We would like to end by requesting you to kindly reconsider your final rating as it strictly decides the outcome of this paper, which in your own self-opinion is well-timed, and a fresh approach to data optimization for autoregressive tasks like NLP.

[1] Sachdeva, Noveen, et al. "Infinite recommendation networks: A data-centric approach." Advances in Neural Information Processing Systems 35 (2022): 31292-31305.

We really thank you for your time and effort in reviewing the paper and providing us with your valuable feedback! Below are the responses to your well-motivated questions, which we hope will shed more clarity on Farzi’s design. Please let us know if we can add / expand on something:

Q: I do not see how the time complexity increases with the vocabulary size.

A: Looking at Farzi’s time complexity in Section 3, both the time and space complexities depend directly on the vocabulary size, i.e., $\dim(\mathcal{V})$. Let me provide further intuition on this:

• Taking the naive, non-parameterized Farzi data parameterization as a simpler reference, each entry in the $\mu \times \xi \times \dim(\mathcal{V})$ data summary is a parameter in the outer-loop optimization. Hence, both the sequence length ($\xi$) and the vocabulary size ($\dim(\mathcal{V})$) directly contribute to the space complexity, as well as the time complexity, e.g., for computing and updating the meta-gradient for each of these parameters, that too for each of the $T$ steps in the inner loop.

Q: What is making things O(1) in memory? Is it essentially gradient checkpointing?

A: Thank you for the question! Definitely makes sense to provide an intuition for this. The key technological advancement here is Algorithm 1, and to be more specific Lines 13-15, which derive the reverse-mode meta-gradient calculation equations, when Adam is used in the inner-loop. Having access to these equations, we can now only keep the $d\mathbf{m}, d\mathbf{x}, d\mathbf{w}$ vectors as loop carry-over variables and compute the exact meta-gradient for updating the data summary (i.e., $d\mathbf{x}$) in $O(T)$ time using the for-loop in Algorithm 1, and making the space complexity independent of $T$. However, naive autodiff and standard meta-gradient libraries like higher would store all intermediate variables of the inner-loop making the space complexity linear in terms of $T$.

Further, regarding the comparison to (Deng and Russakovsky), please note that they first of all leverage prior work [1] which developed reverse-mode SGD (similar to our reverse-mode Adam, albeit much simpler due to being a first-order optimizer). Next, reverse-mode Adam is completely different from gradient checkpointing in the case of meta-learning where it’s non-trivial and suboptimal [2]. On the other hand, our reverse-mode Adam is correct (optimal) because of our hand-derived meta-gradient calculations (Proposition 3.2).

Q: How is $\Omega$—the trajectories from the real data—incorporated in the DD process?

A: Thanks for pointing this out! There is definitely a clarification to be made here: $\Omega$ contains the flat set of all episodic checkpoints for the training trajectories and when we sample $\theta_0$ (0 here represents the inner loop time-index), it can come from the union of all the intermediate checkpoints available to us ($\Omega$). We’ll make sure to add the clarification in the pdf as well.

Q: Rank regularization has been done in previous work, and not be presented as a novelty.

A: Thanks for pointing this out! We do not intend to claim latent parameterization of data, in its entirety as our novelty. However, we sincerely believe the following pieces of development to be critical and also novel to our propositions:

• Previous work regarding parameterized data distillation (not just (Deng and Russakovsky)) have been motivated to achieve better performance using fewer total bytes to store the data. However, in our setting of autoregressive data distillation, we would like to explicitly note that naive data parameterization is infeasible computationally due to the vocabulary size dimension in our data.
• We provide formal connections as to why rank regularization improves performance by reducing overfitting and promoting implicit regularization. Please note that while this conclusion is intuitive, we are the first to showcase a formal proof.

We would like to end by requesting you to kindly reconsider your final rating as it strictly decides the outcome of this paper, which in your own self-opinion is well-timed and contains strong empirical evidence.

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

[2] Hascoet, Laurent, and Mauricio Araya-Polo. "Enabling user-driven checkpointing strategies in reverse-mode automatic differentiation." arXiv preprint cs/0606042 (2006).

We really thank you for your time and effort in reviewing the paper and providing us with your valuable feedback! Below are the responses to your well-motivated questions, which we hope will shed more clarity on Farzi’s design. Please let us know if we can add / expand on something:

Q: It is not so obvious that Farzi would be very helpful for much larger text corpora and large language models

A: While we absolutely agree that large language models would be an exciting application for Farzi, we would like to mention that most research labs do not have the compute necessary for such line of work. Hence, we believe penalizing for such reasons (that is only if considered while rating this paper) might be a very strict requirement that implicitly undermines exploratory or academic research.

Q: Could the authors elaborate more on the total runtime.

A: Great question! Since this question was also raised by other reviewers, we have put the results in a separate comment up top for the sake of not repeating, and are happy to present our impressive results!

Q: $\theta_i$ in the definition of $\Omega$ has nothing to do with the update rule for $\theta_t$ in Equation (2)

A: Thanks for pointing this out! There is definitely a clarification to be made here: $\Omega$ contains the flat set of all episodic checkpoints for the training trajectories and when we sample $\theta_0$ (0 here represents the inner loop time-index), it can come from the union of all the intermediate checkpoints available to us ($\Omega$). We’ll make sure to add the clarification in the pdf as well.

Q: It was not clear how exactly the authors chose the final hyperparameters for each setting

A: Just like any deep learning setup, we test a random subset (as per a reasonable grid search budget) of the grid of combinations listed in Tables 4 & 5 and report the best result after selecting the best hyper-parameters using a validation set.

Q: Is a new synthetic batch created at the beginning of each outer-loop step based on the latent factorization?

A: Thanks for the question! We list the sequence of steps in Farzi for some added clarity and will add this as an algorithm block in the appendix:

1. Before any optimization we randomly initialize $\tilde{\mathcal{D}}_{syn}$ and $\mathbf{M}$
2. Sample a batch of fake sequences from $\tilde{\mathcal{D}}_{syn}$
3. Perform one outer loop step as demonstrated in figure 2
4. Repeat from step 2, till convergence

We would like to end by requesting you to kindly reconsider your final rating as it strictly decides the outcome of this paper, which in your own self-opinion contains important technical contributions, and contains strong empirical evidence.

We really thank you for your time and effort in reviewing the paper and providing us with your valuable feedback!

Below are the responses to your important and well-needed questions, and please let us know if we can add / expand on something:

Q: It is not clear whether this method will be practical and scale for larger language models and larger datasets.

A: We split the response into two parts, one for larger datasets, and one for larger models:

• The model is out-of-the-box scalable to as large datasets as we want, simply because the model only requires batches of the dataset we want to summarize. For more context, the largest dataset used in the paper (Netflix) has almost half a million data-points, and we only use a batch size of $512$ in the outer loop. Even with such an extremely small batch size, only $2,000$ sequences synthesized by Farzi can train models equally well as training on all half a million sequences in the original dataset.
• Being restricted to an academic research environment, we don’t have the compute necessary to experiment with large language models (LLMs). Even though it would be a well-timed application to LLMs, we believe Farzi paves a clear way for future research in data distillation for LLMs. Nonetheless, we would like to point out that Farzi’s reverse-mode Adam derivation makes it possible to perform data distillation for LLMs using no more than the same hardware that is used to train such LLMs. In other words, the memory complexity for Farzi is the same as training the base model.

Q: There is not a clear analysis of the total time gains of this method.

A: Great question! Since this question was also raised by other reviewers, we have put the results in a separate comment up top for the sake of not repeating, and are happy to present our impressive results!

We would like to end by requesting you to kindly reconsider your final rating as it strictly decides the outcome of this paper, which in your own self-opinion is of good practical value, and contains strong empirical evidence.