Task-Adaptive Pretrained Language Models via Clustered-Importance Sampling
Authors: David Grangier
Simin Fan
Skyler Seto
Pierre Ablin
Venue: ICLR 2025 Poster (2025)
Keywords: task-adaptive pretraining
language models
importance sampling
domain adaptation
Primary area: foundation or frontier models, including LLMs
TL;DR: Scalable clustered importance sampling for LM pretraining: End-task accuracy is improved when pretraining on a generalist dataset resampled with task-specific importance sampling weights.

Abstract: Specialist language models (LMs) focus on a specific task or domain on which they often outperform generalist LMs of the same size. However, the specialist data needed to pretrain these models is only available in limited amount for most tasks. In this work, we build specialist models from large generalist training sets instead. We adjust the training distribution of the generalist data with guidance from the limited domain-specific data. We explore several approaches, with clustered importance sampling standing out. This method clusters the generalist dataset and samples from these clusters based on their frequencies in the smaller specialist dataset. It is scalable, suitable for pretraining and continued pretraining, it works well in multi-task settings. Our findings demonstrate improvements across different domains in terms of language modeling perplexity and accuracy on multiple-choice question tasks. We also present ablation studies that examine the impact of dataset sizes, clustering configurations, and model sizes.

Published as a conference paper at ICLR 2025
TASK-ADAPTIN PRETRAINED LANGUAGE MODELS VIA CLUSTERED IMPORTANCE SAMPLING
David Grangier, Simin Fan, Skyler Seto, Pierre Ablin Apple
ABSTRACT
Specialist language models (LMs) focus on a specific task or domain on which they often outperform generalist LMs of the same size. However, the special- ist data needed to pretrain these models is only available in limited amount for most tasks. In this work, we build specialist models from large generalist train- ing sets instead. We propose novel method, ClusteRed Importance SamPling (CRISP). CRISP clusters the generalist dataset and samples from these clusters based on their frequencies in the smaller specialist dataset. able for both pretraining and continued pretraining, and works well in multi-task settings. CRISP performs favorably compared to other methods that adjust the training distribution of the generalist data with guidance from the limited domain- specific data. Our findings demonstrate improvements across different domains terms of language modeling perplexity and accuracy on multiple-choice question tasks. We also present ablation studies that examine the impact of dataset sizes, clustering configurations, and model sizes.
scalable, suit-
INTRODUCTION
Generalist language models (LMs) can address a wide variety of tasks, but this generality comes at a cost (Brown et al. 2020). It necessitates large training set representative of all prospective tasks, as well as large model to fit such comprehensive dataset. Specialist models forgo this generality and fit model for limited domain task. In their narrow specialty, such models can achieve better accuracy at given model size (Kerner 2024).
Pretraining specialist is interesting when two conditions are met: (i) the targeted task justifies the cost of training a dedicated model and (ii) a specialist dataset large enough for pretraining is available. Condition (i) is dependent the targeted application and its potential cconomic benefit. Condition (ii) is more limiting since modern LMs are commonly pre-trained on datasets larger than 100B tokens"] an amount that cannot be commissioned for most applications.
This work considers relaxing condition (ii) and studies methods to train a specialist model when specialized data is scarce. Given large generalist dataset and small specialist dataset, propose to modify the distribution over the generalist dataset guided by the scarce specialist dataset. Training model the modified distribution gives specialist model with better accuracy than generalist model of the same size.
We study this setting across different specialization tasks including domain-speeific language mod- cling (medical, encyclopedic domains) and end-tasks (scholar exams in science and humanities, reasoning questions). We compare different strategies manipulate the pretraining distribution. We evaluate strategies based on text classifiers, saliient and importance sampling (IS). Although IS is rarely used for LM data selection, we build upon on simple IS recipe based on clustering (Grangier et al. 2024b) and report that the resulting method systematically outperforms alternatives. Our IS recipe clusters the generalist set and computes the cluster histogram over the specialist data. Then, for pretraining, generic data is sampled according to the specialist histogram, see Figure T We show the empirical benefit of this method varying model sizes (350m to 7B pa- rameters), the amount of generalist data and the amount of specific data. We assess both perplexity gains for language model adaptation and accuracy improvements for multiple choice question tasks.
100B tokens Im books 60x the annual publication of the top English language publisher (Lee 2021

---

Published as a conference paper at ICLR 2025
This paper presents exhaustive comparison over different model sizes (350m, 1.3B, 6.8B) and different numbers of clusters (scaling from 64 to 16m clusters with hierarchical clustering). We consider different tasks, both for language modeling and multiple-choice questions. We also explain the impact of hyperparameters such as the clustering representation and number of clusters. We study IS in the context of multitasking and continued pretraining- We also perform ablations with respect to the generic pre-training set size and the specialization data size.
RELATED WORK
Generalist vs Specialist LMs Generalist LMs address tasks for which they have not been explicitly trained (Brown et al. 2020) provide good initialization for fine-tuning dedicated model (De- vlin et al. 2019). ). Nowadays generalists compete with dedicated models on many tasks (Jiang et al.) 2024 Dubey al.]
2024). Success, however, comes at a price: generalist must be much larger than specialist for the same accuracy. For instance, on English-to-German translation, the 175-B parameter generalist GPT-3 (Brown et al., 20200) is less accurate than a 136m-parmmetrr special- ist (Sennrich et al. 2016a). For neural LMs, the parameter count directly impacts training and inference costs.
Specialist large LMs exist in domains where large amounts of specialized texts are available. Cor- pora with billions of tokens enable pretraining or continued pretraining, a generalist pretraining phase followed by a specialist one (Gururangan et al., 2020; Parmar et al. (024 Domains with specialist models include medicine and biology (Lewis et al. 2020) Labrak et al. 2024, Bolton
et al. 2024), computer programming and mathematics (Lewkowycz
al. [022] Rozière et al.)
2024 Azerbayev et al., 2024) and finance (Wu et al. 2023] Xie et al. 2023a). When specialist
data available in limited amount, task-adaptive data-selection methods train specialist models on generalist data instead.
These selection methods over-sample generalist data that aids model generalization in the target domain. For masked LMs. Gururangan et al. (2020) observe that contin- ued pretraining improves the performance end-tasks when using data with high vocabulary over- lap with the targeted task. For machine translation (MT), Aharoni & Goldberg (2020) show that a task-adapted pretraining dataset can be selected from generalist dataset using the nearest neighbors of small specialist set. Their nearest neighbor classifier relies on BERT sentence distance (Devlin (et al.] (2119). Still for MT, other works have used other types of classifiers. In particular, contrasting the scores of two I LMs (generalist and specialist) is popular (Moore & Lewis 2010: Axelrod et al.,
Task-Adaptive
2011 Wang et al. fine-tuned BERT
2018 JJunzyss-owmmntee Dowmunt 2018). Other classifiers include logistic regression or (Iter & Grangier (001)) Outside classification, Xie et al.) (2023c) proposed to
use importance sampling for continued pretraining. They improve classification tasks by selecting pretraining data with similar distribution to the targeted domain in terms of hashed-ngrams. Im- portance sampling is also used in (Grangier et al, 2024b) and we build upon that work which adjusts the frequency of generalist clusters informed by specialist data: scale the method to millions of clusters, show that it works with larger models, and extend it beyond language modeling tasks.
A third type of methods for task-adaptative selection relies on bilevel optimization and gradient aligment (Pruthi et al. 2020. Xia et al.) 20244 Grangier et al. 2023). The pretraining distribution is selected such that the reweighted gradients from the generalist dataset mimies the expected gradi- ent from the small specialist dataset. Gradient-alignment for data selection has also been used for other purposes such as data summarization (Borsos et al. 2024), pretraining acceleration (Xie al.
2023b Fan et al. 2024) or auxiliary task weighting (Wang et al. 2020 Raghu et al. 2022 Finally is also worth mentioning data selection methods based on reinforcement learning (Liu al. 2019
Yoon et al. 2020), bayesian optimization (Ruder & Plank 2017). data models (Ilyas et al. 2022
and influence models (Yu et al. 2024).
Pretraining Data Quality Outside of domain aspects, the quality of pretraining data is also an
important topic (Wenzek et al., 2020 Dodge et al. 2021 Penedo et al. 2023 Li et al. 2024). Data quality includes removing data in other languages (Cook & Lui 2012) 2), text formatting (Xu et al.)
2024). favoring long form text Gao et al. 22021; Gunasekar et al. (0023), removing duplicates (Lee) al. 2022). It also involves balancing different sources of data with the goal of reaching better generic pretraining loss Xie et al. (2023b): Fan et al. (2024): Vo et al. (2024). Recent work also
considered filtering KKng al. 2024), correcting Chen & Mueller (024) generating Maini
2

---

