Task-Adaptive Pretrained Language Models via Clustered-Importance Sampling

## Abstract

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.

## 1 Introduction

Generalist language models (LMs) can address a wide variety of tasks, but this generality comes at a cost (Brown et al., ). It necessitates a large training set representative of all prospective tasks, as well as a large model to fit such a comprehensive dataset. Specialist models forgo this generality and fit a model for a limited domain or task. In their narrow specialty, such models can achieve better accuracy at a given model size (Kerner, ).

Pretraining a 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 on the targeted application and its potential economic benefit. Condition (ii) is more limiting since modern LMs are commonly pre-trained on datasets larger than 100B tokens(^1^11100B tokens $\simeq$ 1m books $\simeq$ 60x the annual publication of the top English language publisher (Lee, ).), 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 a large generalist dataset and a small specialist dataset, we propose to modify the distribution over the generalist dataset guided by the scarce specialist dataset. Training a model on the modified distribution gives a specialist model with better accuracy than a generalist model of the same size.

We study this setting across different specialization tasks including domain-specific language modeling (medical, encyclopedic domains) and end-tasks (scholar exams in science and humanities, reasoning questions). We compare different strategies to manipulate the pretraining distribution. We evaluate strategies based on text classifiers, gradient-alignment and importance sampling (IS). Although IS is rarely used for LM data selection, we build upon on a simple IS recipe based on clustering (Grangier et al., ) 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 1. We show the empirical benefit of this method varying model sizes (350m to 7B parameters), 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.

## 2 Related Work

Generalist vs Specialist LMs Generalist LMs address tasks which they have not been trained for explicitly (Brown et al., ), or provide a good initialization for fine-tuning a dedicated model (Devlin et al., ). Nowadays generalists compete with dedicated models on many tasks (Jiang et al., ; Dubey et al., ). Success however comes at a price: a generalist must be much larger than a specialist for the same accuracy.
For instance, on English-to-German translation, the 175-B parameter generalist GPT-3 (Brown et al., ) is less accurate than a 136m-parameter specialist (Sennrich et al., ).
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. Corpora with billions of tokens enable pretraining or *continued pretraining*, a generalist pretraining phase followed by a specialist one (Gururangan et al., ; Parmar et al., )). Domains with specialist models include medicine and biology (Lewis et al., ; Labrak et al., ; Bolton et al., ), computer programming and mathematics (Lewkowycz et al., ; Rozière et al., ; Azerbayev et al., ) and finance (Wu et al., ; Xie et al., ). When specialist data is available in limited amount, there are methods to train specialist models on generalist data instead.

Task-Adaptive Data-Selection The selection methods over-sample generalist data that aids model generalization in the target domain. For masked LMs, Gururangan et al. () observe that continued pretraining improves the performance on end-tasks when using data with high vocabulary overlap with the targeted task. For machine translation (MT), Aharoni & Goldberg () show that a task-adapted pretraining dataset can be selected from a generalist dataset using the nearest neighbors of a small specialist set. Their nearest neighbor classifier relies on BERT sentence distance (Devlin et al., ). Still for MT, other works have used other types of classifiers. In particular, contrasting the scores of two LMs (generalist and specialist) is popular (Moore & Lewis, ; Axelrod et al., ; Wang et al., ; Junczys-Dowmunt, ). Other classifiers include logistic regression or fine-tuned BERT (Iter & Grangier, ).
Outside classification, Xie et al. () proposed to use importance sampling for continued pretraining. They improve classification tasks by selecting pretraining data with a similar distribution to the targeted domain in terms of hashed-ngrams. Importance sampling is also used in (Grangier et al., ) and we build upon that work which adjusts the frequency of generalist clusters informed by specialist data: we 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., ; Xia et al., ; Grangier et al., ). The pretraining distribution is selected such that the reweighted gradients from the generalist dataset mimics the expected gradient from the small specialist dataset. Gradient-alignment for data selection has also been used for other purposes such as data summarization (Borsos et al., ), pretraining acceleration (Xie et al., ; Fan et al., ) or auxiliary task weighting (Wang et al., ; Raghu et al., ). Finally, it is also worth mentioning data selection methods based on reinforcement learning (Liu et al., ; Yoon et al., ), bayesian optimization (Ruder & Plank, ), data models (Ilyas et al., ) and influence models (Yu et al., ).

