CoLoR-Filter: Conditional Loss Reduction Filtering for Targeted Language Model Pre-training

Summary

This work propose a novel data selection method, FRHOzen (Frozen Reducible Hold Out Loss), which leverages an empirical Bayes-inspired approach to derive a simple and computationally efficient selection criterion based on the relative loss values of two auxiliary models. They provide an empirical evaluation of FRHOzen on two language modeling tasks: (1) selecting data from C4 for domain adaptation to evaluation on Books and (2) selecting data from C4 for a suite of downstream multiple-choice question answering tasks.

Strengths

1.They analyze a family of loss-based approaches for targeted selection of pre-training data, propose a simple approach that outperforms existing methods, and provide some preliminary evidence of favorable scaling properties. 2. They analyzed and compared the computation cost of FRHOzen with other similar methods and proved its computational efficiency.

Weaknesses

1.I think this work could not match its objective of selecting optimal subsets of data for language model pre-training. The evaluation is done by fine-tuning an already-trained LM using downstream data, which could not prove the selection could work for LM pre-training. In fact, a pre-trained LM usually should not be specifically optimized for a particular downstream task (with regard to its domain-specific data). 2.In addition, the evaluation are based on olmo, which is an decoder-only LM. Exisiting pipeline don’t fine-tune for an decoder-only LMs, usually sft is needed for prompt alignment. I doubt the evaluation results may not be applicable for the real scenario. 3.The assumption of a pre-set budget n is less valid (in Algorithm 1), and the total number of optimal n is hard to predict during pre-training

Questions

1.Could you provide more details for how you fine-tune the olmo using downstream data? What’s the loss function? 2. What is the unit of data point x in Equation 6? 3.Is the selected optimal data used to optimize the fine-tune phase or the pre-train phase of training? 4.What’s the difference of data points and sequences (Section 4.2 Line 215)?

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

As the authors have stated comprehensive in the limitation Section, I have no further comments for this part.

Summary

This work proposes a simple, intuitive approach for data selection based on empirical bayes formulation minimizing the difference in the likelihood assigned to a candidate training sample by a model trained on a base distribution and a model trained on the base distribution plus a smaller sample of high quality target (test) data of interest. The method is scalable in that it leverages only standard training computation plus additional forward passes and empirical results suggest that it enjoys scale transfer properties where small scale experiments can be used to select data for larger training runs.

Strengths

- The setup is easy to follow, intuitive. Theoretically principled from a bayesian perspective while still feasible in practice is rare :] - Connection to related algorithms is well detailed including complementarity and different tradeoffs against contemporary RHOLoss. - Use of actual conditional-marginal loss gap is a strength of this method over DSDM and other influence function techniques which have been demonstrated to generalize poorly to realistic scenarios for both computational and broken assumption reasons. - The scaling results, headline figure, but last section... are impressive. Suggests that this simple method is worth further empirical exploration and expenditure of compute in the future.

Weaknesses

- The domain transfer experiment should be broadened. Does the result of figure 2 only hold this favorably for the Project Gutenberg Books downstream target? it is possible that improvements are only strong when there exist very distinct sub-distibutions in the prior that match the target. It would improve this empirical section if the authors considered more than one downstream target distribution. - (More minor, academic constraints assumed) Only one model architecture and scale extrapolation test setting are considered. It would be more convincing if another model family were considered and a few more scales, especially beyond 1B, as some trends in O(100M) models change dramatically beyond a few billion parameters.

Questions

- A diversity term seems reasonably easy to incorporate. Did the authors experiment with a diversity regularizer term or step of any sort? - Can we explain why conditional-only worsens with relaxation of the pre-sampling efficiency constraint tau? (Figure 2)

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

- See weakness.

Summary

The paper presents FRHOzen, a new data selection method for targeted pre-training of language models, which uses an empirical Bayes-inspired approach to derive a simple, efficient selection criterion based on the relative loss values of two auxiliary models. Evaluated on tasks such as domain adaptation from C4 to Books and multiple-choice question answering, FRHOzen consistently outperforms training on eight times more randomly selected data. It also scales effectively across model sizes, with data selected by 150 million parameter models yielding improvements when used to train a 1.2 billion parameter model.

