Approximating Language Model Training Data from Weights

Thanks for the suggestions. We have changed our setup to mimic instruction tuning, added a new task (SFT), and scaled to larger models. Full experimental results are available in the general response, but we provide a point-by-point rebuttal below:

The paper assumes access to both the base and fine-tuned model weights, the training algorithm, optimizer type, and training steps. In practice, many open models do not release all these details just the weights.

This is a fair criticism. We do provide empirical data in our response to Reviewer N6vu showing our method works reasonably well regardless of optimizer choice. And our method does not require access to the true amount of training steps. We do require both sets of model weights and knowledge of the training task, which is not always the case, although it is somewhat common (e.g. the LLAMA base and instruction-tuned models, or DeepSeek base vs. R1).

The experiments are restricted to small-scale classification tasks and relatively small models. Unclear if the findings will translate to realistic scenarios such as instruction-tuned LMs.

We have rewritten the main experiments to use an setup closer to instruction tuning for classification tasks, as well as larger models for a pretraining data selection task. Our results seem to scale to new tasks as well as 10x larger models. Please see the general response for empirical data.

Optimal transport on sentence embeddings is used to assess recovery, but the authors do not discuss how well this correlates with actual overlap or usefulness of recovered data. It’s possible to achieve good OT scores without semantically meaningful overlap.

In section 6, our analysis shows precisely this: that increasing the percentage of leaked data (i.e., overlap of recovered data) decreases our OT scores, whereas task accuracy increases (see Figure 3).

Importantly, it is not possible to obtain good OT scores without semantically meaningful overlap, because the optimal transport distance is computed between sentence embeddings.

For example, given two sets of sentence embeddings ${e(x_i)}^n_{i=1}$ and ${e(y_j)}^m_{j=1}$ and their respective empirical distributions $\mu$ and $\nu$, the OT score (Wasserstein-1 distance) can be calculated as:

$$W_{1}(\mu, \nu) \;=\; \min_{\pi \in \Pi(\mu, \nu)} \sum_{i=1}^n \sum_{j=1}^m \pi_{ij} \left\lVert e(x_i) - e(y_j) \right\rVert_{2}$$

The OT score is fundamentally an embedding metric: If embeddings for A and B are far apart, OT will register a large cost, and vice versa.

The comparison to “random” and “top-k gradient” baselines seems insufficient. More sophisticated baselines from data distillation or data selection literature (e.g., influence functions) are not considered. I understand they are slow, but the experimental setup considered in this paper will work just fine with influence functions.

Importantly, most data-selection methods--including influence function formulation--rely on reference/validation data to perform selection (Pruthi et al., 2020; Xia et al., 2024).

Specifically, they compute:

$\nabla \ell\bigl(w_{t_i},\,z\bigr)\;\cdot\;\nabla \ell\bigl(w_{t_i},\,z'\bigr)$

wher $z$ is a training example and $z'$ is the reference/validation example.

In our setup, we only have access to model weights but not reference data. The only other baselines that don’t depend on reference data use perplexity (i.e., selecting examples based on the model’s loss). We have included additional perplexity-based baselines to approximate training data in the general response (please refer to the table in the main response).

Reference:

Pruthi, Garima, et al. "Estimating training data influence by tracing gradient descent." Advances in Neural Information Processing Systems 33 (2020): 19920-19930.

Xia, Mengzhou, et al. "Less: Selecting influential data for targeted instruction tuning." arXiv preprint arXiv:2402.04333 (2024).

What happens when the seed dataset is large and diverse (e.g., Common Crawl)? Does your method still recover task-relevant data?

We haven't evaluated with exactly Common Crawl due to the large computational requirement of computing per-example gradients for the entire web. But we did evaluate with data from the MS Marco corpus, which also includes randomly sampled documents from webpages, similar to a typical Common Crawl distribution. We still observe that our method is able to find task-relevant data when using MS Marco as seed data.

Thanks for the review and helpful comments. We performed a new optimizer ablation showing that our assumption of optimizer used is not necessary; we also added a new task (SFT) which is close to pretraining data selection. Please see our point-by-point response below for more explanation.

The paper started with absense of training data and availability of model weights but then claims to be a good coreset selection algorithm which is orthogonal to the main story. Refer page 2 para 1 "disjoint from true training set ...”

To clarify, our method performs coreset selection on a seed dataset that is distinct from ground truth training data. We assume no information about the ground-truth training data: in the best case, the training data is a subset of the seed data; in the worst case, the two datasets are disjoint. Regardless, our goal is to select from the seed datasets to reach same model performance as final model checkpoint.

Based on my understanding AdamW is widely used for LLM training not sure why the authors are using adam.

Thanks for the insightful question. We ran a new ablation comparing the optimizer used to train the ground-truth model (horizontal) vs. the expert model, after selection (vertical). We show results for SELECT (top) and our random baseline (bottom).

Data selected by SELECT

Data selected randomly

Numbers are test perplexity from SELECT in our new pretraining setting with LLAMA-1B. We observe that using AdamW to train the expert model slightly deteriorates performance on our method, likely due to the weight-level signal lost to weight decay. The default SGD, which has no momentum term, provides even better selection signal than raw Adam. (We plan to add this experiment to the ablation section, space-permitting, in the camera-ready version.)

The baselines are quite weak -- for instance show us first what are the open weight model performance on the test set of the datasets you are using. Then claim that even without having the train set we have achieved a certain level of performance. This method is beating random subset selection baseline so, it has some merit already.

To clarify, the task is classification: most open-weights models, before finetuning, will perform poorly compared to our random subset selection baselines, which are allowed some finetuning data. However, we do show the performance of the final model (trained on the true, e.g. not extracted, training data) on the test set in the last row of our main experiment (please see Table 1).

We included additional baselines for using model’s perplexity to approximate training data in the table in the main response. We show that SELECT still generally outperforms other baselines, although the max-likelihood method is competitive.

Can you use Openwebtext or any pre-train data and use your method and show how much perplexity difference it has with the open weight base model. As all the results are shown in smaller scale data. Openwebtext is also a small size data should fit in academic compute budget. This will help us gauge the method's impact on pre-training.

That’s a great observation. We rewrote our codebase to accomodate larger models and alternative (i.e. language modeling) loss and ran experiments on the 1B and 3B LLAMA 3.2 models. While our method is applicable to pre-training, we believe it is somewhat unrealistic to assume access to the weight initialization. Instead, we demonstrate results on supervised fine-tuning (SFT) using the same language modeling loss, on top of the pretrained model. Our method outperforms all baselines on both the 1B and 3B models.

Why computing gradients is neccessary and can't a similar result be shown with just computing loss which is much cheaper.

This was a great suggestion– thanks. We added both highest-loss and lowest-loss selection baselines in both sets of new experiments in the general response.

For gradients computation the base model is used and for pseudo labelling the finetuned model is used. Can we use the finetuned model for both?

We have tried this but found the gradients of the finetuned model (perhaps surprisingly?) are not useful for the task, and perform worse than random selection. We will provide these results in the appendix of our final paper.

---

The full set of improved main Table 1 and Table 2 results are available in the general response. Please let us know if you have any questions.

Hi! Thanks for the helpful review. We provide a response to each of your suggestions below:

While using the last layer of language model can be high-resolution enough for synthetic data generation, I don't immediately see how this translates to your method.

Recall our formulation of the method in section 4:

$\arg\max_{x \in \mathcal{D}} \Bigl[\,\nabla \ell\bigl(x; \theta\bigr)\;\cdot\;\bigl(\theta_{f} - \theta\bigr)\Bigr].$

Computing the above equation requires taking inner product of the gradients with respect to the training data $x$ and model weight’s difference $\theta_{f} - \theta$.

However, directly using such high-dimensional gradient vectors as features for dataset selection is very computationally expensive, so we use only the last layer alongside random projection to construct meaningful low-dimensional gradient features.

At least from the way the method is presented, it seems to rely a lot on the seed selection, so if the proposed method wants to be used to approximate training data, this seems like a major challenge as most of the datasets for current LLMs might not even be publicly available, so your method is not applicable.

Indeed, seed selection is crucial, as seed datasets must be close to the ground-truth training data. Lacking access to the true training data is a challenge not only for our method but for the entire field of training-data recovery.

However, our method remains applicable even when ground-truth training data is not publicly available. In our experiments, we are able to identify lexically and semantically relevant data and recover performance comparable to that of the final model weight.

Even though you say that even in the cases when the training dataset is not available, your method also works for data selection, there is a lack of comparison between your method an more traditional techniques for data selection.

Importantly, traditional data selection relies on reference/validation data to perform data selection: lexical methods or embedding-based methods select training data based on lexical or semantic similarity with reference data [Albalak et al., 2024].

In our problem set up, we only have access to model weights but not reference data. The other baselines that didn’t depend on is perplexity data selection, which select data based on model’s loss. We included additional baselines for using model’s perplexity to approximate training data in the table in the main response.

Reference:

Albalak, Alon, et al. "A survey on data selection for language models." arXiv preprint arXiv:2402.16827 (2024).

There is an important discussion missing in this paper of how the method scales computationally as the model size and dataset sizes increase, this is not evident from the paper as the models and datasets used were rather small compared to modern standards.

We agree. We have since provided additional experiment on instruction-tuning on a larger 1B language model in the main results. Our method seems to scale well to 10x larger models, as well as a more modern task setup of selecting pretraining data.

How does your method scale to larger models?

We have new experiments with larger models in the general response, please let us know if you have any questions.

Figure 2 is completely unreadable for a colorblind person. ... Please refrain from using notation like \theta^T ...

Thanks for these small suggestions. We will fix the figure style and change the superscript to subscript in the final version (as we can't update the PDF during this rebuttal period).

Thank you for the constructive feedbacks and and careful reviews. We provide point-to-point response to your concern.

The current setup may be a bit synthetic.

1. First only having the task classification task seems to be a bit limited for such language models -- also the training setup doesn't use the typical instruction fine-tuning setup (as it doesn't use a textual prompt but simply requiring the last token to be a class label).

Thanks to your suggestion for how to improve our experimental setup: we have adjusted our experiments to use a more conventional instruction tuning setup and scaled to larger models (1B and 3B parameters). Changing the task format slightly improved performance. More importantly, we rewrote our codebase since submission to accommodate arbitrary tasks, which made it simple to test our models on supervised finetuning data selection. We show results in the general response; our method is useful for supervised finetuning. And although the new format makes the baselines better too, SELECT still outperforms all baselines, including max-likelihood data selection.

2. Furthermore, perhaps what's interesting in this setup is how to avoid some sensitive (exact) training data being recovered; however this method does not have the guarantee for recovering verbatim copies of the training data---all it can do is "approximating" the data---unless it can have access to a good set of seed data. In that spirit, it is closer to dataset distillation, which the authors seem to be aware of. In that case, I suggest the authors change the framing of the paper a bit.

Thank you for this suggestion. You are right that our method approximates rather than recovers exact training data, and that the connection to dataset distillation is more apt than our current framing suggests. We will revise the introduction to better reflect this distinction and clarify how our approach differs from traditional dataset selection (given the absence of validation data).

I think there are several additional experiments for this setup:

1. As one important baseline, one can select the data based on the logprobs: for example, we can score the sequences using the fine-tuned model, and higher logprobs can indicate more confidence or closer to the training data. Running a top-k experiment using the logprob seems to be an important baseline and I suspect it can achieve nice results.

This is a great point. We added this baseline for all main experiments: for classification, we use the likelihood of the autolabeled class to select data; for SFT, we use the perplexity of the sequence.

In general max-likelihood data selection is a strong baseline, especially for SFT. In some cases its performance comes exceptionally close to that of the gradient-based SELECT (for example, 2.63 vs. 2.61 test perplexity on SFT data selection). See the main response for full results.

2. Since the authors only use the last layer gradient, during the training of the final model, one should also consider freeze all the layers expect for the last one. Also the authors can consider a layer selection method -- identifying the best layers to choose.

This is true: in some sense, the data is selected to have a useful gradient only for the last layer of the model. We ran a small experiment during this rebuttal to examine the effect of unfreezing the last layer while training the expert model.

In the SFT experiment detailed in the main response, we compared the default setting (all layers unfrozen) to freezing the last layer. We observe that freezing the last layer worsens perplexity from 2.09±0.07 to 2.26±0.11, indicating that the expressivity of full-finetuning may be necessary important for optimal performance.

3. An interesting analysis for the experiment in figure 3 is that the authors should report how much of the leaked data is selected using different algorithms.

Due to time and compute constraints, we are not able to run this full sweep over data sizes during the rebuttal period. However, we can add this to the final version of the paper (i.e. we can add two additional lines to figure 3: Top-K and random selection).

Questions To Authors:

• In figure 3, the x axis should be 10, 20.. instead of 0.1, 0.2, .., if I understand it correctly.
• the reference to table 2/3 is not correct (it writes section 7 in the last paragraph on page 8)

Thank you for the attention to detail. We will fix typos and references in camera ready.