Pretraining Data Quality Outside of domain aspects, the quality of pretraining data is also an important topic (Wenzek et al., ; Dodge et al., ; Penedo et al., ; Li et al., ). Data quality includes removing data in other languages (Cook & Lui, ), text formatting (Xu et al., ), favoring long form text Gao et al. (); Gunasekar et al. (), removing duplicates (Lee et al., ). It also involves balancing
different sources of data with the goal of reaching a better generic pretraining loss Xie et al. (); Fan et al. (); Vo et al. (). Recent work also considered filtering Kong et al. (), correcting Chen & Mueller () or generating Maini et al. () pretraining data with LMs.
These data quality considerations are orthogonal to domain concerns: quality filters are applied alongside domain adaptation decisions (Albalak et al., ).

## 3 Data Selection for Task-Adaptive Pretraining

We consider three methods for task-adaptive pretraining of LMs. Classification and gradient alignment have been evaluated in different contexts before but not for end-tasks like multiple-choice question answering. Clustered-based importance sampling at scale is a contribution of this work, building upon recent work from Grangier et al. ().

### 3.1 Notations

$D^{\rm g}$ is the training dataset sampled from the generalist distribution ${\cal D}^{\rm g}$. $D^{\rm s}$ is the specialist dataset representative of the final task,
sampled from ${\cal D}^{\rm s}\neq{\cal D}^{\rm g}$. The loss of model $\theta$ on $D$ is

$$ ${\cal L}(D;\theta):=\frac{1}{|D|}\sum_{x\in D}\ell(x;\theta)=-\frac{1}{|D|} \sum_{x\in D}\frac{1}{|x|}\sum_{i}\log p(x_{i}|x^{i-1}_{1};\theta)$ $$

where $|D|$ denotes the cardinality of a set $D$ and $|x|$ denotes the length of sequence $x=x^{|x|}_{0}=(x_{1},\ldots,x_{|x|})$. The perplexity of model $\theta$ on the dataset $D$ is ${\cal P}(D;\theta):=\exp({\cal L}(D;\theta)).$

### 3.2 Classification

A binary classifier is trained to estimate the probability that a generalist pretraining document belongs to the targeted domain. The classifier $\phi$ is learned using positive examples from $D^{\rm s}$ and a subset of $D^{\rm g}$ as negative examples. $\phi$ then builds a domain-specific pretraining set

$$ $C(D^{\rm g},t):=\{x\in D^{\rm g}\text{ such that }\phi(x)>t\}.$ $$

which restrict the generic dataset $D^{\rm g}$ to the examples with an estimated probability to be in-domain above threshold $t$.
The threshold is validated as a trade-off between focusing on data close to the domain of interest while keeping $C(D_{g},t)$ large enough to train an LM of the targeted capacity.
In our case, we rely on a logistic regression classifier trained over sentence BERT (SBERT) text embeddings (Reimers & Gurevych, ), an established classification method (Minaee et al., ). The SBERT representation is also commonly used in data selection (Albalak et al., ; Xie et al., ; Zhang et al., ; Su et al., ). This representation is also used in the alternative selection strategies we consider. As an ablation, we also evaluate the impact of the choice of SBERT (Section 5.1).

### 3.3 Gradient-Alignment

Gradient-Alignment (GA) methods are common when the generic pretraining set $D^{\rm g}$ originates from $n_{\rm g}$ different data sources
$S$, i.e. $D^{\rm g}=\bigcup^{n_{\rm g}}_{i=1}D^{\rm g}_{i}$. These methods select weights for the different sources by considering two functions of $\theta$: the pretraining reweighed loss,

$$ ${\cal L}((w,D^{\rm g});\theta):=\sum^{n_{\rm g}}_{i=1}w_{i}{\cal L}(D^{\rm g}_ {i};\theta),$ $$