Strengths

1. This paper is well-written. The comparison and discussion is very sound, e.g., section 4.2 on computational cost. 2. The proposed method is quite effective -- outperforms training on 8x as much randomly selected data. 3. The method's effectiveness in data selection is transferable across models of different sizes, making it scalable.

Weaknesses

1. Section 2.1, Bayesian Data Selection, lacks rigorous derivation. This part can be considered an intuitive understanding of Bayesian optimization, but it does not constitute a strict derivation. 2. The acceleration effect brought by this paper is quite significant, but the experimental setting uses Books as D_{down} and tests it on the Books held-out set. However, in real-world scenarios, the downstream dataset used and tested should be more general, such as a combination of multiple corpora or multiple end tasks. I hope the authors can confirm that the acceleration effect brought by this paper is applicable to more general datasets.

Questions

I am curious, when the model is larger (for example, at 7B), will the acceleration still be as significant （8x) ?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Summary

This paper proposes a data selection method that can improve the performance of language models on downstream tasks. The method uses two auxiliary models: one pretrained on the pretraining dataset, and the other is finetuned on the downstream task using that pretrained model. Then the method selects the data that has the largest difference in loss between the two models. Authors formulate the method using Bayesian optimization and show that the method is trying to maximize the likelihood of sequences that will appear in the downstream task. The authors conduct experiments using a 150M parameter model and show that it can be transferred to a 1.2B model. Using the selcted data, the model can achieve better performance on tasks, even when it's pretrained on much fewer tokens.

Strengths

1. The method is novel and intuitive. The idea of selecting data based on the difference in loss between two models is interesting and can be easily understood. The method is also well-motivated and well-explained in the paper. 2. The method is well-formulated. It can be viewed as maximizing the posterior likelihood of the downstream sequences. 3. The authors have thoroughly discussed the relevance of the method to other data selction methods. 4. The experiments can very well support the claims. The method is shown to be effective in improving the performance of the model on downstream tasks. Moreover, the method seem to reduce the computational cost of pretraining because the model can be pretrained on much fewer tokens when using selected data.

Weaknesses

1. The paper has an aim of "Pre-training", but the models being tested are relatively small (150M and 1.2B). It's unclear whether the method can scale to larger models (eg, 10B, 100B). The authors should test the method on larger models to show the scalability of the method. 2. In order to select the pretraining data that can improve the downstream performance, the method requires knowing the downstream tasks ahead of time. This is a limitation of the method because in practice, we might not always know the downstream tasks when we are pretraining the model. And people using want the pretrained model to be versatile and not just good at a few tasks. 3. It's not very clear what will happen to the downstream tasks that are not included in the data selection process. It's unclear whether there are side effects of the data selection process on unknown tasks. 4. As the authors mentioned, "In particular, the FRHOzen objective no longer encourages the selection of a diverse dataset". So a lack of diversity and overfitting might be a concern when using the selected data. 5. Using a probability-based view, the method is effective in improving the likelihood of the downstream sequences. However, it does not explain or guarantee anything about RLHF/preference-based fine-tuning. And it's unclear how it will impact the model's safety and robustness.

Questions

1. What will happen to the downstream tasks that are not included in the data selection process? Will the model perform worse on those tasks? 2. How do you choose the downstream tasks? Any criteria? 3. Do you only consider maximal likelihood for the downstream tasks? What about other downstream finetuining methods, such as RLHF, DPO, etc?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations with respect to approximations, data diversity, and computational cost are discussed.

Response to Rebuttal

I appreciate the authors thorough response to my initial review. While I am happy with my strong score of 7 as is, I reiterate to the review pool that I believe the work is high quality, the method is simple and well motivated, and the experimental results are promising and thus recommend it for acceptance. I think that future research would benefit from exploring the author speculation as to question 2. regarding why FRHOzen's selection signal seems more useful than the conditional method. Separating easy tokens from useful tokens is certainly at the crux of all data selection and valuation work in language modeling (throwing out trash tokens is just the table stakes) and so mechanistic analysis of why a method works at this fine-grained a level is just as important as increased scale and scope of the training experiments.