Paper: Farzi Data: Autoregressive Data Distillation (ICLR-2024)

Reviewer comment (Writing):
While the actual prose of the paper was generally clear and easy to read, there are some major concerns with notation/presentation that limit understanding of some of the main contributions of the paper. There are many cases where important notation is not defined. For instance, $\mathrm{Rep}(\mathcal{F}, \mathcal{D})$ is defined in the Appendix, but not the main text, and is critical to interpreting Theorem 3.1. It is not stated what the terms $d\mathbf{m}$, $d\mathbf{x}$, and $d\mathbf{w}$ in Algorithm 1 are supposed to be, so it is impossible to determine if the expressions are correct or not. How to construct the output of the algorithm from these quantities is also not clear. What is the correspondence of the quantities in Alg. 1 to the DD problem, i.e., what will we actually update using the meta-gradient once we know how to compute it? Some (but not all) of these details can be found in the Appendix, but as they are critical to being able to understand the results, they should be moved to the main text and given appropriate explanations.

Author response:
**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!

Reviewer comment (Writing):
Stylistically, there is also some nonstandard notation. For instance, $\mathcal{O}(100)$ (3rd bullet point, pg. 2). I suppose the authors meant "on the order of 100x", but big-O notation has a mathematically precise meaning that doesn't make sense here. Another instance is Proposition 3.2. "Correctness of Algorithm 1, Line 13" is not a complete mathematical statement (or a complete sentence). The result should be stated completely and precisely.

Author response:
**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!

Reviewer comment (Theory):
The proof of Proposition 3.2 is suspect, and this relates to the notation problems in Algorithm 1. What is meant by $d\mathbf{m} = d\mathbf{m} + \frac{\partial w_t}{\partial m_t} \cdot d\mathbf{w}$? Is $w_t$ supposed to be $\mathbf{w}_T$, or is this expression meant to be a recursive formula? What about the formulas for the other quantities, and how are these combined to compute the meta gradient?

Author response:
**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.

Reviewer comment (Experiments):
The authors mention that training with the reference trajectories $\Omega$ is important for obtaining the best performance, as compared with training only from randomly initialized networks. However, it is unclear if this might just have been the result of a greater number of training steps when learning the distilled dataset. Specifically, are the results in Fig. 6(b) obtained with the total number of meta-gradient steps held constant, or do the additional precomputed trajectories result in more meta-gradient steps?

Author response:
**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.

Reviewer comment (Reproducibility):
It is not clear exactly how the precomputed trajectories were used. The assumption is that instead of training the network in the inner loop only from random initializations, the network in the inner loop is initialized with parameters from one of the training trajectories. Is this correct?

Author response:
**Q: How exactly were the precomputed trajectories used?**

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.

Reviewer comment (Experiments):
Why isn't FMLP also used as a teacher network in Table 1?

Author response:
**Q: Why isn't FMLP also used as a teacher network?**

A: Great catch! This is motivated primarily by FMLP being very similar to SASRec in its architecture and inductive biases and we’re mostly interested in ascertaining cross-architecture generalization, e.g., Transformer $\mapsto$ RNN and vice-versa. This allows us to save experiment time and not compute the full 3x3 matrix in Table 1. We’ll make sure to add this clarification in the main text.

Reviewer comment (Theory):
It is said that the complexity of the dataset distillation algorithm scales by the size of the vocabulary (page. 4) and the size of the sequence that we wish to model. I can see the latter to be the case, since the loss will now be summed over the entire sequence as opposed to one forward pass (so the complexity of the forward pass is increased). However, I do not see how the time complexity increases with the vocabulary size. Do we mean space complexity? Also, more than the forward pass the dominant factor in dataset distillation is the computation of a bunch of hessian vector products in the meta gradient. Those terms do not depend on the vocabulary size either… please clarify.

Author response:
**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.

Reviewer comment (Theory):
It would be nice to provide an intuition for what is saving the memory, making things O(1) in memory. Currently the big algorithm block does not provide an intuition for how this approach is O(1) in memory regardless of the number of timesteps of unrolling. This is important to clarify, since this is an important contribution, if clearly explained. If this approach is essentially gradient checkpointing, then it is worth noting that Deng and Russakovsky already implement a version of this in their code.

Author response:
**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).

Reviewer comment (Theory):
Looking at Eqn. 2, I am a bit puzzled as to how \Omega, namely the trajectories from the real data are incorporated in the DD process. From what I am able to understand, \theta_0 \sim Omega -- namely the init is sampled from the pretrained trajectories, and then from the right hand side of eqn. 2 I understand that the rest of the trajectory is obtained using Adam on the synthetic data. Where is the role of the pretrained trajectories then? Please explain.

Author response:
**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.

Reviewer comment (Experiments):
The authors touch on language datasets as a motivation, however do not study this (or other large-sequence tasks) due to practical model/sequence length scaling constraints. Are there reasonable paths forward that would allow this to scale to longer sequence lengths/larger models?

Author response:
**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.

Reviewer comment (Experiments):
Have the authors thought about cases where there is significant noise in the training corpus? Given that the loss is computed with respect to the original dataset, it seems like this could be a problem if one ever tried to directly filter a noisy web-crawl.

Author response:
**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.

Reviewer comment (Experiments):
Even though this paper makes interesting contributions to the DD literature for autoregressive tasks, it is not so obvious that it would be very helpful for much larger text corpora and large language models with millions or billions of parameters. The memory footprint might end up being very large, rendering the whole framework infeasible. Furthermore, a compression rate of 0.1% may not be extremely helpful for very large datasets consisting of billions of sentences. This may limit the applicability of FARZI to settings consisting of "reasonably large but not very large" language corpora.

Author response:
**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.

Reviewer comment (Writing):
In Equation (2), \Omega is a set containing initializations for the inner loop, if I understand correctly. But instead of picking the initialization randomly, these come from a small number of training trajectories on the full dataset. If that is true, then the \theta_i in the definition of \Omega has nothing to do with the update rule for \theta_t in Equation (2). This may still be confusing to some readers though because the same symbols are used (theta with a subscript), so the authors may want to clarify this point (i.e., what exactly is in \Omega).

Author response:
**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.

Reviewer comment (Experiments):
I was not clear how exactly the authors chose the final hyperparameters for each setting. Did they exhaustively try all corresponding combinations in the hyperparameter table and pick the best one?

Author response:
**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.

Reviewer comment (Reproducibility):
Is a new synthetic batch created at the beginning of each outer-loop step based on the latent factorization?

Author response:
**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

Reviewer comment (Experiments):
It is not clear whether this method will be practical and scale for larger language models and larger datasets. It would be great if the authors can elaborate on this.

Author response:
**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.

Reviewer comment (Evaluation):
There is not a clear analysis of the total time gains of this method in comparison with training from scratch. Providing some values would make the case for this method more compelling.

Author response:
**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.