and the targeted loss, i.e. the loss on $D^{\rm s}$ in our case. The weights, on the simplex,
can be inferred via a bilevel formulation of the data selection problem (Dagréou et al., ):
the minimum $\theta^{\star}(w)=\arg\min_{\theta}{\cal L}((w,D^{\rm g});\theta)$ depends on $w$
and task-dependent pretraining is interested in weights $w$ that minimize ${\cal L}(D^{\rm s};\theta^{\star}(w))$ wrt $w$. This formulation results in algorithms that select weights during pretraining to align the gradients of these two functions wrt $\theta$ (Xie et al., ; Grangier et al., ). In our case, we rely on the DoGE (Fan et al., ) algorithm. Compared to classifiers, GA is harder to scale to large model size. This limitation is commonly addressed by finding the mixture weights with a small model before transferring them to a larger model.

In this work, we consider a generic setting where the pretraining dataset $D^{\rm g}$ is
not pre-segmented into few data sources. Instead, we rely on 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, ) and domain adaptation (Grangier et al., ).

### 3.4 CRISP: ClusteRed Importance Sampling for Pretraining

Figure: Figure 1: Clustered-based Importance Sampling with CRISP.
Refer to caption: x1.png

We sketch our strategy from Figure 1. Initially, we divide the space of text into clusters. We decompose the specialist loss and the generalist loss as a weighted sum of losses over clusters. Then we make an independence assumption that implies that the specialist and generalist loss per cluster are identical. The specialist loss is then computed as the generalist loss with a reweighing of each cluster.

Specifically, we want to identify a model with a low loss on the
specialist distribution ${\cal D}^{\rm s}$,

$$ ${\cal L}({\cal D}^{\rm s};\theta)=\mathop{\mathbb{E}}_{x\sim{\cal D}^{\rm s}}[ \ell(x;\theta)]=\sum_{x}\ell(x;\theta)P(x|{\cal D}^{\rm s})$ $$

We marginalize over a discrete latent variable $c$, the cluster variable, and write

$$ ${\cal L}({\cal D}^{\rm s};\theta)=\sum_{x}\sum_{c}\ell(x;\theta)P(x|c,{\cal D} ^{\rm s})P(c|{\cal D}^{\rm s})\underset{(2)}{=}\sum_{x}\sum_{c}\ell(x;\theta)P (x|c)P(c|{\cal D}^{\rm s})$ (1) $$

where the second equality ${=}_{(2)}$ makes the independence assumption
$P(x|c,{\cal D}^{\rm s})=P(x|c)$. If we make a similar assumption for the
generalist loss $P(x|c,{\cal D}^{\rm g})=P(x|c)$, we can write both losses as

$$ ${\cal L}({\cal D}^{\rm s};\theta)=\mathop{\mathbb{E}}_{c\sim(c|{\cal D}^{\rm s })}[{\cal L}(c;\theta)]\quad\text{and}\quad{\cal L}({\cal D}^{\rm g};\theta)= \mathop{\mathbb{E}}_{c\sim(c|{\cal D}^{\rm g})}[{\cal L}(c;\theta)]$ (2) $$

where we define ${\cal L}(c;\theta)=:\sum_{x}\ell(x;\theta)P(x|c)$.
We now apply importance sampling to these expectation, defining the importance weights as $w(c)=P(c|{\cal D}^{\rm s})/P(c|{\cal D}^{\rm g})$,

$$ ${\cal L}({\cal D}^{\rm s};\theta)=\sum_{c}{\cal L}(c;\theta)P(c|{\cal D}^{\rm s })=\sum_{c}{\cal L}(c;\theta)\frac{P(c|{\cal D}^{\rm s})}{P(c|{\cal D}^{\rm g} )}P(c|{\cal D}^{\rm g})=\mathop{\mathbb{E}}_{c\sim(c|{\cal D}^{\rm s})}[w(c){ \cal L}(c;\theta)].$ $$