Published as a conference paper at ICLR 2025
al. 2024) pretraining data with LMs. These data quality considerations are orthogonal to domain concerns: quality filters are applied alongside domain adaptation decisions (Albalak et al. 2024).
DATA SELECTION FOR AAAP-IDAETIVE PRETRAINING
We consider three methods for task aas--adeptive pretraining of LMs. Classification and gradient align- ment have been evaluated in different contexts before but not for end-tasks like multiple-choiee question answering. Clustered-baat importance sampling at scale is a contribution of this work, building upon recent work from Grangier et al. (2024b).
3.1 NOTATIONS
DE is the training dataset sampled from the generalist distribution DE, D is the specialist dataset representative of the final task, sampled from the specialist distribution D DE, The loss of model 0 on dataset D is
((; = D
Σ (x; x
1 D
Σ RED
Σ log p(,
where [D] denotes the cardinality of D and |z| denotes the length of sequence (#1...., 2|x
The perplexity of model 0 the dataset D P(D;0) == exp(C(D:0)).
3.2 CLASSIFICATION
A binary classifier trained to estimate the probability that a generalist pretraining document be- longs to the targeted domain. The classifier is learned using positive examples from D and subset of DE as negative examples. & then builds a domain-specific pretraining set
C(DR,0) () = {0 E DR such that ((x) t).
which restricts the generic dataset D to the examples with an estimated probability be in-domain above threshold t. The threshold is sensitive hyperparameter that impacts the downstream model. is validated as a trade-off between focusing on data close to the domain of interest while keeping C(Ds, f) large enough to train an LM of the targeted capacity. In our case, rely a logistic regression classifier trained over sentence BERT (SBERT) text embeddings (Reimers & Gurevych] 2019), an established classification method (Minaee et al. 2021). The SBERT representation also commonly used in data selection (Albalak et al., 2024 Xie et al. 2023c Zhang et al. 2024 Su et al. 2023). This representation is also used in the alternative selection strategies consider. As an ablation, we also evaluate the impact of the choice of SBERT (Section 5.1
3.3 GRADIENT- ALIGMMENT
Gradient-Alignment (GA) methods are common when the generic pretraining De originates from Ng different data sources S, i.e. D = - D. These methods select weights for the different sources by considering two functions 6: the pretraining reweighed loss,
ngg.
((W,DD)) = wc C(D
and the targeted loss, i.e. the loss on D in our case. The weights, on the simplex, can be in- ferred via bilevel formulation of the data selection problem (Dagréou et al. 2022): the mini- mum 8 (w) = arg ming L((w, () depends on w and task-dependent pretraining interested in weights w that minimize L(Ds; 0* (w)) wrt This formulation results in algorithms that select weights during pretraining to align the gradients of these two functions wrt (Xie et al., 2023b rely the DoGE (Fan al.) 2024) algorithm. Compared
Grangier et al. 2023). In our case, to classifiers, GA is harder to scale to large model size. This limitation is commonly addressed by finding the mixture weights with small model before transferring them larger model.
In this work, we consider generic setting where the pretraining dataset D is not pre-segmented into few data sources. Instead, we rely the k-means clustering of the Sentence BERT embeddings to identify data clusters. Clustering based on text embeddings has been used for data selection, both for quality filtering (Kaddour 2023) and domain adaptation (Grangier al. 2024a).
3

---

Published as a conference paper at ICLR 2025
Learn clusters generalist data
Get cluster histogram
Resample generalist data
from specialist data
Figure 1: Task-adaptive data selection with Clustered Importance Sampling (CRISP).
3.4 CRISP: CLUSTERED IMPORTANCE SAMPLING FOR PRETRAINING
We sketch our strategy from Figure 1 Initially, we divide the space of text into clusters. We de- compose the specialist loss and the generalist loss as a weighted sum losses over clusters. Then make independence assumption that implies that the specialist and generalist loss per cluster are identical. The specialist loss then computed as the generalist loss with a reweighing of each cluster.
Specifically, we want to identify model with a low loss on the specialist distribution D
L(D":0)
E [((:; x~D
Σ (x:; ((=
We marginalize over discrete latent variable c, the cluster variable, and write
L(D"; = ΣΣ "x: ("
(2) ΣΣ E(x; ((x:))))))))((()))
(1)
where the second equality (2) makes the independence assumption P(rle, D*) = P(xle). we make a similar assumption for the generalist loss P(z|c, DR) P(z|c), we can write both losses as
(C(c;)) and [(DS; 0) L(D²;0) co(c)())
E c< D2)
[C(c; 0)]
where we define (c;;0)= (x: 0) P(rle). We now apply importance sampling to these expec- tation, defining the importance weights as w(e) P(cDD)
P)c L(D";0)~ = Σ ((c:))(((D)) Σ C(c DR) P(c[D)=
E [c(c) cn(o)Dr)
In our experiments, we estimate the terms w(c). [(c; 0) from the finite training sets D ~ D= and DE DR. We count the number of examples in each cluster to estimate P(c) The expected loss over cluster L(c, 0) is estimated as the average loss over the generalist examples cluster e, L(Ds n K K(c); 0), where ((c) denotes the examples cluster e. This strategy there- fore only estimates D)) on the small D The term L(e; 0) is estimated over the large set as C(DE n K (c); f) which uses many more samples and hence has less variance than the estimator C(D* n (c)) over the small D
Adam) and propose Algorithm T Here, do not explicitly reweigh the loss. We instead sample clusters from their importance. This avoids frequently visiting clusters with less weight. This strategy has less variance in its gradient estimates, which help convergence (Seiffert et al. 2008 An al. 2021). This algorithm is simple and efficient when one groups the generalistexamples examples by cluster prior to training.
We train CRISP models with stochastic optimization (Kingma Ba 2015
EXPERIMENTS & RESULTS
We perform experiments with transformer LMs (Vaswani et al. 2017]. Most our experiments use models with 1.3B parameters (trained 120B tokens) and conduct ablations with 350m and 7B models (resp. trained on 40B, 350B tokens). Our settings for architectures and optimization are Appendix D
borrowed from Brown et al. (2020).
Our generalist training set is Redpj2 (Together Al Team 2023). We select this dataset as it con- tains only web-crawled data without additional interventions to help evaluation tasks (e.g. adding
4

---

Published as a conference paper at ICLR 2025
Algorithm CRISP Training
1: Parameters: T (number of steps), B (batch size)
2: Input: D (specialist set), De (generalist set)
3: Jff +- DP), VeJ
Count cluster frequency on the specialist set D > Initialize the model.
4: 80 + --nit.Model))
5: for t= port=l..... do
6:
for i 1...,,., do
7:
C: Categorical(h1)
Sample cluster id from the specialist histogram. b Sample generalist example in the selected cluster.
8:
2; ~ Uniform ( O (c))
9: end for
10: 0r f Adam AdamUpdate(OA-1.. f, (tI,...,,nI)
end for
encyclopedias, books or academic articles). Redpj2 contains over 30T tokens with our 32k byte-pair encoding tokenizer (Sennrich et al., 2016b),, Table (i Appendix We segment the dataset into non-overlapping 1,024 token windows (the model context limit) and compute SBERT embedding for every window. We cluster the generalist dataset hierarchically with a clustering tree with branch- ing 64 for 4 levels, see Appendix B The levels therefore have 64, 4,096 (= 649), 260k (= 649) and 16.7m (= 64*) clusters with an average of 540B, 8.4B, 130m and 2m tokens per cluster respectively. As an alternative to SBERT embeddings, we also consider Latent Semantic Index (LSI), i.e. singular
value decomposition over if-idf representations (Deerwester et al. 1990
Dumais 2004).
For our specialist tasks, we consider 3 language modeling tasks (LMM) and 3 multiple-choice-quese tasks (MCQ). For LM, we use Pile subsets from different domains (Gao et al., [021]: medical (Pubmed Central), programming Q&A (Stackexchange), and encyclopedic (Wikipedia). For MCQ answering, we use A12 Reasoning Challenge (Clark et al., 2018 ARC), Massive Multitask Lan- guage Understanding (Hendrycks al., 2021 MMLU), and Reward Bench Reasoning (Lambert) jet al., 2024 RWDB-R). (....)).. ARC focuses on science questions, MMLU focuses on interdisciplinary knowledge, RWDB-R focuses correct vs incorrect solutions math and programming problems. To provide a representative specialist train set D ~ D=, we split the questions into a train and test split, see Table in Appendix C
Our main results are reported with unified settings. For the classifier, the classification threshold is the main parameter. A threshold accepting 2.5% of De worked best for the runs with 1.3B models over 120B tokens. For DoGE, the method is costly to apply over many data sources/clusters and we applied over 64 clusters, i.e. learning mixture weight of dimension 64. For importance sampling, the results presented in this section relies on 260k clusters. Later, Section 5 studies ablations and parameter sensitivity. Details on hyperparameters can be found in Appendix
4.1 LANGUAGE MODELING TASKS
We evaluate specialist LMs on three domains from the Pile (Gao et al. 2021): medical (PubMed Central), encyclopedic (Wikipedia) and programming Q&A (StackExchange). We limit specialist training data from 14m tokens to the full Pile subset, up to 26.7B tokens, see Table Appendix C
As baselines, we either train only on the in-domain (specialist) data without pretraining or we fine- tune model pre-trained on Redpj2. We refer to the Redpj2 pretraining distribution as the base dis- tribution. For task-dependent pretraining, we resample the Redpj2 pretraining set using a classifier, DoGE or importance sampling for each domain. The three methods have access to 14m special- ist training tokens. In each case, the resampled pretraining set is used to train paraaeere transformer model with the same hyperparameters the Redpj2 baseline.
We report pretraining results in Figure 2 and the fine-tuning results Figure 3 For each domain, the pretraining results evaluate models trained using the resampled Redpj2 examples. The fine tuning results evaluate models where each model pretrained on (resampled) Redpj2 has been further trained on the in-domain data itself (PubMed. StackExchange, Wikipedia). All experiments consider the same optimization effort and we validate the fraction of steps spent in fine-tuning, from 3%-ft with 14m tokens (97% pretraining) to 100%-ft with 26.7B tokens (no pretraining).
5

---

Published as a conference paper at ICLR 2025
The pretraining results in Figure show that the in-domain perplexity is better with task-dependent pretraining than with generic pretraining (base Redpj2) for all methods. This gain in perplexity comes as model training focuses data close to the targeted domain: the model capacity is not used to fit the filtered out training data. Table TO in Appendix shows, for instance, that CRISP outperforms base 97.3% of PubMed but reports worse perplexity 95.9% of Redpj2.
When we fine tune the pretrained models, the advantage of task-dependent pretraining is preserved, shown in Figure 3 Task-specific pretraining checkpoints are better starting points for fine-tuning than generic ones. This shows the complementarity between task-dependent pretraining and fine- tuning. Figure also shows the necessity of pretraining: below 7B tokens, the "only specific" 1.3B model shows high perplexity. When comparing task-dependent pretraining methods, importance sampling consistently performs better after fine-tuning, even when the pretraining results are close (e.g. classifier on PubMed, Wikipedia).
Base redpo2
15:45
Clasttee CRUSP
17.0
Speait Popiesit
Spenis Hapicstt
Spectt Deggeedy
11.6
11:44
7.44
13
11.01
6,88
12
11:45
10.5
10.29
160
(a) PubMed
(b) StackExchange
(c) Wikipedia
Figure 2: Pretraining perplexities for language modeling tasks
7.0
9.5
Base mápo
DoCE
75
8.0
6.5
Classifer
7.0
Paylestry
6.0
0.5
6.5
6.0
5.0
5.5
7.0
5.0
8.5
107
102
10300
1/"
1010
10"
(a) Pubmed Central
(b) StackExchange
(c) Wikipedia
Figure 3: Fine-tuned perplexities for language modeling tasks. Task-dependent pretraining is always better than generic pretraining. The ordering of the methods is unchanged from pretraining.
72
71.2
34
70-
70,0
35.7
35-
67.9
68-
66.2
30-
64-
31.0
84-
63.5
00
DoGE:
Base relp2
245
62-
Claotfer
60
(a) Arc-E
(b) Are-C
(c) MMLU
(d) RWDB-R
Figure 4: Accuracy for multiple choice question tasks. Light colors indicate fine tuning improve- ments if any. The ordering of the methods consistent across datasets.

---

Published as a conference paper at ICLR 2025
4.2 MULTIPLE CHOICE QUESTION TASKS
Compared to LM, MCQ has much smaller specialist training sets per task, i.e. between 200k and 2m tokens, see Table in Appendix C The MCQ evaluation is also different: it uses accuracy and not perplexity. For each example, the model scores the concatenation of the question and a possible answer, for each proposed answer. The model is accurate when the correct answer is assigned the highest score (probability or normalized probability, see Appendix ED. For MCQ tasks, unlike for LM tasks, the training loss (negative log likelihood) is therefore not closely tied to the test metric.
Despite these differences, we observe a similar benefit for task-dependent pretraining compared to task-agnostic (base) pretraining. Figure 3 displays displaydssplays a similar method ordering and CRISP is consis- tently the best method. As difference with LM tasks, see Figure 4 Fine-tuning improves the base method on all datasets except ARC-E, but not enough to outperform task-specific pretraining, see Table 12 in Appendix G
observe limited benefits from fine tuning,
ANALYSIS
5.1 CLUSTERING
We study the impact of the text representation for clustering and the number of clusters. We consider two representations for clustering, the SBERT embeddings used in all other experiments and LSI We report their performance with 64, 4096, 262k and 16.7m clusters.
embeddings, see Section 4
The representation is important: examples in the same cluster are close in the embedding space. Our independence assumption, Equation assumes that the loss in cluster is the same regardless whether its data originates from D or D*, i.e.
L(D* n K(c):0) ~ L(D" (cc)
(3)
In practice, it is sufficient that the embedding space reflects the similarity of the loss gradient, i.c. if the gradients of the loss a specialist cluster DP n (cc); f), the model trained on the former improves on the later. Figure shows that the SBERT representation yields better results than LSI for all settings.
generalist cluster D n K (c) is correlated with the gradient over
The number of clusters is a trade-off, and its optimum is at 260k for most of our experiments. There are multiple factors at play when the number of clusters varies. A smaller number of clusters implies larger clusters: our hypothesis, Equation is then stronger, as it assumes loss similarity on large areas of the embedding space. At the limit, with one cluster, this hypothesis assumes that the specialist loss and generalist loss are identical everywhere. Conversely, as the number of clusters gets larger, the estimation of the cluster density on the small specialist set P(c[D*) ~ P(c gets less accurate. The estimator risks overfitting, i.e. favoring clusters frequent in the training set D but not frequent on other samples from DS. Increasing the number of clusters also risks reducing the effective training set size: the specialist data could be mostly concentrated in a few clusters, corresponding to a small fraction of the overall generalist set DE
We explore these aspects on MMLU. We first measure the number of repeated examples when training models with
20
16.7m
22
260k
4096
CRISP pretraining for different number of clusters. Figure 6
Perplexity
15
64
Perplcitty
hase
18-
shows the number of repeti-
10
tions for each quantile of the training set. Even for 16.7m
clusters, only small minority of training examples are re- peated beyond 10 times and the average number of occurrences of the training is examples 1.95,
0
500
1000
0
500
1000
Stops
Stops
(a) Perplexity reweighted Redpj2
(b) Perplexity MMLU train (plain) and test (dotted) sets.
Perplexity for CRISP on MMLU with different number of clusters. Y-scales on (a) and (b) are different.
Figure 8:
well within commonly recom- mended values (Muennighoff
jet al.) 2023a; Xue et 2023).
Even if exact repetitions do not account for the poorer performance of the 16.7m cluster setting, its
7

---

Published as a conference paper at ICLR 2025
34
70
33
65
Accuracy (19)
135
Accuncy (%)
Accuracy
32
31
300
1.51 CRISP
30
50
SBERT CRISP
29
64
4096
260k
16.7m
4096
260k
16.7m
64
4096
260k
16.7m
Num. clusters
Num. clusters
Num. clusters
(a) Arc-E
Are-C
(c) MMLU
Figure 5: Accuracy for multiple choice question tasks varying the text representation for clus- tering and the number of clusters. SBERT is more effective than LSI in all cases.
7m clusters
max=330
0.4
260k
0.4
Nom. repelitios
10
moan=1.95
16.7m
4096 clusters
insrrovemenn
0.0
maxeil 1I
mean=1.06 .04
-0.2
-0.2
-04
-044
10
maxe1
100
0.0
0.5
1.0
10
10-2
Billion
Distance quantiles
Distance quantiles (log scale)
Figure 6: Number of occurrences of each training example for CRISP on MMLU. Repeated examples increase with the number of clusters.
Figure Loss improvement Redpj2 (valid) wrt base as a function of the SBERT distance to MMLU train. Models with large number of clusters are better than base in small area near MMLU train. The gray area indicates the 25-75% quantiles for the MMLU test set.
training set might be less diverse and the model might generalize well only in small neighborhood of its training set. We evaluate if the Redpj2 examples with good perplexity concentrate around D the MMLU training set. Figure shows that the benefit of CRISP over base indeed correlated with the distance to D As the number of cluster increases to 16.7m, the benefit over base concentrate an area with very few samples. For comparison, we plot the middle quartiles [0.25, 0.75] where most of the MMLU test data concentrate in gray. We remark that MMLU test data mostly lies an area where the perplexity of IS 16.7m is low.
Figure 8 shows the perplexity for CRISP runs MMLU. On Figure 8a the perplexity is computed from the reweighed loss on Redpj2. This is the loss optimized during pretraining. shows that when the number of cluster increases the sampled training set is less diverse and corresponds to easier learning problem (< PPL). On Figure 8b) the perplexity is computed on the MMLU data itself, on the training set (plain) and on the test set (dotted). The scale of both plot is different: the resampled perplexities on Redpj2 are therefore not a good approximation of the MMLJU perplexities. This quantifies the eигог resulting from our assumption, Equation 3 We also see overfitting for 16.7m clusters, the only case with better MMLU perplexity for train than for test. Finally, we notice that the gray area in Figure fails to show that 260k cluster would have the best perplexity, which shows that SBERT distance to the training data is not the only factor explaining model performance.
5.2 MODEL SIZE
This section compares CRISP and base at model sizes. The benefit of task-dependent training is consistent across model sizes, see Figure 9 We consolidate results across sizes to report the training cost in GPUh vs accuracy in Figure 10 GPUh are measured in training hours per graphic processor (Nvidia H100). We evaluate multiple checkpoints across model sizes and sort the checkpoints by training cost. The big dots mark transitions between model sizes: they show that the the 1.3B I.S. model outperforms the 6.7B base model on ARC. This shows substantial training speedups (~-30x). Of course, smaller model is also beneficial inference.
8

---

Published as a conference paper at ICLR 2025
40
10
70
Accurncy (5)
Accuracy (5)
Accuracy
35
30
Base redpj2
50
CRISP
25
350m
1.3B
6.7B
350m
1.3B
6.7B
350m
1.3B
6.7B
Nom. trainable parameters
Num. trainable parameters
Num. trainable parameters
(a) Arc-E
(b) Are-C
(c) MMLU
Figure 9: Accuracy for multiple choice question tasks across model sizes.
33.9% faster
28.5x faster
7.0x Taster
10
37.5
Accuracy (8)
35.0
35
60
Accuracy
Accuncy
32.5
30
30.0
50
Base redpi2
25
Base redpj2
27.5
Base redpi2
CRISP
CRISP
CRISP
40
25.0
12
8
12
0
12
Cost in & GPUh
Cost in k GPUb
Cost in & GPUh
(a) Arc-E
(b) Arc-C
(c) MMLU
Figure 10: Accuracy for multiple choice question tasks as a function of training cost. The large dots mark the transition between model sizes (350m -> 1.3B - 6.7B).
max=3,070
10a
120B tokens
mcan=22.46
Nom. repetitions
700B tokens
max=512
32-
312
Full redj2
meun=3.74 74
103
Accmncy
30
Accunay
30-
baso
Basc redpi2
10a
max=11
28
288
mean=1.04
samples
10a
7008 tokeens
samples
50
500
1000
500
1000
Billion tokens
Stops
Steps
(a) Repetitions for less generic data
(b) Acc. for less generic data
(c) Acc. less specific data
Figure 11: MMLU with less training data. When the generalist set D is small (a.b), the impor- tance sampling method will samsample a small part of DR and this part will be seen multiple times during training. When this part is small, the benefit of data selection vanishes. When the spe- cialist set D is small (c), the importance sampling weights are poorly estimated and the importance sampled data might not be representative of the targeted task.
5.3 DIFFERENT AMOUNT OF TRAINING DATA
This section varies both the amount of generalist data available to sample the CRISP dataset from and the amount of specialist data for inferring the CRISP weights. When specialist data concen- trates on a few clusters, CRISP often samples generalist data from the same clusters, which can be problematic when the generalist set is small. We restrict the pretraining set 700B and 120B tokens (downsampling Redpj2 by resp. ~ 50x and ~300xx. Our pretraining runs use 120B tokens, a base never repeats in all settings. When CRISP is applied, some tokens are repeated. Fig- ure Taa shows that, when restricting to 120B tokens, the number of repetition becomes high (22.5 average) and CRISP is ineffective after 256k steps.

---

Published as a conference paper at ICLR 2025
Table I: Accuracy (%) for Task Transfer and Multitasking, Importance Sampling on MMLU and on multitask improves all tasks compared to baseline.
Model
Evaluation Tasks
ARC-E
ARC-C
MMLU
RWDB-R
Multi
Base
Redpj2
58.4
27.5
30.1
62.2
45.1
CRISP
ARC
71.3
38.6
28.9
60.9
48.2
MMLU
63.4
28.7
33.4
65.2
48.2
RWDB-R
42.4
23.4
26.4
70.1
43.1
CRISP
Multi
68.6
34.1
31.1
70.9
51.1
When the specialist dataset smaller, Figure 1cc shows that the errors in estimating cluster fre- quencies P(c[D*) negatively impact end task accuracy. This suggests future work to improve this estimation for tasks with small D% e.g. specific set augmentations or task grouping.
5.4 KASK-TRANSEER AND MULTITASKING
We perform cross-task evaluation, i.e. targeting task and evaluating on task B. we also pre- train a multitask models with CRISP averaged weights from multiple tasks. Our results for the 1.3B models are in Table [ we also report cross-task evaluation results for different model sizes in Ap- pendix D Cross-task evaluations show that, perhaps unsurprisingly, the best results on task A are obtained when pretraining for task Transfer differs across tasks: CRISP targeting MMLJU gives better results than base for all tasks, which is not the case for CRISP targeting ARC or RWDB-R. The multi-task result which mixes the histograms with the same weight (1/3 for ARC, MMLU and RWDB-E R) gives the best result averaged multitask accuracy. Surprisingly, RWDB-R R, this setting slightly outperforms targeting RWDB-R R itself.
Base rodpi2
5.5 TASK- DEEEDEENN CONTINUED PRETRAINING
CRISP from scratch
CRISP from (512, 800, 928, 992]
We have seen the benefit of pretraining model per task with CRISP in Figure 4 For tasks where pretraining cost concern, shorter pretraining runs still provide benefits, see Figure TO] Pretraining a multi-task model is also a cost- effective option, see Table This section evaluates third cost-effectivv option when targeting multiple tasks: continued pretraining. In this case, pretraining is divided into a generic pretraining phase and task-dependent continued pretraining phase using CRISP. The compute cost of the generic pretrain- ing can be shared across multiple tasks. Our results in Fig- ure 12 show that even 10% of CRISP continued pretraining (i.e. generic pretraining for 928 steps out of 1,024) gives an accuracy (32.9%) close to a full CRISP run (33.4%). We also remark that the impact continued pretraining stronger than fine tuning generic model MMLU (31.0% accuracy), see Figure 4
33
Ascurccy
31
29
27
500 SScpps
1000
Figure 12: Continued Pretraining on MMLU
6 CONCLUSIONS
A small specialist LM is interesting since it can outperform larger generalist LM its targeted domain while having a lower inference cost. We explore pretraining specialist LMs when little specialization data is available, common setting that prevents pretraining of dedicated LMs. We evaluate different methods that modify the distribution of a generic training set guided by little spe- cialist data. Our experiments highlight the benefit of clustered importance sampling: i.e. resampling the generic set such that its cluster histogram matches the specialist data. Our findings show that pretraining with this method provides strong models both for LM and question answering tasks. We also explore ways to lower the training cost of specialist models by showing their benefit on shorter training runs, continued pretraining and multitask settings. Our work shows that a simple, scalable importance sampling method can provide effective specialist LMs, even from little specialization data. Since clustered importance sampling is modality-agnostic, we foresee extensions of this work to other modalities, including vision and audio.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Angelos Katharopoulos, Matteo Pagliardini and Anastasiia Filippova for their advice throughout this project. We thank the anonymous reviewers for their suggestions and comments.
REFERENCES
Roee Aharoni and Yoav Goldberg. Unsupervised domain clusters in pretrained language models. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.). Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7747-7763, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/20200 ac-.o20..cllmmm maa.elmmmi..922.. URL
692
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, Colin Raffel, Shiyu Chang, Tatsunori Hashimoto, and William Yang Wang. A survey on data selection for language models, 2024. URL BtLLtttps:/ arxiv. org/abs/////ooo www.v.ooorg/ab///222222122 240..168277
J An, Ying. and Y Zhu. Why resampling outperforms reweighting for correcting sampling bias with stochastic gradients. In International Conference on Learning Representations, 2021.
David Arthur and Sergei Vassilvitskii. k-means++:
The advantages of careful seeding. Technical
report, Stanford, 2006.
Amittai Axelrod, Xiaodong He, and Jianfeng Gao. Domain adaptation via pseudo in-domain data lection. In Regina Barzilay and Mark Johnson (eds.), Proceedings of the 201 Conference on Em- pirical Methods Natural Language Processing, pp. 355-3622 362, Edinburgh, Scotland. UK., July 2011. Association for Computational Linguistics. URL https: /tEps::/aalanthhooogy org/
D11- 1133
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Al- bert Q. Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics, 2024. URL
10631
Elliot Bolton, Abhinav Venigalla, Michihiro Yasunaga, David Hall, Betty Xiong, Tony Lee, Rox- ana Daneshjou, Jonathan Frankle, Percy Liang Michael Carbin, and Christopher D. Man- ning. Biomedlm: A 2.7b parameter language model trained on biomedical text, 2024. URL
httpss:/aarxt www.////////ooooooooooo/oooeoo// 2403. 18421
Zalán Borsos, Mojmír Mutný, Marco Tagliasacchi, and Andreas Krause. Data summarization via bilevel optimization. Journal of Machine Learning Research, 25(73):1-53, 2024. URL http:
77jmlr www.ooo/pppprr/////// 113522 1132.html
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle. M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL htttes://pproceeeddngg. neurips cc/paper locterr_fllee files/paaprr/222// f Flee/15sccbbboo39911118888111212
Paperaaap..r. pdf
Jiuhai Chen and Jonas Mueller. Automated data curation for robust language model fine-tuning,
2024. URL BuR//ttpss//axii..oo//a///////
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 05557, 2018

---

Published as a conference paper at ICLR 2025
Paul Cook and Marco Lui. langid.py for better language modelling. In Paul Cook and Scott Nowson (eds.), Proceedings of the Australasian Language Technology Association Workshop 2012, pp. 107-112, Dunedin, New Zealand, December 2012. URL https:
org/
Mathicu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. Advances in Neural Information Processing Systems, 35:26698-266 10, 2022.
Scott Deerwester, Susan T Dumais, George W Furnas, Thomas K Landauer, and Richard Harshman. Indexing by latent semantic analysis. Journal of the American society for information science, 41 (6):391-407, 1990.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pp. 4171-41866 Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: (0.18653/v1/N19-11 1423.. URL https::/accantthee
N1 N19-1423
Jesse Dodge, Maarten Sap, Ana Marasović, William Agnew, Gabriel Ilharco, Dirk Groeneveld, Margaret Mitchell, and Matt Gardner. Documenting large webtext corpora: A case study on the colossal clean crawled corpus. In Scott Wen-tau Yih (eds.), Proceedings the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 1286-1305. Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguisties. doi:
Moens, Xuanjing Huang, Lucia Specia, and
10t/5/3/1/222.emm//pp https://aclaatthl..ee wwwp::/hwww..o.o.../22222//2 wwt.:221111111111..mmmmmmmmm11 ma maain... 98
URL
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, ..., Zach Rait, Zachary Viti,,, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd of models, 2024. URL bRthttts: /atttpa://aaiiiiia http:://aavvvv....ia org abs 21783]
Susan T Dumais. Latent semantic analysis. Annual Review of Information Science and Technology (ARIST), 38:189-222 2008999200000 2004.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation, 2024. URL hattpss:/ar.ii..rrg/aa////////
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.00021 2021.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le ho Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa. Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. htttihttt:///eeee 12608602
David Grangier, Pierre Ablin, and Awni Hamnun. Adaptive training distributions with scalable online bilevel optimization, 2023. URL / 20.11.///11.... 11973
David Grangier, Angelos Katharopoulos, Pierre Ablin, and Awni Hannun. Projected language mod- els: A large model pre-segmented into smaller ones. In ICML Workshop on Foundation Models in the Wild, 2024a.
David Grangier, Angelos Katharopoulos, Pierre Ablin, and Awni Hannun. Specialized language models with cheap inference from limited domain data, 2024b. /Rtt/ttt:///xxii org/ abs/ /2tp.2402..2 01093
12

---

Published as a conference paper at ICLR 2025
Suriya Gunasekar, Yi Zhang. Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang. Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need, 2023. URL www.:tttpss/.ttm https///rxvv.rro
abs /20/2 11..16644
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith. Don't stop pretraining: Adapt language models domains and tasks. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 8342 8360, Online, July 2020. Association for Computational Linguistics. doi:
URL
main. 740
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations (ICLR), 2021.
Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Data- models: Predicting predictions from training data, arXiv preprint (.XXi:2202.00622 2022.
Dan Iter and David Grangier. On the complementarity of data selection and fine tuning for domain adaptation, 2021. /ttphttts///aaxivvcoo
07591
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William Sayed. Mixtral of experts, 2024. URL htttps://axxiv.rrtt abs / 2401.040000 04088
Marcin Junezys-Dommunt. Dual conditional cross-entropy filtering of noisy parallel corpora. In Ondrej Bojar, Rajen Chatterjee, Christian Federmann, Mark Fishel, Yvette Graham, Barry Had- dow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Christof Monz, Matteo Negri, Aurélie Névéol, Mariana Neves, Matt Post, Lucia Specia, Marco Turchi, and Karin Verspoor (eds.), Proceedings of the Third Conference on Machine Translation: Shared Task Papers, pp- 888- 895, Belgium, Brussels, October 2018. Association for Computational Linguistics. doi: 10665///1/11864777777 UtL http://acaantthoogo... htt.://///...aaa////oaooooooooo. orgp/llloooooo//1111 6478
Jean Kaddour. The minipile challenge for data-efficient language models, 2023. URL https:
T Kerner. Domain-specinf pretraining of language models: comparative study in the medical field. arXiv preprint arXiv:2407.14077 2024.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Xiang Kong, Tom Gunter, and Ruoming Pang. Large language model-guided document selection, 2024. URL https:t
org. abs /00.b///240..2 04638
Yanis Labrak, Adrien Bazoge, Emmanuel Morin, Pierre-Antoine Gourraud, Mickael Rouvier, and Richard Dufour. Biomistral: A collection of open-souree pretrained large language models for medical domains, 2024. URL http::httt/rrrtt
org.
10373
Nathan Lambert, Valentina Pyatkin, Jacob Morrison, LJ Miranda, Bill Yuchen Lin, Khyathi Chandu. Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, Noah A. Smith, and Hannanch Hajishirzi. Rewardbench: Evaluating reward models for language modeling, 2024. URL
https://arxx
http::://rrdv.oor//aa////////// wwtp:vv.orggab////.//o
13787
13

---

Published as a conference paper at ICLR 2025
Edmund Lee. What happens when a publisher becomes megapublisher?
New
York Times, 2021. URL
venguin-randort house- simon- wtw.:mmossccuutteepu-i...oo..a.a
Katherine Lee, Daphne Ippolito. Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison- Burch, and Nicholas Carlini. Deduplicating training data makes language models better. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th An- nual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp- 8424-8445, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. URL https://aclanneeet:: org 2022 mall-longg
577
Patrick Lewis, Myle Ott, Jingfei Du, and Veselin Stoyanov. Pretrained language models for biomed- ical and clinical tasks: Understanding and extending the state-of-the-art. In Anna Rumshisky, Kirk Roberts, Steven Bethard, and Tristan Naumann (eds.), Proceedings of the 3rd Clini- cal Natural Language Processing Workshop, pp. 146-157, Online, November 2020. Associ- ation for Computational Linguistics. doi: 10.18653/v1/20000 12000000m000iiaaall-111117. URL https: www.aatttooooo...rg///0 http://00..iiiiiaann............
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Sooe Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with lan- guage models, 2022. URL httthhtps:////c //axivv org/ abs/ 2206. 14858
Jeffrey Li, Alex Fang, Georgios Smymis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Rein- hard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Al- balak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel II- harco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Se- woong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang. Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kol- lar. Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: In search of the next generation of training sets for language models, 2024. URL https: http://aaxvv../.. org/ / 2245..1999
Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang. Towards general text embeddings with multi-stage contrastive learning, 2023. URL https:::/// //atxiv.
org/abs/ 2308 03281
Miaofeng Liu, Yan Song. Hongbin Zou, and Tong Zhang. Reinforced training data selection for do- main adaptation. In Anna Korhonen, David Traum, and Lluís Mårquez (eds.). Proceedings of the 57th Anmual Meeting the Association for Computational Linguistics, pp. 1957-1968, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/11/11--11883 URL https: // ac www.cannnhhooggg.
Pt9/ 1189
Pratyush Maini, Skyler Seto, Richard Bai, David Grangier, Yizhe Zhang, and Navdeep Jaitly. Rephrasing the web: A recipe for compute and data-efficient language modeling. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Anmual Meet- ing of the Association for Computational Linguistics (Volume Long Papers), pp. 14044- 14072, Bangkok, Thailand, August 2024. Association for Computational Linguistics. URL.
2oo.:/ooooo00222..11-ooooo 757
Shervin Minaee, Nal Kalchbrenner, Erik Cambria, Narjes Nikzad, Meysam Chenaghlu, and Jianfeng Gao. Deep learning based text classification: comprehensive review. ACM computing surveys (CSUR), 54(3): 2021.
Robert C. Moore and William Lewis. Intelligent selection of language model training data. In Jan Hajič, Sandra Carberry, Stephen Clark, and Joakim Nivre (eds.), Proceedings of the ACL 2010 Conference Short Papers, pp. 220-224 Uppsala, Sweden, July 2010. Association for Computa- tional
URL http//////////aaaaaaoooooooooo/p
14

---

Published as a conference paper at ICLR 2025
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Noua- mane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. arXiv preprint arXiv:2305.16264 2023a.
Niklas Muennighoff, Nouamane Tazi, Loic Magne, and Nils Reimers. MTEB: Massive text embedding benchmark. In Andreas Vlachos and Isabelle Augenstein (eds.), Proceedings of the 17th Conference of the European Chapter of the Association for Computationall Linguis- tics, pp. 2014-2037, Dubrovnik, Croatia, May 2023b. Association for Computational Linguis- tics. doi: 10. http:///22222mmmmme. 148. URL. https:/ /acl aeceelannhhlgggo wtwphellogy.orr/2222.... org/2023. eacl ma in. 148/
Jupinder Parmar,
Sanjev Satheesh, Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Reuse, don 't retrain: A recipe for continued pretraining of language models, 2024. URL hetps:
Guilherme Penedo, Quentin Malartie, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon Ilm: Outperforming curated corpora with web data, and web data only, 2023.
URL https: //axxii//raxv..../////
arxiv. org/ abs/ /200s/2305 2306. 01116
Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. In Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processsing Systems, volume 33, pp. 19920-19930.
Curran Associates, Inc., 2020. URL ht htttps//ssocceednnss.eeer..pp.pp cc/pape. e688 /ttp:///...ec..113333333aaa.3...
wwwees/paprr/22220
Paper. pdf
Aniruddh Raghu, Jonathan Lorraine, Simon Kornblith, Matthew McDermott, and David K Duve- naud. Meta-learning to improve pre-training. Advances in Neural Information Processing Sys- tems, 34:22333-212-224,, 2021.
Nils Reimers and Iryna Gurevych. Settence-BERTE Sentence embeddings using Siamese BERT- networks. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th ternational Joint Conference Natural Language Processing (EMNLP-IJCNLPP, pp. 3982- 3992, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v 1/11--1410/
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Ev- timov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Seialom, and Gabriel Synnaeve. Code Ilama: Open foundation models for code, 2024.
URL /
o abs /208. 12950
Sebastian Ruder and Barbara Plank. Learning to select data for transfer learning with Bayesian optimization. In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.), Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing, pp. 372-382, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10.18653/w1/D17-1003
Chris Seiffert, Taghi M Khoshgoftaar, Jason Van Hulse, and Amri Napolitano. Resampling or reweighting: A comparison of boosting implementations. In 2008 20th IEEE international con- ference on tooly with artificial intelligence, volume pp. 445-451. IEEE, 2008
Rico Sennrich, Barry Haddow, and Alexandra Birch. Edinburgh neural machine translation sys- tems for WMT 16. In Ondíej Bojar, Christian Buck, Rajen Chatterjee, Christian Federmann, Liane Guillou, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Aurélie Névéol, Mariana Neves, Pavel Pecina, Martin Popel, Philipp Koehn, Christof Monz, Matteo Negri, Matt Post, Lu- cia Specia, Karin Verspoor, Jörg Tiedemann, and Marco Turchi (eds.), Proceedings of the First Conference on Machine Translation: Volume 2, Shared Task Papers, pp. 371 171-376, Berlin, Ger- many, August 2016a. Association for Computational Linguistics. doi: 10. 11863//1//16/ 1/W31-2333. URL httpt::///// //acl antho ogy. (ttg///6-23222 23233
15

---

Published as a conference paper at ICLR 2025
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1715- 1725, Berlin, Germany, August 2016b. Association for Computational I Linguistics. doi: 10.18653/ 11/P16-1162 htt..t.ps:: antho oropoooooo//11111122222
Hongjin Su, Jungo Kasai, Chen Henry Wu, Weijia Shi, Tianlu Wang, Jiayi Xin, Rui Zhang, Mari Ostendorf, Luke Zettlemoyer, Noah A. Smith, and Tao Yu. Selective annotation makes lan- guage models better few-shot learners. In International Conference on Learning Representations (ICLR), 2023.
Together AI Team. Redpajama--dta-v2: An open dataset with 30 trillion tokens for train- ing large language models, October 2023. URL https: wwwp://wwww..omm hwwp:/.ttotteeer.a/aaa.
Sedpajama-dee data-v2
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Huy V. Vo, Vasil Khalidov, Timothée Darcet, Théo Moutakanni, Nikita Smetanin, Mare Szafraniec, Hugo Touvron, Camille Couprie, Maxime Oquab, Armand Joulin, Hervé Jégou, Patrick Labatut, and Piotr Bojanowski. Automatic data curation for supervsee learning: A clustering-basee approach, 2024.
org/ abs/ 24 405.. 15613
Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang. Daxin Jiang, Rangan Ma- jumder, and Furu Wei. Text embeddings by weakly-superriised contrastive pre-training. arXiv preprint arXiv2222..05533, 2022. URL
221.. 03533
Wei Wang, Taro Watanabe, Macduff Hughes, Tetsuji Nakagawa, and Ciprian Chelba. Denoising neural machine translation training with trusted data and online data selection. In Ondfej Bojar, Rajen Chatterjee, Christian Federmann, Mark Fishel, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Christof Monz, Matteo Negri, Aurélie Névéol, Mariana Neves, Matt Post, Lucia Specia, Marco Turchi, and Karin Verspoor (eds.). Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 133-143, Brussels, Belgium, October 2018. Association for Computational Linguistics. doi: 0011188111111-6661 ac http:///aaaathhooyyy.. org. A18-6314
URL
Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models, 2020. URL https:
05874
G. Wenzek, M. A. Lachaux, A. Conneau, V. Chaudhary, F. Guzmán, A. Joulin, and E. Grave. Ccnet: Extracting high quality monolingual datasets from web crawl data. In Proceedings of The 12th Language Resources and Evaluation Conference, pp.
2020.
Shijie Wu, Ozan Irsoy, Steven Lu, Vadim Dabravolski, Mark Dredze, Sebastian Gehrmann, Prab- hanjan Kambadur, David Rosenberg, and Gideon Mann. Bloomberggpt: large language model for finance, 2023. URL
org/ abs /203. 17564
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS: Selecting influential data for targeted instruction tuning. In International Conference on Machine Learning (ICML), 2024.
Qianqian Xie, Weiguang Han, Xiao Zhang, Yanzhao Lai, Min Peng. Alejandro Lopez-Lira, and Jimin Huang. Pixiu: large language model, instruction data and evaluation benchmark for finance, 2023a. httphhtts////tttt
org/ abs / 2306.0060 05443
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang Quoe V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimiz- ing data mixtures speeds up language model pretraining. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural In- formation Processing Systems, volume 36, pp. 19798-69818 Curran Associates, Inc.,
16

---

Published as a conference paper at ICLR 2025
2023b. URL http:////ccceee......e..eeccs/e.
File/ fl12bbbecsss2222cc22aaaa--a--eer
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language
models via importance resampling. CoRR, abs/2302.03169, 2023c. doi: 10.48500/RXIVV23002. 03169, URL httpht:/s//oiii
10./10.... (088850/aaiia 2302. 03169
Zhipeng Xu, Zhenghao Liu, Yukun Yan, Zhiyuan Liu, Ge Yu, and Chenyan Xiong. Cleaner pre- training corpus curation with neural web scraping, 2024. URL https://axxi..rr///aa
2402. 14652
Fuzhao Xue, Yao Fu, Wangchunshu Zhou, Zangwei Zheng, and Yang You. To repeat or not to repeat: Insights from scaling Ilm under token-crisis.
A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 59304-59322. Curran Associates, Inc.. 2023. URL http://ss/ppooeeednnn...rrrrr.rp ptapppppprrr.lllessppp//////333 E22.0:04221110022223333300220020
pdf
Jinsung Yoon, Sercan Arik, and Tomas Pfister. Data valuation using reinforcement learning, In International Conference on Machine Learning, pp. 10842-10851. PMLR, 2020.
Zichun Yu, Spandan Das, and Chenyan Xiong. Mates: Model-aware data selection for efficient pre- 06046
training with data influence models, 2024. URL /.LL:tttpss//axxx..ooo//////////
Shaokun Zhang, Xiaobo Xia, Zhaoqing Wang, Ling-Hao Chen, Jiale Liu, Qingyun Wu, and Tongliang Liu. IDEAL: Influence-driven selective annotations empower in-context learners in large language models. In International Conference on Learning Representations (ICLR), 2024.
17

---

Published as a conference paper at ICLR 2025
APPENDIX
A SCALABLE CLUSTERING
We cluster the generic dataset (Redpj2) with hierarchical clustering. We build a clustering tree. Each node in the tree is associated with cluster centroid. The examples traverse the tree from top to bottom, selecting the node corresponding to the closest centroids among the current node's children.
The training of the tree proceed from root to leaves. Iteratively, a new level is built by applying k-means a subset of the examples belonging to each node. We built a tree of depth up always splitting nodes in 64 clusters. For k-means, we normalize the Fuclidean norm of the vectors prior to clustering. We train the model via Expectation Maximization using k-means++ initialization (Arthur & Vassilvitskii 2006). At each step, sample 6,400 new examples. With 20 steps, visit 128k examples. To ensure a cluster distribution close to uniform, we monitor the cluster sizes at each assignment steps. If cluster is larger than our balancing limit (0.022 1.5 1 1/64), split evenly at random its assignments with the smallest cluster, as suggested by Jegou et al. (2010). The clustering
can be found in Table
B CLUSTERING & EMBEDDING COST
The computational cost of k-means clustering is negligible compared to the cost of com- puting text embeddings. Most of the experiments in this work are performed with SBERT MiniLM-L6-v2 Reimers & Gurevych (2019). Embedding the 34.6T tokens of Redpj2 amounts to 5.4k GPU hours on a reference NVidia H100 GPU. Other models can provide better embeddings higher cost. We examined the results of clustering accuracy. MTEB evaluation Muennighofi et al. (023b)), versus embedding cost (for Redpj2 in GPU hours on H100) in Table
Table 2: Embedding Cost versus Clustering Accuracy.
Clustering method
Cost (GPUh)
Accuracy (%)
all-Mi LMM.1666262
Reimers & Gurevych 2019
5.4k
41.94
e5 e5-1arge-v2
Wang ct al. 2022)
91 4k
44.26
e5-base-v2
Wang ct 2022
27.4k
44.10
all-mpnet-base-v2
Reimers & Gurevych
2019
28.6k
43.69
gte-base-vl.5
2023
87.4k
47.90
gte-small
Li et al.
2023
13.3k
44.89
In this cost-benefit table, gte-small stands out. We clustered the Redpj2 dataset with embedding from this model and report MCQ results with CRIPS this clustering. We compare these results with LSI and SBERT from Figure in the main text. The results in Table 3 show that these embeddings are beneficial, especially with small number of clusters.
C DATASET STATISTICS
Our generic pretraining set is Redpj2 (Together Al Team 2023 We use the head+middle English version of the dataset, i.e. web-documents with high density of English text. Our specialization datasets for language modeling are much smaller, see Table Compared the LM tasks, the multiple choice question tasks have even smaller specialization training set, i.e. between 200k and 2m tokens, see T5 For the LM data, we rely on the train split provided by Pile Gao al. (202T). For the MCQ data, split each evaluation set into equal sized train and test set uniformly at random. This provides a representative specialist train set D D This also avoids cross-cov cross-contaminn coos-cootamination between tasks, e.g. the official training set of MMLU contains ARC which would prevent the task transfer experiments in Section 5.4.
D ARCHITECTURES & HYPERPARAMETERS
Our architecture configurations are borrowed from Brown et al. (2020) and described in Table We report the data selection hyperparameters in Table and the clustering hyper-parameters in Table
18

---

Published as a conference paper at ICLR 2025
Table 3: MCQ accuracy with CRISP for different embedding methods.
Clusters
Emb.
Arc-E
Are-C
MMLU
64
LSI
65.0
31.8
31.1
SBERT
65.2
33.0
31.2
GTE
67.8
33.8
31.5
4096
LSI
67.6
35.4
31.6
SBERT
69.7
36.5
32.6
GTE
69.9
37.0
32.8
262k
LSI
66.5
36.4
31.3
SBERT
71.2
38.6
33.4
GTE
69.3
37.6
33.4
16m
LSI
53.8
29.6
29.0
SBERT
62.3
33.8
30.8
GTE
N/A
N/A
31.7
Table 4: LM Datasets.
Redpj2
PubMed
StackExchange
Wikipedia
Dataset role
generalist
specialist
specialist
specialist
Train
Num. tokens
34.6T
26.7B
10.3B
4.68B
Num. documents
24.0B
2.94m
15.4m
5.79m
Test
Num. tokens
359m
52.4m
20. 1m
14.1m
Num. documents
248k
5.82k
29.9k
17.4k
Table 5: MCQ Datasets.
ARC-E
ARC-C
MMLU
RWDB-R
Train
Num. tokens
143k
79.6k
2.05m
426k
Num. questions
1.18k
578
6.95k
736
Avg. tokens per choice
30.3
34.5
73.5
289
Test
Num. tokens
144k
87.7k
2.09m
408k
Num. questions
1.19k
593
7.09k
695
Avg. tokens choice
30.2
37.0
73.6
293
Num. choices per question
4
4
Table 6: Model Hyperparameters
Num. parameters
350m
1.3m
6.7B
Architecture
Embedding dim.
1,024
2,048
4,096
Latent dim.
4,096
8,192
16,384
Num. heads
16
16
32
Depth
24
24
32
Context limit
1,024
1,024
1,024
Optimization
Batch size
96k
115k
1.04m
Learning rate
le-4
le-4
3e-4
Grad clipping
5.0
5.0
0.1
Steps
400k
1m
340k
Num. train tokens
40B
120B
350B
MCQ EVALUATION
For multiplechoocee questions, we use the LM eval hamness (Goo et al. 2024). For each task, the eval- uated model estimates the (log) probability of each answer given the context c, that log P(ale).
19

---

Published as a conference paper at ICLR 2025
Table 7: Data-Selection Hyperparameters
Method
Parameter
Range
Classifier
Regularization strength
None, 1000, 100, 10, 0.1, 0.01, 0.001) (0.5, 0.6, 0.7, 0.75, 0.8, 0.9, 0.95. 0.975, 0.98, 0.9875, 0.99, 0.995, 0.9975}
Threshold quantiles
DoGE
Num. clusters
64
Proxy model size
Transformer base, 110m parameters 32k batch size, le-4 learning rate, 100k steps
Proxy model optimization
Bregman coefficient
5e-4
Transferred weights
{run average, last step average)
Importance
Num. clusters
(64, 4096, 262k, 16.7m}
Table 8: Hierarchical Clustering Hyperparameters
Parameter
Range
Tree depth
Tree arity
64
Balancing limit
0.022
Number of samples per step
6,400
Number of steps
20
SBERT model
MiniLM-L6-v2
SBERT emb. dim.
384
LSI dim.
256
The question contains the task prompt concatenated with the current question, while the answer contains the answer text. With this strategy, the model has no access to alternative answer choices proposed in the prompt. Table reports without example questions, that is, a zero-shot evaluation (Brown et al.] 2020). Accuracy calcu- lated by verifying whether the highest score is assigned to the correct answer. The scores correspond to log probabilities for ARC-E and RWDB-R, while ARC-C, MMLU uses normalized scores, i.e. log probabilities divided by the number of characters in the answer.
prompt. For all evaluations,
use the above prompt
Table Task prompts (non-bold) for the multiple-choice-uusstoo tasks.
A12 Reasoning Challenge (ARC) Easy and Challenge Question:
<question>>>> in
Answer: <answer>
Massive Multitask Language Understanding (MMLU)
The following are multiple choice questions (with answers) about <topic>..) ) Question: <question>>>
Answer: <answer>
Rewardbench Reasoning (RWB-R)
Follow the instructions below. |n
Instructions: <question>>> |n
Answer: <answer>
F SUPPLEMENTARY RESULTS FOR LANGUAGE MODELING TASKS
We measure the fraction of examples where the pre-trained model is better than the base model. We measure this rate both on the held-out data from DE (measured the 360m tokens from Redpj2 valid) and on held-out data from D (measured on the full Pile validation set). The results in Ta- ble show that the model trained with importance sampling improves perplexity on most specialist
20

---

Published as a conference paper at ICLR 2025
Table 10: Fraction of examples with lower perplexity with importance sampling than with base. Compared to base, CRISP models specialize: they performs better most specialist examples and worse on most generic examples.
Generalist 2 (Redpj2)
Specialist 2
(Pile subset)
PubMed
6, %
97.3%
StackExchange
2.9%
92.6%
Wikipedia
12.4%
86.7%
documents (right column). Its training the importance sampled distribution utilize model capacity mostly on data close to the domain of interest, this relieves the model from fitting well most of the generic data, and hence most generic documents have higher perplexity with CRISP (left column).
For completeness, we also report the perplexity numbers Figure
Table II
Table 11: Perplexity on language modeling tasks after fine-tuning. perplexity numbers from Figure 3
These tables reports
(a) PubMed
Specific tokens
14m
100m
500m
2.5B
7.5B
26.7B
Only Specific
25.73
10.09
6.64
5.08
4.47
4.20
Base redpj2
6.34
5.98
5.38
4.79
4.47
4.20
DoGE
5.41
5.19
5.02
4.72
4.47
4.20
Classifier
5.18
5.08
4.90
4.67
4.45
4.20
CRISP
5.11
5.00
4.83
4.63
4.44
4.20
(b) StackExchange
Specific tokens
15m
133m
1.2B
10.3B
Only Specific
23.93
9.60
5.79
4.35
Base redpj2
8.41
6.66
5.35
4.35
DoGE
7.04
6.21
5.30
4.35
Classifier
7.15
6.34
5.39
4.35
CRISP
6.67
6.12
5.38
4.35
(c) Wikipedia
Specific tokens
14m
93m
668m
4.7B
Only Specific
57.13
18.22
9.97
6.76
Base redpj2
8.53
7.99
7.41
6.43
DoGE
8.13
7.71
7.26
6.39
Classifier
7.83
7.53
7.23
6.50
CRISP
7.66
7.37
7.09
6.40
G SUPPLEMENTARY RESULTS FOR MULTIPLE CHOICE QUESTIONS
Table 12 reports the MCQ results before and after fine-tuning, i.e the accuracy numbers from Fig- ure 4 Fine-tuning on the small MCQ train sets optimizing log-likelihood does not always benefit end-task accuracy.
H COMPARING THE RESULTS OF DOGE AND IMPORTANCE SAMPLING
We observe in Table 13 that the pretraining results of DoGE and importance sampling 64 clusters are close. Both methods pretrain models by sampling the clustered generalist data according to the cluster weights. If both methods would infer the same cluster weights, their pretraining runs would
21

---

Published as a conference paper at ICLR 2025
Table 12: MCQ Accurary. Fine tuning results dashed when not improved from pretraining.
This table reports the accuracy numbers from Figure
ARC-E
ARC-C
MMLU
RWDB-R
Pretr.
+ft
Pretr.
+ft
Pretr.
+ft
Pretr.
+ft
Base redpjv2
58.4
27.4
28.8
30.0
31.0
62.1
63.5
DoGE
66.2
33.3
34.9
31.0
31.5
67.4
67.9
Classifier
68.2
36.7
32.4
69.4
CRISP
71.2
38.6
33.4
33.5
70.0
Table 13: DoGE & CRISP on 64 Clusters
LM PPL PubMed
MCQ Ace (%) f
ARC-E
ARC-C
MMLU
DoGE CRISP
7.44
66.2
33.3
31.0
7.28
65.2
33.0
31.2
be identical. We therefore ask if the similar results are due to similar cluster weights. Figure 13 compares the cluster weights for both methods. The top clusters for both methods are similar, but their histograms are not identical. This shows that similar pretraining results can be obtained with different weights.
0.6
DeGE
0.5
Claster weight
0.4
Claster weight
0.15
0.3
0.2
0.2
0.05
0.1
0.0
0.00
Clustens
Clusters
Clusters
(a) PubMed
(b) ARC
(c) MMLU
Figure 13: DoGE CRISP weights with 64 clusters. We report the top-16 clusters sorted by mean weight across methods.
I COMPARING CRISP AND CROSS-ENTROPE DIFFERENCE (CED)
Contrasting the scores of two LMs (generalist and specialist) is popular method for data selee- tion (Moore & Lewis 2010 Axelrod et al. [201]] Wang et al. 2018 Jamezys-Dowmuut 2018). We considered this method based Wang al. (018); generalist LM on the specialist training set. We rely on a 350m parameter model for the selection. One should note that this method particularly expensive since requires scoring the entire Redpj2 dataset twice with an LM, which is more expensive than embedding and clustering the dataset. Our results show that CED improves the classifier method but CRISP is significantly better, Table 14
obtain the specialist LM by fine-tuning
TASK TRANSFER FOR 350M, 1.3B AND 7B MODELS
Table 15 complements the task-transfer results from Table in Section 5.4 with the results across different model sizes. The importance sampling models trained with MMLU histograms outperform the base models on all tasks for all model sizes.
22

---

Published as a conference paper at ICLR 2025
Table 14: Comparison with Cross-Entropy Difference for MCQ Accuracy (%), 1.3B model.
Arc-E
Arc-C
MMLU
Base
58.4
27.4
30.0
CED
58.9
30.5
31.1
Doge
66.2
33.3
31.0
Classifier
68.2
36.7
32.4
CRISP
71.2
38.6
33.4
Table 15: Accuracy (%) for Task Transfer on 350m, IB and 7B models.
Model
Evaluation Tasks
ARC-E
ARC-C
MMLU
RWDB-R
Multi
350m
Base
49.5
24.5
27.0
57.6
40.5
CRISP ARC
65.3
31.5
27.4
58.4
44.7
CRISP MMLU
55.6
26.3
29.8
61.3
44.0
1B
Base
58.4
27.5
30.1
62.2
45.1
CRISP ARC
71.3
38.6
28.9
60.9
48.2
CRISP MMLU
63.4
28.7
33.4
65.2
48.2
7B Base
69.9
35.9
34.4
64.9
50.7
CRISP ARC
74.5
42.2
32.6
62.4
51.1
CRISP MMLU
70.0
37.6
38.0
67.5
53.1
23


Reviews:
[{"review_id": "4Eg7Yjaixn", "reviewer": "Reviewer_RQp5", "summary": "This work proposes a novel method to pre-train a specialist language model, where the amount of specialist data is scarce, and the amount of generalist data is rich. The high-level idea in this method is to sample from the generalist data and select the ones that are more similar to the specialist data. The authors target it with a simple cluster-then-importance-sampling, showing improvement in perplexity and accuracy on language model and multiple choice QA tasks.", "questions": "It would be helpful if the authors could provide a more detailed association with previous work.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The proposed method is simple while effective, showing competitive performance compared to vanilla pre-training and other baselines.\n\n* The ablation study in Section 5 is comprehensive and well-explored.", "weaknesses": "* Previous work [1] has proposed a similar method of adjusting generalist data according to clusters informed by specialist data. The major contribution of the method is that this work shows that that method will work when scaled up.\n\n[1] Specialized Language Models with Cheap Inference from Limited Domain Data", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "98XcUNasG0", "reviewer": "Reviewer_oNof", "summary": "This paper addresses the challenge of building specialist language models (LMs) when domain-specific data is limited. The authors introduce Clustered Importance Sampling (CRISP), which adapts a generalist dataset for specialist pretraining by adjusting its distribution to align with limited domain-specific data. CRISP’s effectiveness is demonstrated across various domains, yielding improvements in language modeling perplexity and multiple-choice question accuracy. The approach is scalable, applicable to both small and large models, reduces the need for costly domain-specific data, and consistently outperforms other adaptation techniques, such as classifiers and gradient alignment, across a range of tasks and domains.", "questions": "I have the following additional questions and comments:\n\n- As shown between SBERT and LSI experiments, the choice of clustering model impacts performance. Would using a more powerful embedding model like E5-large instead of the current MiniLM-L6-v2 improve performance, considering factors such as computational cost (as using E5-large would be more costly) and potential performance gains?\n-  How long does this process typically take to build the clustering tree for different configurations? I think Figure 10 only accounts for the training cost and it would be nice to compare it with the overall training cost.\n\nI have also identified some potential typos:\n- Currently, the equation on page 4 after equation (2) gives $\\mathcal{L}(\\mathcal{D}^{s}; \\theta) = \\sum\\limits_{c} \\mathcal{L}(c; \\theta) \\frac{P(c | \\mathcal{D}^{s})}{P(c | \\mathcal{D}^{g})} P(c | \\mathcal{D}^{g}) = \\underset{c \\sim (c | \\mathcal{D}^{s})}{\\mathbb{E}} [w(c) \\mathcal{L}(c; \\theta)]$. Shouldn't this be $\\mathcal{L}(\\mathcal{D}^{s}; \\theta) = \\sum\\limits_{c} \\mathcal{L}(c; \\theta) \\frac{P(c | \\mathcal{D}^{s})}{P(c | \\mathcal{D}^{g})} P(c | \\mathcal{D}^{g}) = \\underset{c \\sim (c | \\mathcal{D}^{g})}{\\mathbb{E}} [w(c) \\mathcal{L}(c; \\theta)]$, that is $c \\sim (c | \\mathcal{D}^{g})$ instead of $c \\sim (c | \\mathcal{D}^{s})$? In addition, if correction is needed, would this change the interpretation of the Importance Sampling as the equation changes?\n- In Algorithm 1, I think line 8 should be $x_i \\sim \\text{Uniform}(D^{g} \\cap K(c_i))$ as it samples a generalist example over a samples cluster id in the large set as explained in previous paragraphs. In addition, in line 10, the set should start with $x_1$ instead of $x_0$ to be consistent with the indexing.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. The problem is clearly defined and well-motivated, addressing the need for effective specialist language models when domain-specific data is limited.\n\nS2. CRISP demonstrates superior performance compared to both untrained and pre-trained models on the Redpj2 baseline, as well as classifier-based and DoGE approaches. It achieves lower perplexity on language modeling tasks and higher accuracy on multiple-choice question (MCQ) tasks across four different domains.\n\nS3. The paper provides a comprehensive analysis, exploring the effects of cluster numbers, model sizes, varying amounts of training data, task transferability, and multitask training. CRISP has also been shown to be more efficient in terms of training cost and scales well with larger models.", "weaknesses": "W1. While the paper provides a comprehensive analysis of different cluster numbers, it is limited in terms of guidance on determining the optimal number of clusters beyond empirical trial and error, even though this is a crucial parameter for CRISP’s performance. Results indicate that the number of clusters influences the sampling distribution and domain adaptation effectiveness, concerning token repetition or overfitting in certain configurations.\n\nW2. Although CRISP improves model performance in non-pretrained scenarios, it offers limited additional benefits for MCQ tasks after fine-tuning compared to task-specific pretraining.\n\nW3. When the specialist dataset $\\mathcal{D}^s$ is very small, CRISP appears to be less effective in enhancing end-task accuracy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qQjTaKYSmM", "reviewer": "Reviewer_Z9XR", "summary": "This work focuses on training the specialist LM to avoid the limitation of scarce  specialist data.  To this end, this work explore three data selection method during pre-training in which the clustered importance sampling (named CRISP in this work) performs best. The CRISP aims to modifies the distribution of a generic training set guided by little specialist data. To be specific, the method resamples the generic set such that its cluster histogram matches the specialist data. Therefore, the model can always select the samples with the same or similar distribution with the targeted domains. Experimental results demonstrate that pre-training with this method provides better specialist models both for LM and question answering tasks. Besides, the additional analysis experiment benefits readers more, e.g., investigating the impact of model size, training set size, the number of clusters, clustering methods, continued pre-training and multi-task settings. All of the results show the advantage of the proposed method.", "questions": "See the weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The authors extend the importance sampling (Grangier et al., 2024b) through adjusting the frequency of generalist clusters guided by specialist data. Besides, it scale the method to millions of clusters, show that it works with larger models, and extend it beyond language modeling tasks (i.e., adding the downstream tasks).\n2. Extensive experiments show the importance of sampling data and the effectiveness of the cluster-based importance sampling method.\n3. The cluster-based importance sampling method may benefit the researchers that devoted into the end-to-side small LLMs.", "weaknesses": "1. The cluster-based importance sampling method is not proposed by this work and the authors only show it works in larger models, and extend it to MQA tasks. Therefore, the contribution of technology is limited. \n2. From my perspective, the CRISP method frequently selects samples (tokens) that share the same distribution as the target domain. What if you were to iterate over generalist samples/tokens—without distinguishing whether they belong to the same domain—the same number of times as you would when training specialist LMs?\n3. When a new domain or downstream task coming, it has to be trained a new specialist LMs from scratch. It maybe time-cost and laborious. \n4. If you apply the CRISP method based on a pre-trained generalist LMs, I wonder whether it can gain more benefit. If not, does it mean that you are still not using enough data in current pre-training setting?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "90FVFPpZdZ", "reviewer": "Reviewer_TBtV", "summary": "In this work, the authors investigate the data selection problem during the pre-training and continued pre-training stages of language models. They propose CRISP, a method that clusters both the generalist and specialist datasets, and then selects data from the generalist dataset based on the cluster distribution of the specialist dataset. The authors conduct extensive experiments on six benchmarks to demonstrate the effectiveness of CRISP.", "questions": "1. If I understand the proposed approach CRISP correctly, should line 8 of Algorithm 1 be $x_i \\sim \\text{Uniform}(D^{g} \\cap K(c_i))$?\n2. Typically, we use accuracy to measure performance on MMLU, while the authors compute perplexity on MMLU. How are the inputs formatted? Are the question and the corresponding correct option concatenated?\n3. As shown in Algorithm 1 and Figure 1, CRISP resamples the generalist data based on the specialist data histogram. Is the specialist dataset not used in the training process?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The authors conduct comprehensive experiments, utilizing significant computational resources, to validate the effectiveness of the proposed CRISP method.\n2. Although the computational cost of CRISP is not explicitly discussed, based on its methodology, CRISP appears to be computationally efficient.", "weaknesses": "1. The paper is somewhat hard to follow. For instance, CRISP requires a K-means clustering model to select task-specific data, but the authors do not mention on which dataset the K-means model is trained. Furthermore, many technical details are placed in the appendix, making the paper hard to follow and not self-contained.\n2. CRISP essentially selects data from the generalist dataset based on the distribution of the specialist dataset. There are several baseline approaches that should be considered, such as Cross-Entropy Difference [1], as mentioned in the related work section.\n3. Although the author conduct extensive experiments with substantial amount of compute, the scientific contributions of this work appear to be marginal and the findings are not particularly surprising. Like previous works, CRISP aims to align the selected data more closely with the targeted domain. It is unclear what unique strengths CRISP offers compared to prior methods. For example, given the vast amount of training data, it would be beneficial if the authors can demonstrate CRISP is significantly more efficient than existing approaches. Furthermore, the data quality in the selection process might be another factor to consider.\n\n\n### Additional Comments (do not affect my rating)\n\n1. Recent works [2, 3] have demonstrated that results on RewardBench have a negative correlation with real downstream performance. I suggest the authors consider using other benchmarks for evaluation. Although these findings are very recent and were reported after the paper submission deadline—and it's not the authors' fault—I believe it would be beneficial if they could adjust their evaluation settings accordingly.\n2. The authors utilize significant computational resources for the experiments in this work. I appreciate these efforts. It would be helpful if the authors could present a comparison regarding the compute cost to better understand the pros and cons of the approaches discussed.\n\nReference:\n\n[1] Moore, Robert C., and William Lewis. \"Intelligent selection of language model training data.\" In Proceedings of the ACL 2010 conference short papers, pp. 220-224. 2010.\n\n[2] Zhou, Enyu, Guodong Zheng, Binghai Wang, Zhiheng Xi, Shihan Dou, Rong Bao, Wei Shen et al. \"RMB: Comprehensively Benchmarking Reward Models in LLM Alignment.\" arXiv preprint arXiv:2410.09893 (2024).\n\n[3] Frick, Evan, Tianle Li, Connor Chen, Wei-Lin Chiang, Anastasios N. Angelopoulos, Jiantao Jiao, Banghua Zhu, Joseph E. Gonzalez, and Ion Stoica. \"How to Evaluate Reward Models for RLHF.\" arXiv preprint arXiv:2410.14872 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)