In our experiments, we estimate the terms $w(c),{\cal L}(c;\theta)$ from
the finite training sets $D^{\rm s}\sim{\cal D}^{\rm s}$
and $D^{\rm g}\sim{\cal D}^{\rm g}$. We count the number of examples in each
cluster to estimate $P(c|{D}^{\rm s}),P(c|{D}^{\rm g})$.
The expected loss over a
cluster ${\cal L}(c;\theta)$ is estimated as the average loss over the generalist
examples in cluster $c$, ${\cal L}({D}^{\rm g}\cap K(c);\theta)$,
where $K(c)$ denotes the examples in cluster $c$.
This strategy therefore only estimates $P(c|{D}^{\rm s})$ on the small ${D}^{\rm s}$. The term ${\cal L}(c;\theta)$ is estimated over the large set as
${\cal L}({D}^{\rm g}\cap K(c);\theta)$
which has less variance than the estimator ${\cal L}({D}^{\rm s}\cap K(c);\theta)$ over the small ${D}^{\rm s}$.

We train CRISP models with stochastic optimization (Kingma & Ba, , Adam) and propose Algorithm 1. Here, we do not explicitly reweigh the loss. We instead samples cluster from their importance. This avoids frequently visiting clusters with less weights. This strategy has less variance in its gradient estimates,
which can help convergence (Seiffert et al., ; An et al., ).
This algorithm is simple and efficient when one groups the generalist
examples by cluster prior to training.

Figure: Algorithm 1 CRISP Training

## 4 Experiments & Results

We perform experiments with transformer LMs (Vaswani et al., ). Most of our experiments use models with 1.3B parameters (trained on 120B tokens) and we conduct ablations with 350m and 7B models (resp. trained on 40B, 350B tokens). Our settings for architectures and optimization are borrowed from Brown et al. (), see Appendix C.

Our generalist training set is Redpj2 (Together AI Team, ). We select this dataset as it contains only web-crawled data without additional interventions to help evaluation tasks (e.g. adding encyclopedias, books or academic articles).
Redpj2 contains over 30T tokens with our 32k byte-pair encoding tokenizer (Sennrich et al., ), see Table 2 in Appendix B.
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 branching $64$ for $4$ levels, see Appendix A. The levels therefore have 64, 4,096 ($=64^{2}$), 260k ($=64^{3}$) and 16.7m ($=64^{4}$) 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 tf-idf representations (Deerwester et al., ; Dumais, ).

For our specialist tasks, we consider 3 language modeling tasks (LM) and 3 multiple-choice-question tasks (MCQ). For LM, we use Pile subsets from different domains (Gao et al., ): medical (Pubmed Central), programming Q&A (Stackexchange), and encyclopedic (Wikipedia). For MCQ answering, we use AI2 Reasoning Challenge (Clark et al., , ARC), Massive Multitask Language Understanding (Hendrycks et al., , MMLU), and Reward Bench Reasoning (Lambert et al., , RWDB-R). ARC focuses on science questions, MMLU focuses on interdisciplinary knowledge, RWDB-R focuses on correct vs incorrect solutions to math and programming problems. To provide a representative specialist train set $D^{\rm s}\sim{\cal D}^{\rm s}$, we split the questions into a train and test split, see Table 3 in Appendix B.

Our main results are reported with unified settings. For the classifier, the classification threshold is the main parameter. A threshold accepting 2.5% of $D^{\rm g}$ worked best in 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 it over 64 clusters, i.e. learning a 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 C.

### 4.1 Language Modeling Tasks

We evaluate specialist LMs on three domains from the Pile (Gao et al., ): 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 2 in Appendix B.

As baselines, we either train only on the in-domain (specialist) data without pretraining or we fine-tune a model pre-trained on Redpj2. We refer to the Redpj2 pretraining distribution as the base distribution. 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 specialist training tokens.
In each case, the resampled pretraining set is used to train a 1.3B-parameter transformer model with the same hyperparameters as the Redpj2 baseline.

We report pretraining results in Figure 4, and the fine-tuning results in Figure 4. 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 to 100%-ft with 26.7B tokens.

The pretraining results in Figure 4 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 on data close to the targeted domain: the model capacity is not used to fit the filtered out training data. Table 8 in Appendix E shows, for instance, that CRISP outperforms base on 97.3% of PubMed but reports worse perplexity on 95.9% of Redpj2.

When we fine tune the pretrained models,
the advantage of task-dependent pretraining is preserved, in Figure 4. 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 4 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).

Figure: (a) PubMed
Refer to caption: x2.png

### 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 3 in Appendix B. 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 D). 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 4 displays a similar method ordering and CRISP is consistently the best method. As a difference with LM tasks, we observe limited benefits from fine tuning, 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 10 in Appendix F.

## 5 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 embeddings, see Section 4.
We report their performance with 64, 4096, 262k and 16.7m clusters.

The representation is important: the examples in a cluster are close in the embedding space. Our independence assumption, Equation 1, assumes that the loss in a cluster $c$ is the same regardless whether its data originates from ${D}^{\rm g}$ or ${D}^{\rm s}$, i.e.

$$ ${\cal L}({D}^{\rm g}\cap K(c);\theta)\simeq{\cal L}({D}^{\rm s}\cap K(c); \theta).$ (3) $$

In practice, it is sufficient that the embedding space reflects the similarity of the loss gradient, i.e. if the gradients of the loss over a generalist cluster ${D}^{\rm g}\cap K(c)$ is correlated with the gradient over a specialist cluster ${D}^{\rm s}\cap K(c);\theta)$, the model trained on the former improves on the later. Figure 7 shows that the SBERT representation is more adapted and yields better results than LSI for all settings.

Figure: (a) Arc-E
Refer to caption: x12.png

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 3, 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|{\cal D}^{\rm s})\simeq P(c|{D}^{\rm s})$ gets less accurate. The estimator risks overfitting, i.e. favoring clusters frequent in the training set ${D}^{\rm s}$ but not as frequent on other samples from ${\cal D}^{\rm s}$. 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 ${D}^{\rm g}$.

We explore these aspects on MMLU. We first measure the number of repeated examples when training models with CRISP pretraining for different number of clusters. Figure 7 shows the number of repetitions for each quantile of the training set. Even for 16.7m clusters, only a small minority of training examples are repeated beyond 10 times and the average number of occurrences of the training is examples 1.94, well within commonly recommended values (Muennighoff et al., ; Xue et al., ). Even if exact repetitions do not account for the poorer performance of the 16.7m cluster setting, its training set might be less diverse and the model might generalize well only in a small neighborhood of its training set.
We plot Figure 7 to evaluate if the Redpj2 examples with good perplexity concentrate
around $D^{\rm s}$, the MMLU training set. This plot shows that the benefit of CRISP over base is indeed correlated with the distance to $D^{\rm s}$. As the number of cluster increases to 16.7m, the benefit over base concentrate in an area with very few samples. For comparison, we plot the 2 middle quartiles [0.25, 0.75] where most of the MMLU test data concentrate in gray. We remark that MMLU test data mostly lies in an area where the perplexity of IS 16.7m is low.

Figure: (a) Perplexity on reweighted Redpj2
Refer to caption: x18.png

Figure 8 shows the perplexity during training for CRISP runs on MMLU. On Figure 8(a), the perplexity is computed from the reweighed loss on Redpj2. This is the loss optimized during pretraining. It shows that when the number of cluster increases the sampled training set is less diverse and corresponds to an easier learning problem ($<5$ PPL). On Figure 8(b), 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 MMLU perplexities. This quantifies the error resulting from our assumption, Equation 3. We also remark 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 7 fails to show that 260k cluster would have the best perplexity, which highlights that SBERT distance to the training data
is not the only factor explaining model performance.

### 5.2 Model Size

Figure: (a) Arc-E
Refer to caption: x20.png

This section compares CRISP and base at 3 model sizes. The benefit of task-dependent training is consistent across model sizes, see Figure 10. 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 ($\sim$30x). Of course, a smaller model is also beneficial at inference.

### 5.3 Different Amount of Training Data

Figure: (a) Repetitions for less generic data
Refer to caption: x26.png

This section varies both (i) the amount of generalist data available to sample the CRISP dataset from and (ii) the amount of specialist data for inferring the CRISP weights.

When specialist data concentrates 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 to 700B and 120B tokens (downsampling Redpj2 by resp. $\sim$ 50x and $\sim$300x). Our pretraining runs use 120B tokens, so a base run never repeats in all settings. When CRISP is applied, some tokens are repeated.
Figure 11(a) shows that,
when restricting to 120B tokens, the number of repetition becomes high (22.5 on average) and CRISP is ineffective after 256k steps.

When the specialist dataset is smaller, Figure 11(c) shows that the errors in estimating cluster frequencies $P(c|{\cal D}^{\rm s})$
negatively impact end task accuracy. This suggests future work to improve this estimation for tasks with small ${D}^{\rm s}$: e.g. specific set augmentations or task grouping.

### 5.4 Task-Transfer and Multitasking

We perform cross-task evaluation, i.e. targeting a task A and evaluating on a task B, we also pretrain a multitask models with CRISP averaged weights from multiple tasks. Our results for the 1.3B models are in Table 1, we also report cross-task evaluation results for different model sizes in Appendix H.
Cross-task evaluations show that, perhaps unsurprisingly, the best results on a task A are obtained when pretraining for task A. Transfer differs across tasks: CRISP targeting MMLU 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-R) gives the best result on averaged multitask accuracy. Surprisingly, on RWDB-R, this setting slightly outperforms targeting RWDB-R itself.

**Table 1: 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 |

### 5.5 Task-Dependent Continued Pretraining

Figure: Figure 12: Continued Pretraining on MMLU
Refer to caption: x29.png

We have seen the benefit of pretraining a model per task with CRISP in Figure 4.
For tasks where pretraining cost is a concern, shorter pretraining runs still provide benefits, see Figure 10. Pretraining a multi-task model is also a cost-effective option, see Table 1. This section evaluates a third cost-effective option when targeting multiple tasks: continued pretraining.
In this case, pretraining is divided into a generic pretraining phase and a task-dependent continued pretraining phase using CRISP. The compute cost of the generic pretraining can be shared across multiple tasks. Our results in Figure 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 of continued pretraining is stronger than fine tuning a generic model on MMLU (31.0% accuracy), see Figure 4.

## 6 Conclusions

A specialist LM is interesting since a small specialist LM can outperform a larger generalist LM on its targeted domain while having a lower inference cost. We explore pretraining specialist LMs when little specialization data is available, a common setting that prevents pretraining of dedicated LMs. We evaluate different method that modifies the distribution of a generic training set guided by little specialist 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. These benefits are studied across model size, training set size and clustering methods. We also explore ways to lower the training cost of specialist models addressing 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.

## Acknowledgements

We thank Matteo Pagliardini and Angelos Karathopoulos for their help throughout the project. We also thank Maartje ter Hoeve, Yizhe Zhang and Navdeep Jaitly for their suggestions and comments.

## Appendix

## Appendix A Scalable Clustering

We cluster the Redpj2 dataset with hierarchical clustering.
We build a clustering tree. Each node in the tree is associated with a 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 to a subset of the examples belonging to each node. We built a tree
of depth up to $4$, always splitting nodes in 64 clusters. For k-means, we normalize the Euclidean norm of the vectors prior to clustering. We train the model via Expectation Maximization using k-means++ initialization (Arthur & Vassilvitskii, ). At each step, we sample 6,400 new examples. With 20 steps, we visit 128k examples. To ensure a cluster distribution close to uniform, we monitor the cluster sizes at each assignment steps. If a cluster is larger than our balancing limit $(0.022\simeq 1.5*1/64)$, we split evenly at random its assignments with the smallest cluster, as suggested by Jegou et al. ().
The clustering hyper-parameters can be found in Table 6.

## Appendix B Dataset Statistics

**Table 2: 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 3: 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 per choice | 30.2 | 37.0 | 73.6 | 293 |
| Num. choices per question | 4 | 4 | 4 | 2 |  |

Our generic pretraining set is Redpj2 (Together AI Team, ). We use the head+middle English version of the dataset, i.e. web-documents with a high density of English text. Our specialization datasets for language modeling are much smaller, see Table 2.
Compared the LM tasks, the multiple choice question tasks have even smaller specialization training set, i.e. between 200k and 2m tokens, see Table 3. For the LM data, we rely on the train split provided by Pile Gao et al. ().
For the MCQ data, we split each evaluation set into an equal sized train and test set uniformly at random. This provides a representative specialist train set $D^{\rm s}\sim{\cal D}^{\rm s}$. This also avoids cross-contamination between tasks, e.g. the official training set of MMLU contains ARC which would prevent the task transfer experiments in Section 5.4.

## Appendix C Architectures & Hyperparameters

Our architecture configurations are borrowed from Brown et al. () and described in Table 4. We report the data selection hyperparameters in Table 5 and the clustering hyper-parameters in Table 6.

**Table 4: 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 | 1e-4 | 1e-4 | 3e-4 |
| Grad clipping | 5.0 | 5.0 | 0.1 |
| Steps | 400k | 1m | 340k |
| Num. train tokens | 40B | 120B | 350B |

**Table 5: Data-Selection Hyperparameters**
| Method | Parameter | Range |
| --- | --- | --- |
| Classifier | Regularization strength | {None, 1000, 100, 10, 1, 0.1, 0.01, 0.001} |
|  | Threshold quantiles | {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} |
| DoGE | Num. clusters | 64 |
|  | Proxy model size | Transformer base, 110m parameters |
|  | Proxy model optimization | 32k batch size, 1e-4 learning rate, 100k steps |
|  | Bregman coefficient $\mu$ | 5e-4 |
|  | Transferred weights | {run average, last 20 step average} |
| Importance S. | Num. clusters | {64, 4096, 262k, 16.7m} |

**Table 6: Hierarchical Clustering Hyperparameters**
| Parameter | Range |
| --- | --- |
| Tree depth | 4 |
| 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 |

## Appendix D MCQ Evaluation

For the multiple-choice-question tasks, we use the LM eval harness (Gao et al., ). For each task, the evaluated model estimates the (log) probablity of each answer $a$ given the context $c$, i.e. $\log P(a|c)$. The question contains the task prompt concatenated with the current question, while the answer contains the answer text. With this strategy, the model does not have access to all proposed answer choices in the prompt.
Table 7 reports our prompt. For all evaluation, we use the above prompt without example questions, i.e. a zero-shot evaluation (Brown et al., ).
Accuracy is computed by verifying if 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.

**Table 7: Task prompts (non-bold) for the multiple-choice-question tasks.**
| AI2 Reasoning Challenge (ARC) Easy and Challenge |
| --- |
| Question: <question>$\backslash$n |
| Answer: <answer> |
| Massive Multitask Language Understanding (MMLU) |
| The following are multiple choice questions (with answers) about <topic>.$\backslash$n |
| Question: <question>$\backslash$n |
| Answer: <answer> |
| Rewardbench Reasoning (RWB-R) |
| Follow the instructions below.$\backslash$n |
| Instructions: <question>$\backslash$n |
| Answer: <answer> |

## Appendix E Supplementary Results for Language Modeling Tasks

We measure the fraction of examples where the pretrained model is better than the base
model. We measure this rate both on held-out data from ${\cal D}^{\rm g}$ (measured on the 360m tokens from Redpj2 valid) and on held-out data from ${\cal D}^{\rm g}$ (measured on the full Pile validation set). The results in Table 8 shows that the model trained with importance
sampling improves perplexity on most specialist documents (right column). Its training on 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, hence most generic documents have higher perplexity with CRISP (left column).

For completeness, we also report the perplexity numbers from Figure 4 in Table E.

**Table 8: Fraction of examples with lower perplexity with importance sampling than with base. Compared to base, CRISP models specialize: they performs better on most specialist examples and worse on most generic examples.**
|  | Generalist ${\cal D}^{\rm g}$ | Specialist ${\cal D}^{\rm s}$ |
| --- | --- | --- |
|  | (Redpj2) | (Pile subset) |
| PubMed | 6.1% | 97.3% |
| StackExchange | 2.9% | 92.6% |
| Wikipedia | 12.4% | 86.7% |

**Table 9: Perplexity on language modeling tasks after fine-tuning. These tables reports the perplexity numbers from Figure 4.**
| (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 |

## Appendix F Supplementary Results for Multiple Choice Questions

Table 10 reports the MCQ results before and after fine-tuning, i.e the accuracy numbers from Figure 4. Fine-tuning on the small MCQ train sets
optimizing log-likelihood does not always benefit end-task accuracy.

**Table 10: MCQ Accurary. Fine tuning results are dashed when not improved from pretraining. This table reports the accuracy numbers from Figure 4.**
|  | 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 | – |

## Appendix G Comparing the results of DoGE and Importance Sampling

**Table 11: DoGE & CRISP on 64 Clusters**
|  | LM PPL$\downarrow$ | MCQ Acc (%) $\uparrow$ |  |  |
| --- | --- | --- | --- | --- |
|  | PubMed | ARC-E | ARC-C | MMLU |
| DoGE | 7.44 | 66.2 | 33.3 | 31.0 |
| CRISP | 7.28 | 65.2 | 33.0 | 31.2 |

We observe in Table 11 that the pretraining results of DoGE and importance sampling on 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 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.

Figure: (a) PubMed
Refer to caption: x30.png

## Appendix H Task Transfer for 350m, 1.3B and 7B models

Table 12 complements the task-transfer results from Table 1 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.

**Table 12: Accuracy (%) for Task Transfer on 350m, 1B 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 |

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: It would be helpful if the authors could provide a more detailed association with previous work.
soundness: 3
strengths: * The proposed method is simple while effective, showing competitive performance compared to vanilla pre-training and other baselines.

* The ablation study in Section 5 is comprehensive and well-explored.
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.
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.

[1] Specialized Language Models with Cheap Inference from Limited Domain Data

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: I have the following additional questions and comments:

- 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?
-  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.

I have also identified some potential typos:
- 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?
- 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.
soundness: 4
strengths: S1. The problem is clearly defined and well-motivated, addressing the need for effective specialist language models when domain-specific data is limited.

S2. 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.

S3. 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.
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.
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.

W2. 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.

W3. When the specialist dataset $\mathcal{D}^s$ is very small, CRISP appears to be less effective in enhancing end-task accuracy.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: See the weaknesses.
soundness: 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).
2. Extensive experiments show the importance of sampling data and the effectiveness of the cluster-based importance sampling method.
3. The cluster-based importance sampling method may benefit the researchers that devoted into the end-to-side small LLMs.
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.
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. 
2. 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?
3. 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. 
4. 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?

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
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))$?
2. 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?
3. 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?
soundness: 3
strengths: 1. The authors conduct comprehensive experiments, utilizing significant computational resources, to validate the effectiveness of the proposed CRISP method.
2. Although the computational cost of CRISP is not explicitly discussed, based on its methodology, CRISP appears to be computationally efficient.
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.
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.
2. 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.
3. 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.


### Additional Comments (do not affect my rating)

1. 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.
2. 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.

Reference:

[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.

[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).

[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).

Meta-review: This paper proposes a method called Clustered Importance Sampling (CRISP) for pre-training specialist LMs when domain-specific data is limited. CRISP adjusts the distribution of a generalist dataset to align with a limited specialist dataset by resampling the generalist data based on a cluster distribution. The authors demonstrate that CRISP improves performance on language modeling and multiple-choice question accuracy tasks compared to other adaptation techniques.

The authors show that the method achieves competitive performance compared to other pre-training methods on various tasks and domains. The paper provides a comprehensive analysis of the method, exploring the effects of various hyperparameters and training configurations. The method is shown to be scalable to different model sizes. 

The reviewers argue that the paper does not provide clear guidance on determining the optimal number of clusters, which is crucial for CRISP's performance. CRISP offers limited additional benefits for multiple-choice question tasks after fine-tuning compared to task-specific pre-training. When the specialist dataset is very small, CRISP appears to be less effective in enhancing end-task accuracy.

Overall, this paper presents a promising method for pre-training specialist LMs with limited data. Despite some limitations which the authors have agreed to acknowledge, I think this is a good paper and I recommend acceptance.

Decision: Accept (Poster)