Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Authors: Wanyun Xie
Francesco Tonin
Volkan Cevher
Venue: ICML 2025 poster (2025)
Keywords: Data mixture
LLMs
leverage score
pretraining
finetuning
Primary area: deep_learning->large_language_models
TL;DR: We introduce a new efficient data-mixing framework based on leverage scores boosting LM generalization that can flexibly adapt to new domains and can be applied in finetuning

Abstract: Training data mixtures greatly impact the generalization performance of large language models. Existing domain reweighting methods often rely on costly weight computations and require retraining when new data is introduced. To this end, we introduce a flexible and efficient data mixing framework, Chameleon, that employs leverage scores to quantify domain importance within a learned embedding space. We first construct a domain affinity matrix over domain embeddings. The induced leverage scores determine a mixture that upweights domains sharing common representations in embedding space. This formulation allows direct transfer to new data by computing the new domain embeddings. In experiments, we demonstrate improvements over three key scenarios: (i) our computed weights improve performance on pretraining domains with a fraction of the compute of existing methods; (ii) Chameleon can adapt to data changes without proxy retraining, boosting few-shot reasoning accuracies when transferred to new data; (iii) our method enables efficient domain reweighting in finetuning, consistently improving test perplexity on all finetuning domains over uniform mixture. Our code is available at https://github.com/LIONS-EPFL/Chameleon.

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Wanyun Xie
Francesco Tonin
Volkan Cevher'
Abstract
significantly influences the generalization capabilities and downstream performance of LLMs.
Training data mixtures greatly impact the gen- eralization performance of large language mod- els. Existing domain reweighting methods often rely on costly weight computations and require retraining when new data is introduced. To this end, we introduce a flexible and efficient data mixing framework, CHA CHAMELEON, that employs leverage scores quantify domain importance within learned embedding space. We first con- struct domain affinity matrix over domain em- heddings. The induced leverage scores determine a mixture that upweights domains sharing com- mon representations in embedding space. This formulation allows direct transfer to new data by computing the new domain embeddings. ex- periments, we demonstrate improvements over three key scenarios: (i) our computed weights im- prove performance on pretraining domains with a fraction of the compute of existing methods; (ii) CHAMELEON can adapt to data changes with- out proxy retraining, boosting few-shot reason- ing accuracies when transferred to new data; (iii) our method enables efficient domain reweight- ing in finetuning, consistently improving test per- plexity on all finetuning domains over uniform mixture. Our code is available at https:// github. S.O.o//...NSSEEEEE EPFIS EPF..//CPP.// loo.
Domain reweighting, which involves adjusting the relative contributions of different domains in the training dataset, has emerged as critical aspect of LLM training (Gao et al., 2020; Du et al., 2022). However, obtaining optimal domain weights is a challenging problem due to factors such as data quality, diversity, inter-domain overlap, and task-specific complexities (Shen et al., 2023; Longpre et al., 2024).
Early domain reweighting methods relied on manual selec- tion, often favoring high-quality sources like Wikipedia and academic texts (Brown al., 2020; Gao al., 2020). While intuitive, these approaches are neither optimal nor scalable.
Recent work explores computational strategies for optimiz- ing domain mixtures, such as DoReMi (Xie et al., 2023) and DoGE (Fan et al., 2024b), which use small proxy model to derive domain weights for training a larger base model. Though effective, these methods are computationally expen- sive and have limited practical applicability.
We contend that ideal data-mixing method should (i) improve universal generalization, the fundamental goal of domain reweighting; (ii) adapt to domain modifications data naturally evolves between preparation and LLM train- ing. making frequent recalibration impractical; (iit) handle different training stages such as pertaining and fine-tuning. Most existing methods are limited to pretraining scenarios and do not consider either domain changes or the fine-tuning stage where domain specificity often plays a larger role.
In this work, we introduce CHAMELEON, novel and effi- cient framework for data mixing that addresses these chal- lenges. Our method computes domain weights directly from learned embeddings using kernel ridge leverage scores (KRLS), which quantify the importance of each domain based on its contribution the overall embedding space.
1. Introduction
Pretraining large language models (LLMs) relies heavily on vast and diverse datasets, encompassing sources such academic papers, books, and code repositories (Brown et al., 2020; Dubey et al., 2024). The composition of these datasets
Unlike existing approaches, CHAMBLEON reduces the data- mixing compute and eliminates the need for frequent re- training of proxy models when domains change. Instead, it constructs domain affinity matrix to capture relationships between domains and computes leverage scores that guide domain reweighting. Our data-centric approach enables
'Laboratory for Information and Inference Systems, École Polytechnique Fédérale de Lausanne (EPFL), Switzerland. Corre- spondence to: Wanyun Xie <wanyun.xie <vanyun.xoeeepttch... @cp.t.chs.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
seamless adaptation new data and flexibility across both pretraining and fine-tuning stages.
Table I: Data-mixing methods capabilities comparison.
DoReMi
DoGE
CHAMEEEON
Specifically, our contributions are summarized as follows:
Generalization Downstream New data Finetuning
We propose CHAMELEON, an efficient data-mixing framework that leverages KRLS to quantify do- main representativeness from embedded data. In- verse KRLS-based domain weights emphasize general knowledge for pertaining. We empirically demonstrate that CHAMELEON matches DoReMi and DoGE in pre- training performance at fraction of their cost.
X
X
/
In this setting, our work develops an offline method that achieves uniformly strong performance across domains with- out relying on downstream task knowledge. Other offline methods, such as Data Mixing Laws (Ye et al., 2024) and Meg that-in contrast to ours-requires access to the down- stream tasks (Liu et al., 2024), use multiple proxy models to search for optimal data mixtures, revealing that domain weights transfer across model sizes.
As a data-centric method, CHAMELEON can flexibly adapt changes in domain composition without re- training proxy models, enhancing its practicality in real scenarios. It outperforms baselines at 1 1 of the retraining cost, even as domains double.
We extend domain reweighting to fine-tuning, where KRLS-based weights emphasize domain-specific uniqueness. Empirical results show consistent perplex- ity improvements across all domains on both natural language and code datasets.
Additionally, studies on data and model scaling laws provide further insights into domain weighting strategies (Kang et al., 2024; Ye et al., 2024). However, they critically rely on the scaling strategy and do not have an easy way of adapting to domain expansion.
From a practical standpoint, CHAMELEON significantly lowers the computational burden associated with domain reweighting, making advanced LLM training pipelines more accessible to researchers with limited resources.
Kernel Ridge Leverage Scores (KRLS). The notion of statistical leverage score (Gareth et al., 2013) is used in best- rank approximation to define importance score for the rows in a matrix by their influence on the optimal low-rank approximation (Mahoney & Drineas, 2009), with approxi- mation error guarantees for sampling (Li et al., 2013).
Indeed, our KRLS-based scores are computationally inex- pensive, hyperparameter-robee robust, and converge quickly. This efficiency is particularly advantageous when incorporating new data, where our method can be applied directly without re-running the entire proxy pipeline. By bridging the gap between pretraining and fine-tuning, our method provides unified and agile framework for efficient as well as effective data mixing across all stages of LM training.
Ridge leverage scores are also proposed with additional regularization term (Cohen al., 2015; 2017; Rudi et al., 2018). Leverage scores are extended the kernel setting in (Bach, 2013), namely kernel ridge leverage scores.
KRLS sampling is extensively used to accelerate kernel methods (Alaoui & Mahoney, 2015; Musco & Musco, 2017; Rudi et al., 2017; 2018). Inverse KRLS have been related to Christoffel functions (Pauwels et al., 2018), which in machine learning are used for, e-g., landmark sampling (Fanuel et al., 2022), density estimation (Pauwels et al., 2018), and outlier detection (Lasserre & Pauwels, 2019; Beckermann et al., 2021; Ducharlet et al., 2024).
2. Related Work
Domain Reweighting. Domain reweighting improves LLM pretraining by balaneing data contributions from dif- ferent sources. In this setting, online adaptation strategies require frequent recalibration and monitoring (Albalak et al., 2023; Jiang et al.. 2024; Fan et al., 2024a).
Two closely related approaches are DoReMi (Xie et al., 2023) and DoGE (Fan et al., 2024b). DoReMi trains both reference and a proxy model using Group DRO (Sagawa et al., 2020) mitigate excess domain loss, while DoGE tracks domain-specific gradients during proxy training.
3. Data Mixing via CHAMELEON
Setup. We consider a dataset D (DI...., Dp. consist- ing of k distinct domains, each represented by its metadata (e.g., source, topie). Our objective is to determine domain weight vector E AE. where As is the probability simplex, enhancing the performance of LMs. Following prior work (Xie et al., 2023; Fan et al., 2024b), we adopt a two-stage strategy: (1) training a small proxy model to infer domain weights and (2) training large base model using the com-
Both methods are inefficient: DoReMi depends on the refer- ence model's ' quality and requires training two models, while DoGE incurs high gradient tracking costs. Their domain weights fluctuate significantly during training (Figure 3).
2

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Learn Domain Embeddiing
Compute Domain Affinity
Compute KRLS Scores
Train base LMs
Pretraining
10
12
13
-21
.10 -10
Domin Weights: of" suftmax(S
15
12
24
19
90
18
10
Transfer to New Data
15
19
22
22
-10
-2
101
Domain Weights of softmax (5,
-21
20
-22
15
31
-10
10
18
11
21
Finetuning
Demain Weights o! solmax (5x)
-17
.00
Ana
Rock
CC
Co
Githah
Stack
Wiki
Figure 1: Pipeline of domain reweighting via CHAMELEON. The given data is first embedded through the proxy model, previously trained on a corpus D> with uniform weights. Domain embeddings are then determined by averaging the embeddings for each domain. The domain affinity matrix Sip is computed as the pairwise inner products between domain embeddings. Finally, (KRLS) is applied to S2p to obtaín score S2 indicating the degree of inter-domain independency. resampling non- con-uniform distribution is obtained by softmax normalizing the scores. Finally, the target base language model is trained with the obtained data mixture, where the inverse KRLS S is used during pretraining of initial or new data to promote general knowledge learning and the KRLS Sx is used during finetuning emphasize task-specific knowledge.
puted weights. Note that many studies have empirically shown that domain weights transfer across different model sizes (Xie et al., 2023; Fan et al., 2024b; Liu et al., 2024).
3.1. Quantifying domain characteristics
We first introduce domain embeddings capturing domain- level characteristics, representing each domain D; as the embedding vector 2; E IR', with p being the embedding dimension. This embedding is computed by averaging the LM embeddings of data points in the domain:
We approach the data mixture problem from a data-centric perspective. Unlike prior works such as DoReMi (Xie et al., 2023) and DoGE (Fan et al., 2024b), which derive domain weights based the optimization process of proxy model, we focus instead on the intrinsic properties of the data itself.
1 D
Σ GED.
(L) hgg
(a),
31 =
To characterize domain characteristics, we extract embed- dings from hidden layers of the proxy model. These embed- dings encapsulate rich semantic and structural information about the input data in continuous, high-dimensional space. As result, the embeddings not only represent the individual domains but also capture their inter-domain relationships.
where (1) denotes the L-th layer embedding of the proxy model hag, In practice, using a sufficiently large, randomly sampled batch B; C D; provides a robust approximation of the domain embedding and is used in experiments. The proxy is trained on D with uniform domain weights, i.e., ag = 1/k, following (Xie et al., 2023; Fan et al., 2024b). We define the resulting domain embedding matrix as X [x]] e R
Figure 1 presents two-dimensional UMAP (McInnes et al., 2018) projection of mid-layer embeddings derived from the SlimPajama dataset using a proxy language model. The visualization highlights the following key characteris- ties: (i) semantic distinctiveness: similar domains cluster closely, while unique domains stand apart; (ii) centrality and coverage: broad domains like "CC" and "C4" create dominant regions, covering shared semantic space, while more specific domains like "Arxiv" are more distinct.
To quantify properties across domains, we exploit Kernel Ridge Leverage Scores (KRIS). KRLS is well-established tool in data analysis quantifying the influence or importance of data points (Alaoui & Mahoney, 2015). KRLS measure the contribution of each domain to the overall embedding space. First, we define a kernel function k(D;, .ry) = x xj, which captures the similarity between domains D, and Dj. Using this kernel, we construct the domain affinity matrix:
Our observations lead to two central questions:
1. How can we precisely quantify domain characteristics?
Ωp = [n(2;, [n(,,,,)) M)) = XXT.
2. How can such properties inform domain reweighting?
The domain affinity matrix Ωp captures pairwise rela- tionships between domains, with higher values indicating a higher degree of semantic similarity. Note that we
We tackle these questions in the sequel.
3

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
employ the linear kernel as the LM itself already introduces significant non-linearity. An example domain affinity matrix is visualized in Figure 1 @ While Up captures inter-domain relationships, it does not directly provide a measure of individual domain importance for data mixing.
Algorithm I Domain Weighting via CHAMELEON.
1: Input: Training data from k domains D
[Di...., ......,,D), regularization parameter X, embedding layer index L.
2: if Pretraining then
To address this, we propose to employ KRLS defined on Spp to quantify the influence of each domain within the overall embedding space. We compute the scores as defined below.
3: Train proxy hg (a) with uniform weights a; = k.
4: end if
Extract domain embeddings: 2,
1. [D;]] EeeD, n8p (a)
for each domain Dj. Construct the feature matrix: X =
Definition 3.1 (Domain KRLS). For a given regularization parameter À > 0, the KRLS for domain D; is defined as:
[= [x. 2
Compute the domain affinity matrix: UD XXT.
(D;) = (Sp + kJ1) IJa
(KRLS)
8: Compute KRLS (D,) for each domain D, using Stp.
9: if Pretraining then
where / is the k x & identity matrix.
exp(55'(D1)
10: Domain weights of
E)=rexp(55 (D,)))
The KRLS (D.) quantifies the importance of domain D; in the embedding space. Specifically, these scores corre- spond the diagonal entries of the hat matrix (2tp + kXZ)-1 of Kernel Ridge Regression (KRR) (Hastie et al., 2005). They are proportional to the weights assigned to each domain in the dual KRR estimator, and they depend only on the inputs 2; and constant À independently of specific target values (Calandriello et al., 2016; Chen Yang, 2021), )). Inputs with higher KRLS indicate higher contribution to the KRR estimator, i.e., they are more unique in the kernel representation. More discussions on the KRLS and the role of regularization are provided in Appendix A.1 and A.2, respectively.
11: else if Fine- Tuning then
12: Domain weights of PT
(D.))
13: end if
14: Output: Domain weights a or a
space, where higher values indicate denser regions. Detailed remarks are provided in Appendix A.2.
During pretraining, assigning higher sam- pling probability to domains with high S upweights high- density data regions, which are most influential on base LMs' performance (Mallen et al., 2023; Feng et al., 2024). To achieve this, we determine domain weights using the inverse of KRLS:
Pretraining.
In data mixing, a high KRLS value for domain D; indicates that its embedding x; cannot be -ppr-oppro as a combination of embeddings from other domains, implying that D; is relatively distinct or unique in its characteristics. Conversely, low KRLS value suggests that D; is highly well-represented, as it can be readily reconstructed from other domain embeddings, representing broader or more widely shared characteristics across domains.
exp (D.) Σ exp (5,((D))
PT o
where we convert the scores S into the probability distri- bution a PT by appliying softmax normalization E exp())
exp(1
3.2. Incorporating KRLS into LM training
Importantly, the domain weights obtained by CHAMELFON focus on the intrinsic properties of the data and our method does not affect the proxy model's training process. This al- lows to compute importance weights aND of new domains directly without requiring retraining of the proxy model by applying it to the new data. The proxy is used to obtain the new domain embeddings, from which the new (KRLS) score is calculated. In contrast, existing data mixture methods couple domain reweighting with the proxy model's optimization, necessitating costly retraining whenever domains are added. This dependency not only increases computational overhead but also contradicts the goal of im- proving efficiency in large-scale dynamic training pipelines.
An essential question is thus: Which domains should be pri- oritized: the general ones with higher degree of dependency with others or independent ones with more unique charac- teristics? ? Prior work suggests that data mixing strategies should adapt different training phases (Ma et al., 2023; Feng et al., 2024). Therefore, we address this by consider- ing pretraining and fine-tuning separately, as their objectives differ fundamentally (Parthasarathy et al., 2024).
Remark 3.2 (Theoretical motivation). Before introducing our weights, we provide a theoretical motivation for our approach. It is well established that the inverse KRLS ST1 is proportional to the Christoffel function (Pauwels et al., 2018), which is a common measure of density of the data in embedding space. Christoffel functions precisely charac- terize the local density of the data distribution in the feature
Finetuning. Finetuning aims to specialize on a novel spe- cific task (Yang et al., 2024), requiring the model to learn differential features not fully captured during pretraining,
4

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
so we instead prioritize the domains with high KRLS Sx:
DoRcMi
.25
DoGis
FT exp (Sx(D.)
7
Domain Weight
CHAMELEON
ブラ exp (Sx)
.20
.15
Note that the key difference between finetuning and pretrain- ing in the data mixture problem is that, during finetuning, we do not need to train a separate proxy model. Instead, we directly use the pre-trained model for finetuning. This allows us to extract the embeddings from the pre-trained
.10
.05
.00
model for the relevant domains, which are then used to compute the domain weights.
Arxiv
Book
CC
C4
Github
Wikipedia Exchanxe
Algorithm and complexity analysis. The overall domain reweighting process is summarized in Algorithm 1. Our phase-specific strategy ensures that it well adapts to the differing demands of pretraining and fine-tuning. Obtaining embeddings I; requires a single forward pass for each sample D; through the proxy; inference is fast as the proxy is a small model. Given the typically small number of domains k the KRLS computation in Definition 3.1 is cheap in (kk)). We do not add any overhead in proxy training. Our approach is therefore efficient and contrasts with prior methods requiring domain-specnt iterative optimization (Xie et al., 2023; Fan et al., 2024b).
Figure Domain weights on SlimPajama. We compare weights computed by data-mixing methods on SlimPajama.
Table 4: GPU hours for obtaining domain weight.
DoReMi
DoGE
CHAMELEON
684M base model
7.4h
6.3h
0.8h
56h
Following DoGE (Fan al., 2024b), we use a small 82M decoder-only transformer (Vaswani et al., 2017) the aux- iliary models for CHAMELEON, DoReMi, and DoGE. Aux- iliary models for both DoGE and DoReMi are trained for 10k iterations, and the proxy model of CHAMELEON is only trained for 2k steps and we use 4k samples for embedding computation per domain. Detailed training setup demon- strated in Appendix B.2. Domain weight obtained through different methods is reported in Figure 2. Through train- ing small auxiliary models, domain weights are obtained to train larger base models with the size of 684M. Note that we employ the simple, linear kernel is ((2,, xz)) = x zj. which does not need further kernel hyperparameter tuning.
4. Experiments
We show CHAMELEON improves the model's performance through data mixture with less computational costs (Sec- tion 4.1). In addition, CHAMELEON is scalable and can be applied to larger datasets without the need to retrain a proxy model (Section 4.2). Moreover, it can easily applied to fine-tuning tasks (Section 4.3).
4.1. Universal Generalization
Considering universal generalization, the main goal is to im- prove the general performance of the model across domains in the training set and also in various downstream tasks. For the performance on the in-distribution data, we measure the perplexity across all domains. For downstream tasks, we follow RegMix (Liu et al., 2024) selecting 13 tasks that cover various realistic scenarios.
Perplexity and computational cost. Table 2 shows per- domain perplexity the held-out test dataset for 684M base models. CHAMELEON outperforms Uniform and DoReMi, achieving similar performance to DoGE but with significantly lower computational cost. Specifically, DoReMi trains two auxiliary models, while DoGE computes k = 7 gradients per iteration with roughly 1.7x wall-clock time per iteration, both of which are costly. In contrast, CHAMELBON achieves competitive results with just 2k steps of training, and its inference cost is minimal compared to the training. In our setting, training an 82M proxy model requires 10 1017-1018 FLOPs, and extracting embeddings re- quires only 15 FLOPs (< 1% of proxy training). As a re- sult, the total FLOPs of CHAMELEON, including both train- ing and inference for embeddings, is only 10% of DoReMi
Training setup. We experiment on the SlimPajama-6677 dataset (Soboleva et al., 2023) consisting of 7 data domains. We choose Uniform with uniform domain weights, DoReMi, and DoGE as our baselines, which are downstream task agnostic offline methods same CHAMELEON. We include RegMix as an additional reference, as it instead leverages prior knowledge of downstream tasks. Specifically, RegMix optimizes domain weights using the validation loss of the domain most correlated with downstream performance, identified as "CC" in their work (Liu et al., 2024).
5

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Table 2:
Universal generalization perplexity. Per-domain test perplexity for the universal generalization compared with the uniform baseline, DoReMi, DoGE, and RegMix with 684M parameter models. Note that, unlike other methods, RegMix knows the target downstream tasks for optimization. We report the compute (in FLOPs) required to arrive at the data mixture. CHAMELEON boosts generalization a fraction of the computational cost.
Domain
Uniform
DoReMi
DoGE
CHAMELEON
RegMix
Arxiv
8.16
9.16
9.07
8.31
11.35
Book
42.55
46.48
40.30
39.23
41.52
CC
45.26
40.62
38.99
40.11
37.32
C4
49.00
43.92
40.65
42.59
43.85
Github
3.99
4.10
4.09
4.20
4.99
Stackexchange
7.99
8.35
7.39
7.94
10.63
Wikipedia
12.42
10.78
15.74
13.90
20.88
Average PPL (L)
24.20
23.34
22.32
22.31
24.36
# Domains Over Uniform
3/7
4/7
4/7
3/7
FLOPs
1.34 10
6.68 10
1.36
10 (1x)
1.20 10
(10x)
(5x)
(9x)
Table Universal generalization reasoning. Accuracy of downstream tasks in the same settings as Table 2.
Table 5: Universal generalization with 1.2B model rea- soning, Large-scale pretraining experiments.
Benchmark
Unif.
DoReMi
DoGE
Cu..
RegMix
Benchmark
Unif.
DoReMi
DoGE
CiI.
RegMix
ARC-E
36.8
37.6
38.0
37.8
39.1
ARC-E
39.4
41.2
41.9
42.4
43.0
COPA
55.7
59.3
62.3
61.9
63.0
CОРA
64.0
66.0
63.0
61.0
66.0
HellaSwag
26.5
27.0
27.2
27.1
27.0
HellaSwag
27.5
27.7
28.2
28.4
27.6
Lambada
13.5
13.6
14.7
15.1
16.5
Lambada
17.9
17.3
18.7
21.6
20.7
LogiQA
21.7
21.9
22.4
22.6
21.4
LogiQA
22.0
24.0
22.0
21.2
20.7
MultRRC
57.2
55.7
57.3
57.2
56.6
MultiRC
57.2
57.2
57.2
57.2
56.9
OpenBook
14.1
13.3
14.6
14.4
14.7
OpenBookQA
15.0
13.6
13.8
16.4
17.4
PIQA
59.2
59.5
60.0
60.5
57.6
PIQA
61.5
61.9
61.8
63.8
58.7
QQP
36.8
36.8
36.8
39.2
37.1
QQP
36.8
36.8
36.9
36.9
36.8
RACE
26.1
25.3
26.4
26.5
27.3
RACE
26.0
26.7
27.8
29.1
28.4
SciQ
61.8
62.5
64.9
64.3
64.1
SciQ
69.7
68.3
69.0
72.6
72.0
Social IQA
35.0
35.5
35.7
35.7
35.6
SocialIQA
36.2
36.5
35.9
37.2
36.1
WinoGrande
50.5
51.3
52.0
52.1
50.9
WinoGrande
52.8
49.6
48.9
51.5
50.0
Average (1)
37.9
38.4
39.4
39.6
39.3
Average (t)
40.5
40.5
40.4
41.5
41.1
and 20% of DoGE, In addition, we demonstrate GPU hours in Table and discuss its details in Appendix B.4, high- lighting the efficiency and lower computational overhead of CHAMELEON.
2017), SeiQ (Welbl al., 2017), Social IQA (Sap al., 2019), WinoGrande (Sakaguchi et al., 2021). The reported accuracy in Table is the average from 0-shot to 5-shot evaluations following (Liu et al., 2024), scored using the Im-eval-harnes evaluation framework (Gao et al., 2024).
Evaluation on downstream tasks. We apply our method on realistic downstream tasks. We follow RegMix (Liu et al., 2024) selecting 13 tasks that cover various realistic tasks: ARC-E (Clark et al., 2018), COPA (Sarlin et al., 2020), HellaSwag (Zellers et al., 2019), Lambada (Paperno et al., 2016). LogiQA (Liu et al., 2020), MultiRC (Khashabi et al., 2018), OpenBookQA (Mihaylov et al., 2018), PiQA (Bisk al., 2020), QQP (Wang, 2018), RACE (Lai al.,
These benchmarks cover diverse range of tasks, en- abling a comprehensive evaluation of the real-world im- nor-
pact of CHAMELEON. For each benchmark, malized accuracy as the evaluation metric if provided by Im-eval-harnee else we use regular accuracy. Notably, CHAMELEON also shows competitive performance across downstream tasks even compared with RegMix, a task- aware method.
6

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Table 6: Universal generalization with 1.2B model perplexity. Large-scale experiments in the same settings of Table 5.
Domain
Uniform
DoReMi
DoGE
CHIAMBLEON
RegMix
Arxiv
6.30
7.09
7.07
6.33
10.61
Book
28.25
32.66
27.83
24.63
27.55
CC
31.19
29.96
28.11
26.95
24.70
C4
34.74
33.05
31.06
29.58
31.94
Github
2.91
3.03
3.07
2.94
4.08
Stackexchange
6.01
6,44
5.80
5.76
9.54
Wikipedia
8.65
7.93
10.88
9.03
20.08
Average PPL (L)
16.86
17.17
16.26
15.03
18.36
# Domains Over Uniform
3/7
4/7
4/7
3/7
Scale to 1.2B model. Prior works (Xie et al., 2023; Fan et al., 2024b; Liu et al., 2024) have shown that domain weights transfer effectively across different model scales. To validate this for CHAMELEON, to 1.2B models. The detailed results for perplexity and downstream task performance are presented in Table 6 and Table 5, respectively. Our findings confirm such transferabil- ity: the weights derived from an 82M proxy model proved effective when applied to both the 684M and 1.2B models. Importantly, we observed that CHAMELEON yielded even more significant improvements on larger models.
only incurs significant computational costs but also delays the adaptation to changes in domain composition. How can we develop a scalable method to reliably compute domain weights when domain composition changes?
extended our experiments
CHAMELEON has such scalability. Unlike DoReMi and DoGE, CHAMELEON's algorithm does not alter the proxy model's optimization, instead it focuses more on the intrin- sic data characteristics, where the trained proxy model can already capture domain features, even for new unseen data.
DoReMi
Stability and Practicality. We show that domain weights obtained by CHAMELEON remain stable across different training steps of the proxy model, whereas DoReMi and DoGE are sensitive or converge slowly as shown in Figure 3. Additionally, our method is robust to variations in proxy model size, the hyperparameter X, and the number of sam- ples computing embeddings, as detailed in Appendix B.5. This stability significantly reduces the need for extensive hyperparameter tuning, making CHAMELEON tical and resource-efficient for real-world applications. In addition, we report GPU hours in Table 4 for domain weight computation and for training 684M base model, showing that the proxy training cost is non-negligible. Compared to DoReMi and DoGE, we reduce computational overhead to less than 2% of final training cost. This reduction crucial for academic labs and smaller-scale training. More details given in Appendix B.4.
Domain Weight
DoGE
Domain Weight
prac-
CHAMELEON
Domain Weight
3
Arxiv
Book
CC
Github
Stackexchange
Wikipedia
4.2. Scalable to Pile
is common for new data be introduced during the offi- cial training of large base models, particularly when training on diverse and evolving datasets. However, existing meth- ods including DoReMi, DoGE and RegMix require retrain- ing a new proxy model from scratch whenever domains are added or removed, making them both inefficient and incon- venient for dynamic data environments. This process not
Figure 3: KRLS scores converge quickly. Domain weights across methods during proxy training. DoReMi and DoGE require more iterations to stabilize, while CHAMELEON con- verges quickly after 1k iterations. The detailed discussion in Appendix B.5,

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning.
Table 7: Transfer to new domains perplexity. Per-domain test perplexity when adapting to new data. The Human baseline (Gao et al., 2020). When domain composition changes, the other methods need to retrain proxy models from scratch and re-run domain weight optimization, while we can directly compute the KRLS of the new domains by extending the proxy model trained on previous data. We report the extra compute (in FLOPs) required adapt the data mixture the Pile.
Domain
Human
DoReMi
DoGE
CHAMELEON
RegMix
Arxiv
9.76
14.11
10.78
9.73
16.19
Dm_ mmn_mathematies
5.52
6.27
4.52
5.31
19.26
Enron_ eemills
12.82
9.96
9.39
7.77
12.06
Europarl
34.69
24.77
11.62
28.18
131.80
Freelaw
14.12
16.20
17.99
15.04
18.66
Github
5.92
5.84
4.90
6.16
9.95
Gutenberg_PP 19_
39.36
38.36
39.57
33.28
34.43
Hackemews
35.94
29.68
29.87
27.02
29.80
Nih_exporte
22.93
25.89
26.81
24.12
28.69
Philpapers
47.59
36.43
44.56
34.63
51.71
Pile_cc
43.19
32.85
58.17
34.90
32.30
Pubmed abstracts
Pubmed_abstracts
17.87
24.19
25.62
21.87
23.20
Pubmed _entral
9.76
9.43
8.10
7.36
13.80
Stackexchange
10.41
11.48
11.79
11.25
18.96
Ubuntu _irec
36.12
32.10
23.20
29.34
20.71
Uspto sypto_backgroundh
17.22
21.19
20.08
18.25
22.05
Wikipedia en
28.70
25.95
40.42
24.68
29.32
Average PPL (L)
23.05
21.45
22.79
19.94
30.17
# Domains Over Human
9/17
10/17
11/17
4/17
Extra FLOPs
0
1.34 10
6.68 x 10
4.62 10
3.5 X 10
(290x)
(145x)
(1x)
(758x)
Table 8: Transfer to new domains reasoning. Accuracy of downstream tasks in the same settings as Table 7.
used for pretraining LMs, with overlapping data sources such as books, scientific texts, web content, and codehases. The Pile dataset includes more domains than Slimpajama and its data is more diverse. Note that the original Pile dataset includes 22 domains but only 17 are now available due to copyright issues.
Benchmark Hum.
DoReMi
DoGE
CH.
RegMix
ARC-E
37.3
39.3
35.3
39.2
39.5
COPA
56.8
61.5
54.7
60.9
61.2
HellaSwag
26.7
27.3
26.1
27.4
27.3
To obtain domain weights Pile, we input 4k samples per domain of Pile to the proxy model trained on Slimpajama in Section 4.1 and infer embeddings for computing domain weights through (KRLS). The computed domain weights reported in Table 19, where reported in their respective papers for DoReMi and RegMix. We use Human as baseline that is selected manually as in the Pile paper (Gao et al., 2020). As in Section 4.1, we report per-domain perplexity in Table 7 and downstream accuracy Table
Lambada
12.5
15.8
9.3
15.9
15.4
LogiQA
22.5
21.2
22.1
23.8
21.9
MultiRC
56.7
56.5
57.2
57.3
56.2
OpenBook
13.3
13.1
13.1
14.2
14.5
the domain weights
PiQA
57.8
59.3
55.9
59.7
60.4
QQP
37.5
36.8
36.8
37.2
37.6
RACE
25.8
27.2
24.9
26.8
27.2
SciQ
64.1
65.7
58.1
66.0
67.1
Social IQA
35.0
36.0
34.2
36.6
36.3
WinoGrande
50.7
51.2
49.8
50.9
49.9
CHAMELEON outperforms DoReMi, DoGE and even Reg- Mix in both cases. Importantly, FLOPs of CHAMELEON is marginal compared with DoReMi and DoGE (only 1% of training a new proxy model from scratch) since we can reuse the previous proxy model and our extra FLOPs only include inference cost for extracting embeddings, while DoReMi DoGE, and RegMix require retraining the proxy models.
Average (t)
38.2
39.3
36.7
39.7
39.6
To test its scalability, we employ the proxy model trained on Slimpajama in Section 4.1 to Pile dataset (Gao et al., 2020) directly. Both SlimPajama and Pile large-scale datasets
8

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Table 9: Finetuning, Per-domain test perplexity compared with the uniform baseline for finetuning on the 7 languages in Wiki40b. CHAMELEON can flexibly be extended to the finetuning pipeline by computing the KRLS of the finetun- ing domains, improving performance across all domains.
Table 10: Finetuning. Per-domain test PPL vs. the uniform baseline for finetuning on the 7 programming languages in Stack dataset. CHAMELEON improves across all domains.
Domain
Uniform
CHAMBLEON
Python
19.98
16.53
Domain
Uniform
CHAMELEON
Java
19.27
15.53
French
6.86
6.51
C
28.24
22.58
German
10. 12
8.78
C++
25.16
21.09
Italian
13.29
12.42
Go
30.25
19.26
Spanish
8.41
8.04
Ruby
21.78
17.83
Portuguese
8.00
7.78
PHP
9.45
7.43
Dutch
13.98
12.30
Average PPL (1)
22.02
17.18
Polish
5.07
4.21
# Domains Over Uniform
7/7
Average PPL (L)
9.43
8.58
Domains Over Uniform
7/7
5. Conclusion
Note that as the number of domains increases, DoGE com- putes k 17 gradients per iteration, resulting in approxi- mately 2.5x wall-clock time per iteration this dataset.
We introduce CHAMELEON, a novel and efficient framework for data mixing that leverages KRLS to quantify the repre- sentativeness of data domains. We demonstrate that inverse ERRLS-based domain weights effectively identify highly im- portant domains for pretraining LMs. CHAMELEON can adapt to new domains without retraining proxy models, out- performing baselines in downstream tasks. Given that it is computationally inexpensive and stable, CHAMELEON lowers the overall cost of the expensive LLM pretraining pipeline, which can be useful both in industry and within academic budgets. We also extend domain reweighting to fine-tuning with sistent improvements.
4.3. Finetuning
Besides pertaining, CHAMELBON interestingly allows data- mixture optimization during finetuning. We take advantage of the existing pretrained model and can extract embeddings on finetuning data directly for domain weight computation. As discussed in Section 3.2, the goals of pretraining and fine-tuning are distinct, with pretraining aiming for broad generalization and fine-tuning focusing on specialization. Therefore, we directly use leverage scores for computing domain weights, as described in Section 3.2.
weights, demonstrating con-
This work highlights the need to tailor data- strate- gies to different training phases. In future work, we aim to extend our approach to online settings for dynamic optimiza- tion during training. Additionally, we will extend to target specific downstream tasks by modifying the identity matrix within the KRLS to emphasize relevant domains, enhancing our method's 's flexibility for specific downstream tasks.
We fine-tune pretrained model trained on the Pile (from Section 4.2) for 10k steps on two separate datasets: (i) Wiki40b (Guo et al., 2020), which includes multiple lan- guages, for which we select Latin languages, and (ii) Stack-decup (Kocetkov et al., 2022), which covers various programming languages, from which we use 7. The results shown in Table 9 and Table 10. CHAMELEON outper- forms the uniform weights baseline across all domains in both tasks. showing our data mixture can greatly benefit finetuning. Remarkably, our weight computation is compu- tationally cheap in finetuning, as we simply need forward passes through the pretrained model to compute domain embeddings and we can then directly apply (KRLS).
Acknowledgements
We thank the reviewers for their constructive feedback. Thanks to Simin Fan for the helpful discussion. This work was supported as part of the Swiss AI Initiative by a grant from the Swiss National Supercomputing Centre (CSCS) under project ID a06 on Alps. This work was supported by the Swiss National Science Foundation (SNSF) under grant number 200021 1000221205011..11 1. This work was supported by Hasler Foundation Program: Hasler Responsible AI (project number 21043). Research was sponsored by the Army Re- search Office and was accomplished under Grant Number W91 24-22----000
Additionally, present fine-tuning results using PT in- stead of FT in Appendix B.9 for reference. The results demonstrate that fine-tuning with KRLS-based domain weights outperforms using their inverse. This indicates that data data-mixing strategies should be tailored to different training phases.
9

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Impact Statement
sampling. In Proceedings of the Twenty-Eighth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 1758- 1777. SIAM, 2017.
The approach presented in this paper aims at advancing the field of Machine Learning. No other potential societal con- sequences of our work are deemed necessary specifically highlight here.
De Marchi, S., Sommariva, A., and Vianello, M. Mul- tivariate christoffel functions and hyperinterpolation. Dolomites Research Notes on Approximation, 7(Special Issue), 2014.
References
Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., et al. Glam: Efficient scaling of language models with mixture-of- experts. In International Conference on Machine Learn- ing (ICML), pp. 5547-5569, PMLR, 2022.
Alaoui, A. E. and Mahoney, M. W. Fast Random- ized Kernel Methods With Statistical Guarantees, 2015. arXiv:1411.0300
Albalak, A., Pan, L., Raffel, C., and Wang, W. Efficient online data mixing for language model pre-training, In RO- FoMo: Robustness of Few-shot and Zero-shot Learning in Large Foundation Models Workshop, 2023.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Bach, F. Sharp analysis of low-rank kernel matrix approxi- mations. In Conference on Learning Theory, pp. 185-209. PMLR, 2013.
Ducharlet, K., Travé- Massuyès, L., Lasserre, J.-B., B., Le Lann, M.-V., and Miloudi, Y. Leveraging the christoffel func- tion for outlier detection in data streams. International Journal of Data Science and Analytics, 2024.
Beckermann, B., Putinar, M., Saff, E. B., and Stylianopou- los, N. Perturbations of christoffel narrbou kernels: De- tection of outliers. Foundations of Computational Mathe- matics, 21(1):71-124, 2021.
Dunkl, C. F. and Xu, Y. Orthogonal polynomials of sev- eral variables, volume 155, Cambridge University Press, 2014.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In AAAI Conference on Artificial Intelligence, 2020.
Fan, S., Grangier, D., and Ablin, P. Dynamic gradi- ent aligmment for online data mixing. arXiv preprint arXiv:2410.0244 2024a.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pp. 1877-1901, 2020.
Fan, S., Pagliardini. M., and Jaggi, M. DOGE: Domain reweighting with generalization estimation. In Interna- tional Conference Machine Learning (ICML), 2024b.
Fanuel, M., Schreurs, J., and Suykens, J. A. Nyström land- mark sampling and regularized christoffel functions. Ma- chine Learning 1 116):2213322555 2022.
Calandriello, D., Lazaric, A., and Valko, M. Analysis of nyström method with sequential ridge leverage score sam- pling. In Uncertainty in Artificial Intelligence Conference, 2016.
Feng, S., Prabhumoye, S., Kong, K., Su, D., Patwary, M., Shocybi, M., and Catanzaro, B. Maximize your data's potential: Enhancing Ilm accuracy with two-phase pre- training. arXiv preprint arXiv:2412.15588 2024.
Chen, Y. and Yang, Y. Fast statistical leverage score ap- proximation in kernel ridge regression. In International Conference on Artificial Intelligence and Statistics (AIS- TATS), pp. 2935-2943. PMLR, 2021.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T. Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022 2020.
Clark, P., Cowhey, I., Etzioni, O.,, Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1803. 05457,, 2018
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac'h, h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang,
Cohen, M. B., Musco, C., and Musco, C. Ridge lever- age scores for low-rank approximation. arXiv preprint arXiv:1511.07263, 6, 2015.
J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika,
L., Tang, E., Thite, A., Wang, B., Wang. K., and Zou, framework for few-shot language model evaluation, 2024.
Cohen, M. B., Musco, C., and Musco, C. Input sparsity time low-rank approximation via ridge leverage score
10

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Gareth, J., Daniela, W. Trevor, H., and Robert, T. An introduction statistical learning: with applications in R. Spinger, 2013.
Longpre, S., Yauney, G., Reif, E., Lee, K., Roberts, A., Zoph, B., Zhou, D., Wei, J., Robinson, K., Mimno, D., and Ippolito, D. A pretrainer's 's guide to training data: Measuring the effects of data age, domain coverage, qual- ity, & toxicity. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Pp. 3245-3276 Association for Computational Linguis- tics, 2024.
Guo, M., Dai, Z., Vrandečić, D., and Al-Rfou, R. Wiki-40b: Multilingual language model dataset. In Proceedings of the Twelfth Language Resources and Evaluation Confer- ence, 2020.
Hastie, T., Tibshirani, R., Friedman, J., and Franklin, J. The elements of statistical learning: data mining, inference and prediction. The Mathematical Intelligencer, 27(2): 83-85, 2005.
Ma, Y., Liu, Y., Yu, Y., Zhang, Y., Jiang, Y., Wang, C., and Li, S. At which training stage does code data help Ilms reasoning? arXiv preprint arXiv:2309. /6298, 2023.
Jiang, Y., Zhou, A., Feng, Z., Malladi, S., and Kolter, J. Z. Adaptive data optimization: Dynamic sample selection with scaling laws. arXiv preprint arXiv:2410.1188. 2024.
Mahoney, M. W. and Drineas, P. CUR matrix decompo-
sitions for improved data analysis. Proceedings of the National Academy of Sciences, 106(3):697-702, 2009.
Mallen, A., Asai, A., Zhong, V., Das, R., Khashabi, D., and Hajishirzi, H. When not to trust language models: Inves- tigating effectiveness of parametric and non-parametric memories. In Annual Meeting of the Association for Com- putational Linguisties (ACL), Toronto, Canada, July 2023. Association for Computational Linguistics.
Kang, F., Sun, Y.. Wen, B., Chen, S., Song, D., Mah- mood, R., and Jia, R. Autoscale: Automatic prediction of compute-optimal data composition for training llms. arXiv preprint arXiv:2407.20177, 2024.
Khashabi, D., Chaturvedi, S., Roth, M., Upadhyay, S., and
Roth, D. Looking beyond the surface: challenge set for reading comprehension over multiple sentences. In Pro- ceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2018.
McInnes, L., Healy, J., and Melville, J. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv: 2018.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint (rXX:::809.027777 2018.
Kocetkov, D., Li, R., Allal, I.. B., Li, J., Mou, C., Ferrandis, C. M., Jernite, Y., Mitchell, M., Hughes, S., Wolf, T., et al. The stack: th of permissively licensed source code. arXiv preprint
2022.
Musco, C. and Musco, C. Recursive sampling for the nys- trom method. In Advances in Neural Information Process- ing Systems (NeurIPS), volume 30. Curran Associates, Inc., 2017.
Lai, G., Xie, Q. Liu, H., Yang. Y.. and Hovy, E. Race: Large-scale reading comprehension dataset from exami- nations. arXiv preprint arXiv: 1et.:040444... 2017.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q. N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The LAMBADA dataset: Word predic- tion requiring broad discourse context. arXiv preprint /606v//006.0600 2016.
Lasserre, J. B. and Pauwels, E.
The empirical christoffel function with applications in data analysis. Advances in Computational Mathematics, (5(3):1439-1485) 2019.
Li, M., Miller, G. L., and Peng, R. Iterative row sampling. In 2013 IEEE 54th Anmual Symposium on Foundations of Computer Science, pp. 127-136, 2013.
Parmar, J., Prabhumoye, S., Jennings, J., Liu, B., Jhunjhun- wala, A., Wang. Z., Patwary, M., Shoeybi, M., and Catan- zaro, B. Data, data everywhere: A guide for pretraining dataset construction. arXiv preprint arXiv:2407.086 2024.
Liu, J., Cui, L., Liu, H., Huang. D., Wang, Y., and Zhang, Y. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08122 2020.
Parthasarathy, V. B., Zafar, A., Khan, A., and Shahid, A. The ultimate guide to fine-tuning 1lms from basics to breakthroughs: An exhaustive review of technologies, research, best practices, applied research challenges and opportunities. arXiv preprint arXiv:2408, 13296, 2024.
Liu, Q., Zheng, X., Muennighoff, N., Zeng, G., Dou, L., Pang, T., Jiang, J., and Lin, M. Regmix: Data mixture as regression for language model pre-training. arXiv preprint arXiv:2407.01492. 2024.

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Pauwels, E., Bach, F., and Vert, J.-P. Relating leverage scores and density using regularized christoffel functions. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, 2018.
Welbl, J., Liu, N. E,, and Gardner, M. Crowdsoure- ing multiple choice science questions. arXiv preprint TrXiv:177.006299 2017.
Xie, S. M.. Pham, H., Dong, X., Du, N., Liu, H., Lu, Y., Liang, P., Le, Q. V., Ma, T., and Yu, A. W. DoReMi: Optimizing data mixtures speeds up language model pre- training. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Rudi, A., Carratino, L., and Rosasco, L. Falkon: An opti- mal large scale kernel method. In Advances in Neural Information Processing Systems (NeurIPS), volume 30. Curran Associates, Inc., 2017.
Rudi, A., Calandriello, D., Carratino, L.,, and Rosasco, L.. On fast leverage score sampling and optimal learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 31. Curran Associates, Inc., 2018.
Yang, Y., Mishra, S., Chiang, J. N., and Mirzasoleiman, B. SmallToLarge (S2L): Scalable Data Selection for Fine-tuning Large Language Models by Summarizing Training Trajectories of Small Models. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang. P. Distributionally robust neural networks. In International Conference on Learning Representations (ICLR), 2020.
Ye, J., Liu, P., Sun, T., Zhou, Y., Zhan, J., and Qiu, X. Data mixing laws: Optimizing data mixtures by pre- dicting language modeling performance. arXiv preprint arXiv:2403 16952, 2024.
Sakaguchi, K., Bras, R. L.., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64, 2021.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. HellaSwag: Can machine really finish your sentence? In Annual Meeting of the Association for Computational Linguistics (ACL), pp. 4791 1-4800, 2019.
Sap. M., Rashkin, H., Chen, D., LeBras, R., and Choi, Y. Socialiqa: Commonsense reasoning about social interac- tions. arXiv preprint (904:09900007000 2019.
Sarlin, P.-E., DeTone, D., Malisiewicz, T., and Rabinovich, A. Superglue: Learning feature matching with graph neural networks. In Conference on Computer Vision and Pullern Recognition (CVPR), 2020.
Schölkopf, B., Herbrich, R., and Smola, J. generalized representer theorem. Conference on Learning Theory, pp. 416-426. Springer, 2001.
Shen, Z., Tao, T., Ma. L., Neiswanger, W., Liu, Z., Wang, H., Tan, B., Hestness, J., Vassilieva, N., Soboleva, D., et al. Slimpajama-de: Understanding data combinations for 1lm training. arXiv preprini 2309. 10818, 2023.
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, R., Hestness, J., and Dey, N. SlimPajama: A 627B token cleaned and deduplicated version of RedPa- jama, 2023. URL https: https:///ugginggacee co/ Latasets/cereebrase
5 watpeass/ssmpPPe..
627B.
Vapnik, V. N. An overview of statistical learning theory. IEEE Transactions on Neural Networks, 10(5):988-2) 9999 1999.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.,, Gomez, A. N., Kaiser, I.. u. and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Wang, A. Glue: A multi-task benchmark and analysis plat- form for natural language understanding. arXiv preprint arXiv: 1804.804400 1804.07461, 2018
12

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Appendix
Table of Contents
A Additional discussions the KRLS
14
A.1 Details of leverage scores
14
A.2 Discussions
15
Additional experiments
16
B.1 Additional figures for Figure
16
B.2
Experimental setup.
16
B.3 Domain weights Slimpajama
17
B.4 Temporal cost
17
B.5 Stable domain weights of CHAMELEON
17
B.6 Additional baseline: Data Mixing Laws
18
B.7 Domain weights on Pile
19
B.8 Domain weights for finetuning
19
B.9 PPL of finetuning...
19
13

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
A. Additional discussions on the KRLS
A.1. Details of leverage scores
In this work, we employ KRLS to assign scores to data domains, and we focus on the leverage scores as a measure of domain importance, both in pretraining and finetuning language models. The KRLS (Alaoui & Mahoney, 2015) is a cernelized-versionn of the ridge leverage scores, which are used quantify the importance of the rows in a matrix for the best-rank approximation with approximation error guarantees (Mahoney & Drineas, 2009; Li et al., 2013). The KRLS for the i-th domain is defined as
S = (9p + k^
(I)
with kernel matrix SDi = K([T, xz) and regularization constant À > O. In our data-centric approach, we rank domains based on the degree of inter-domain dependency. To better see this, we establish the equivalence of RLS computed in the feature space induced by the finite-dimensionl feature map & and KRLS when employing the corresponding kernel k(x,, xyy) = (xx) o(x,), where we use ((x) = Section 3.1. This result provides insights into the behavior of the employed KRLS for domain reweighting. Let the of domain embeddings {2i) where J; IR, Let Φ R² - Rd
be tiite-ddmensionall feature map with associated p.s.d. kernel 2z) ((x,)) (oxc)""(an )(xy) by kernel trick (Vapnik. 1999), and ((X) E # be the matrix whose rows are the feature mappings <(x;)" T. The ridge regression hat matrix in feature space H ((X) (X) + -x)= with ridge leverage scores diag (H In kernel ridge regression, the kernel ridge hat matrix is H, = (DD A1) and the kernel ridge leverage scores given by diag(HH.)
Lemma A.1. With kernel k(x, 1) 6(z)² o(y), where o: Rdd IR a fifiite-dimenssinn feature map, the RLS in feature space induced by & and the KRLS are s.t. diag(H))) =
Proof. With the kernel s(2, v) = <(x)² (v)), we have lt (XX) Substituting into Hn:
H₂ =
-
We utilize the following matrix identity: for matrices A E mmxn and B E Rexm and scalar À + O.
A(BA + AI (AB + A
This identity can be verified by multiplying both sides by (BA+ XI) from the right, then by (AB + AI) from the left, which yields the same result on both sides.
Let A = (X) and B = T(X)T. Then:
+(X)0(1))0(x +
Therefore,
" = (X))
(X(X)((X)) A1) x HR.
Thus, the diagonal elements, and hence the leverage scores, are equivalent: diag((H)) diagg
The above lemma establishes the equivalence of the ridge leverage scores computed in the feature space induced by the finite-dimensional feature map o and the KRLS when employing the corresponding kernel k(x,, ary) ((x:)) T (T(())))))))))))))())) It is possible to relate the RLS to the ridgless solution, where the regularized solution converges to the least-norm solution as + O. Let the eigendecomposition of 2pp UEUT, then the KRLS of domain i, i = I,..., can be written as
Sx(D,) Σ o, +
Jj U2
where oj is the j-th eigenvalue of Ωd. Therefore, the KRLS is weighted version of the standard statistical leverage (Gareth et al., 2013), i.e., Σj- U2,- with weights depending on the regularization and the eigenspectrum of 2p. We now recall the relationship between the least-norm solution a system of equations to the (ridgless) statistical leverage score = ((x.) (+(X) ((X)) +x(x)))))))))) (Mahoney Drineas, 2009).
14

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Lemma A.2. The least-norm solution characterizes &g as follows:
4 min llvel22
((X) d(a))))
The leverage score of domain measures how important it in composing the row space of (()). If a row (domain) has a component orthogonal to all other rows (domains), its leverage score 1. In data mixture, (2) seeks linear combination of features that best approximates the embedding 2; of the i-th domain. Intuitively, & is highest when I; is linearly independent from the other domain embeddings. In our approach, we use the KRLS to assign scores to data domains both in pretraining and finetuning language models. High scores indicate data domains that are difficult to approximate with a linear combination of other domains, and are thus more unique. On the other hand, low scores indicate data domains that show higher degree of dependency with other domains, identifying more common data characteristics. During pretraining, we rank domains with lower KRL.S as more important, as they are more common and thus more useful for learning general-purpose language representations. During finetuning specific task, we upweight domains with higher KRLS, they are more unique and therefore better suited for learning task-specific representations.
A.2. Discussions
We now discuss additional theoretical insights and properties of our methodology.
Role of regularization. Adding AI to Sip in (I) reduces the influence of the small principal components, resulting in proportionately lower sampling probability. A large À soft-thresholds the low part of the spectrum of Sip and amplifies the contribution of the top eigenvectors of the kernel matrix, focusing on the most dominant domains. In practice, due to the relatively small number of domains, we observe that our algorithm is robust to the choice of X, where small regularization is sufficient and larger values do not significantly alter the computed domain weights.
Adaptability to larger models. Our approach can robustly transfer domain weights obtained from small proxy model hep to a larger target model he, thanks to its use of domain affinities. Specifically, our method computes domain weights based on the kemel matrix lbp which captures pairwise inner products between domain embeddings. While the absolute embeddings z; may vary between the proxy and target model, their inner products show a higher degree of consistency across model sizes. For instance, if the proxy learns that the "github" and "stackexchange" domains are semantically close, a larger, more powerful model typically also maintains this proximity in embedding space. Consequently, the pairwise similarities encoded in 2dp, and therefore the resulting KRLS scores (I) and domain weights, are robust across different model scales. Empirical evidence is presented in Table 16.
Relation with Christoffel functions. In machine learning literature, the Christoffel function a key concept that characterizes the local density of the data distribution in feature space (Pauwels al., 2018). Christoffel functions are known in orthogonal polynomials (Dunkl & Xu, 2014) and approximation theory (De Marchi et al., 2014). They are extended to machine learning (Pauwels et al., 2018), where it makes the connection between inverse leverage scores and the kernelized Christoffel function. In machine learning, they are mainly used for landmark sampling (Fanuel et al., 2022), density estimation (Pauwels et al., 2018), and outlier detection (Lasserre & Pauwels, 2019; Beckermann et al., 2021; Ducharlet et al., 2024). Given samples { the kernelized Christoffel function is defined as the following regularized minimization over a reproducing kernel Hilbert space (RKHS) H with associated kernel Kgg
Cx,(()) inf
L g(zz)) + 21 s.t. g(z) - 1,
gEN
where is a regularization constant, || denotes the RKHS norm of g, and nú > O. The Christoffel function is linked to the ridge leverage scores (RLS) (Alaoui & Mahoney, 2015: Cohen al., 2017: Rudi et al., 2018). which quantify the influence of each sample on the learned model. Specifically, the Christoffel function at point is proportional to the inverse of its RLS. High RLS values indicate data points that are difficult to represent as linear combinations of other points in the feature space. Conversely, a high Christoffel function value (and thus low RL.S) suggests data point lies in a region of high data density and can be better expressed in terms of other points. In our work, focusing on domains with high degree of linear dependency, i.e., high Christoffel function, is shown to enable improved generalization and transfer learning capabilities. The following lemma details the dlosed-form expression for the regularized kernelized Christoffel function each sample.
15

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Lemma A.3 (Paawwels al., 2018)). The regularized kernelized Christoffel function takes the following value sample i. for j ,...,,
(xn(z;)= = !á (K (K + nxdiaa(n) n
(4)
jj
where K IR is the kernel matrix with entries Kiz (2, 2j), and diag(1) diag (oog(m)-... mm...., is diagonal matrix with entries fl the diagonal.
This is derived from the representer theorem applied to (3) (Schölkopf et al., 2001; Pauwels et al., 2018). Lemma A.3 reveals the connection between the Christoffel function and the KRLS (1). The Christoffel function (x,1(2)) is therefore inversely proportional to KRLS of sample 2y with rj 1.
B. Additional experiments
B.1. Additional figures for Figure
10
Ardiv
42.2
-16.9
-14.5
-14.5
5.6
2.9
-4.9
40
Brok
16.99
26.7
12
14.6
-2.
16.1
0.9
UMAP
CC
12
18.7
Relationship Strength
20
C4
14.5
22.1
-1..8
-16.7
-2.3
Gathub
5.8
21.2
-0033
-11.8
45.7
31.1
-19
Succexchange
2.9
16.1
13.3
-16.7
31.1
29
16.9
Wikipedlia
0.9
223
-19
-16.9
48
20
15
UMAPI
Azdy
Book
CC
C4
Gitbab
Stokerchange
Wikipodia
wwikipedia bboo
cc
c4 arxiv github stackexchange
Figure 4: Domains in embedding space. 2D UMAP vi- sualization of embeddings SlimPajama learned by the proxy model. Semantically similar domains occupy sim- ilar regions in embedding space, creating high-density clusters.
Figure 5: Domain affinity matrix. The matrix shows the relationship strength between domains in SlimPajama.
B.2. Experimental setup-
We follow the experimental setup of DoGE (Fan al., 2024b), using small 82M decoder-only Transformer (Vaswani et al., 2017) as the auxiliary model for CHAMELEON, DoReMi, and DoGE. Additionally, we use 684M model the base model for pretraining. Moreover, we set the batch size 128, the cosine learning rate scheduler, weight decay 0.01, and gradient clipping 1.0 for all models. For the training on Slimpajama, Wiki40b, and Stack datasets, we set batch size 128. We increase the batch size to 512 on the Pile dataset since is more noisy and has larger number of domains.
In CHAMELEON, temperature factor т in the softmax normalization for domain weights additionally applicable: a; exp(= /r) / exp(xs/7), where S- (D) for pretraining and (D) for fine-tuning. In our experiments, typically set [5, 10 and PT 0.5).
For the ablation study of CHAMELEON, we also evaluate proxy models with other sizes (60M, 124M, and 210M). The model architectures and their corresponding learning rates are reported in Table 11.
For the setup of RegMix on the Slimpajama dataset, we follow its original setup (Liu et al., 2024) and train 200 1M proxy models with k steps fit a regression model.
16

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Table II: Architecture hyperparameters for various model scales.
Layers
Attention heads
Embed dim
Hidden dim
Max. learning rate (min.)
60M
3
768
3072
5 x 10 (1x 10-1)
82M
6
12
768
3072
5 10 (1x 10
124M
12
12
768
3072
5 10 (1x 10-1)
210M
24
16
768
3072
5 10=*(1x 10-+)
684M
36
24
1200
4800
1.5 x 10="(xx
1.2B
36
25
1600
6400
1.5 x 10 10
B.3. Domain weights on Slimpajama
We report our final domain weights for base model training Table 12. Specifically, DoReMi and DoGE use domain weights through training proxy models with 10k steps. CHAMELEON use the model with 2k steps. For RegMix, we follow its paper (Liu et al., 2024) using "CC" the target domain and train 200 1M proxy models to get the domain weights.
Table 12: Final domain weights.
Domain
DoReMi
DoGE CHAMBLEON
RegMix
Arxiv
0.057
0.041
0.083
0.001
Book
0.002
0.078
0.164
0.025
CC
0.237
0.268
0.202
0.924
C4
0.237
0.283
0.247
0.024
Github
0.130
0.059
0.082
0.019
Stackexchange
0.101
0.230
0.149
0.006
Wikipedia
0.236
0.041
0.073
0.001
B.4. Temporal cost
In Table 13, we report the required GPU (H100) hours for obtaining domain weight regarding the experiments in Table 2. We also show GPU hours training 684M base model for reference. We claim that improving the efficiency of determining the domain mixture is essential because the associated computational cost of proxy training is non-negligible. Compared to DoReMi and DoGE, which add over 10% to base model training costs, we reduce computational overhead less than 2% of final training cost. This reduction is crucial for academic labs and smaller-scale training.
More importantly, the computational cost reported is often optimistic lower bound for the baselines even for larger base models, since DoReMi and DoGE require extensive hyperparameter tuning. It has been shown that DoReMi's weights are unstable or difficult to reproduce (Fan et al., 2024b; Parmar et al., 2024) and DoGE approximations make it more sensitive to learning rate (Kang et al., 2024). Such sensitivity necessitates repeated validation on base models. Nevertheless, CHAMELEON insensitive to hyperparameters and detailed discussion is in Appendix B.5.
Table 13: GPU hours for obtaining domain weight.
DoReMi
DoGE
CHAMEEEEN
684M base model
7.4h
6.3h
0.8h
56h
B.5. Stable domain weights of CHAMELEON
Table 14 corresponds to Figure and reports the specific domain weights obtained by the proxy model with the different number of steps. CHAMELEON is more stable and converges faster than DoReMi and DoGE. In addition, we report perplexity using domain weights derived from a proxy model trained for 2k steps for DoReMi and DoGE in Table 15. This further demonstrates that DoReMi and DoGE converge slower than CHAMELEON
17

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
Furthermore, Table 16 demonstrates that CHAMELEON is also very stable across different model sizes, À values, and number of samples for embedding computations.
Table 14: Domain weights at different steps.
Domain
DoReMi
DOGE
CHAMELEON
1k
2k
5k
10k
1k
2k
5k
10k
1k
2k
5k
10k
Arxiv
0.251
0.172
0.095
0.057
0.116
0.092
0.060
0.041
0.088
0.083
0.072
0.096
Book
0.003
0.008
0.004
0.002
0.139
0.131
0.102
0.078
0.149
0.164
0.174
0.158
CC
0.080
0.114
0.153
0.237
0.177
0.209
0.247
0.268
0.174
0.202
0.188
0.195
C4
0.080
0.115
0.154
0.237
0.176
0.210
0.259
0.283
0.281
0.247
0.220
0.251
Github
0.215
0.138
0.241
0.130
0.121
0.101
0.074
0.059
0.096
0.082
0.094
0.083
Stackexchange
0.095
0.146
0.118
0.101
0.155
0.167
0.200
0.230
0.137
0.149
0.137
0.136
Wikipedia
0.276
0.308
0.235
0.236
0.116
0.090
0.058
0.041
0.074
0.073
0.082
0.080
Table 15: Perplexity using domain weights derived from a proxy model trained for 2k steps.
Domain
DoReMi
DoGE
CHAMELEON
Arxiv
8.08
8.49
8.31
Book
52.07
40.38
39.23
CC
48.69
41.31
40.11
C4
52.98
44.54
42.59
Github
3.99
4.05
4.20
Stackexchange
7.98
7.81
7.94
Wikipedia
10.57
13.98
13.90
Average PPL (4)
26.34
23.01
22.31
Table 16: Domain weights across different model sizes, values, and number of samples.
Model Sizes
À Values
Number of Samples
Domain
60M
82M
124M
210M
==1
X 10
X 100
2k
4k
8k
Arxiv
0.084
0.083
0.087
0.093
0.080
0.083
0.087
0.079
0.083
0.081
Book
0.158
0.164
0.157
0.165
0.159
0.164
0.173
0.165
0.164
0.170
CC
0.170
0.202
0.169
0.170
0.178
0.202
0.201
0.193
0.202
0.209
C4
0.271
0.247
0.288
0.259
0.243
0.247
0.258
0.253
0.247
0.241
Github
0.073
0.082
0.072
0.089
0.097
0.082
0.057
0.079
0.082
0.066
Stackexchange
0.152
0.149
0.153
0.145
0.152
0.149
0.128
0.138
0.149
0.143
Wikipedia
0.072
0.073
0.074
0.078
0.081
0.073
0.095
0.093
0.073
0.090
B.6. Additional baseline: Data Mixing Laws
We an additional baseline, Data Mixing Laws (Ye al., 2024), for comparison. It derives domain weights by leveraging scaling laws of training steps, model sizes, and data mixtures predict the performance large models trained on diverse data from small-scale training. This requires training multiple small proxy models with varying domain weights, making more computationally expensive than ours, which trains just one proxy model.
We use their reported domain weights to train 684M model on Slimpajama. Since their weights are optimized with the Pile as the target, they may be suboptimal for SlimPajama. However, given the alignment of their objectives and overlap in data sources, we consider the comparison meaningful.
As shown in Table and Table 18, CHAMELEON outperforms Data Mixing Laws in both perplexity and downstream tasks
18

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
fraction of the cost. Data Mixing Laws' FLOPs is calculated for 4 different proxy sizes and 20 separate mixtures, where cost is 2 orders of magnitude lower.
Table 17: Data Mixing Laws perplexity. Perplexity in the same settings as Table 2.
Domain
CHAMELEON
Data Mixing Laws
Arxiv
8.31
7.55
Book
39.23
45.06
CC
40.11
44.21
C4
42.59
45.79
Github
4.20
4.01
Stackexchange
7.94
7.96
Wikipedia
13.90
16.20
Average PPL ())
22.31
24.40
# Domains Over Uniform
4/7
4/7
FLOPs
1.36 x 10
5.36 x 10
(1x)
(394x)
Table 18: Data Mixing Laws reasoning, Accuracy of downstream tasks in the same settings as Table 2.
Benchmark CHAMELEON
Data Mixing Laws
ARC-E
37.8
34.5
CОРA
61.9
59.0
HellaSwag
27.1
27.4
Lambada
15.1
14.7
LogiQA
22.6
26.0
MultiRC
57.2
57.2
OpenBook
14.4
25.2
PIQA
60.5
58.5
QQP
39.2
36.8
RACE
26.5
26.4
SciQ
64.3
57.2
Social IQA
35.7
36.1
WinoGrande
52.1
48.4
Average (t)
39.6
39.0
B.7. Domain weights on Pile
We report the domain weights we use on the Pile dataset in Table 19. Note that CHAMELEON and DoGE are from our own experiments, Human is suggested in (Gao et al., 2020), DoReMi uses the same as (Xie et al., 2023), and RegMix uses weights from (Liu et al., 2024).
B.S. Domain weights for finetuning
We report the FT on Wiki40b and Stack datasets separately below, which corresponds to Section 4.3.
B.9. PPL of finetuning.
We report the results of fine-tuning with a PT in Table 22 and Table 23 for reference. It is clear to see that fine-tuning with PRL.S-based domain weights is better than the one with the inverse of KRLS-basee weights.
19

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Table 19: Domain weights on Pile.
Domain
Human
DoReMi
DoGE
RegMix
CHAMELEON
Arxiv
0.134
0.004
0.0608
0.001
0.0386
Dm_ mnn_matemaaiis
0.025
0.002
0.0280
0.000
0.0538
Enron_emaiin
0.004
0.009
0.0239
0.002
0.0085
Europarl
0.005
0,008
0.0407
0,000
0.0048
Freelaw
0.049
0.005
0.0293
0.001
0.0147
Github
0.054
0.022
0.0693
0.000
0.0099
Gutenberg_PP
0.025
0.009
0.0283
0.002
0.0115
Hackernews
0.010
0.016
0.3949
0.012
0.0637
_ih_exporter
0.007
0.008
0.180
0.001
0.0424
Philpapers
0.003
0.034
0.0266
0.000
0.0226
Pile_cc
0.142
0.743
0.0348
0.870
0.4519
0.107
0.014
0.0398
0.024
0.0104
Pubmed central
0.136
0.006
0.0251
0.003
0.1207
Stackexchange
0.118
0.019
0.0266
0.000
0.0226
Ubuntu ire
0.009
0.011
0.0474
0.064
0.0123
Uspto_ backkooudsss
0.053
0.004
0.0366
0.002
0.0212
Wikipedia en
0.117
0.086
0.0425
0.016
0.1075
Table 20: Wiki40b Domain Weights.
Table 21: Stack Dataset Training Weights.
Domain
CHAMELEON
Domain
CHAMELEON
French
0.115
Python
0.125
German
0.163
Java
0.129
Italian
0.127
0.102
Spanish
0.109
C++
0.088
Portuguese
0.090
Go
0.241
Dutch 0.140
Ruby
0.118
Polish
0.257
PHP
0.197
Table 22: Per-domain perplexity compared with the Uniform baseline for fine-tuning with 684M parameter models on the languages in the Wiki40b dataset.
Domain
Uniform
CHAMELEON (a
CHAMELEON (a
French
6.86
6.51
7.14
German
10.12
8.78
10.85
Italian
13.29
12.42
14.42
Spanish
8.41
8.04
8.70
Portuguese
8.00
7.78
8.14
Dutch
13.98
12.30
15.05
Polish
5.07
4.21
5.31
Average PPL (L)
9.43
8.58
9.94
20

---

CHAMELEON: A Flexible Data-mixing Framework for Language Model Pretraining and Finctuning
Table 23: Per-domain perplexity compared with the Uniform baseline for fine-tuning with 684M parameter models on the languages in the Stack dataset.
Domain
Uniform
CHAMELEON (a
CHAMELEON (a )T)
Python
19.98
16.53
20.11
Java
19.27
15.53
19.32
C
28.24
22.58
25.02
C++
25.16
21.09
23.83
Go
30.25
19.26
28.66
Ruby
21.78
17.83
21.75
PHP
9.45
7.43
9.47
Average PPL (L)
22.02
17.18
21.17
21


Reviews:
[{"review_id": "41XQCNaSIN", "reviewer": "Reviewer_jX7F", "summary": "Authors propose a method for data sampling for pretraining and finetuning language models. Their idea is to train a classifier, then extract the middle layer word embeddings of the classifier for each domain in the training data, and then to do matrix factorization to obtain a scalar weight for each domain. These weights are used in two heuristic equations to obtain the probabilities for sampling from the domains during the training of the LM. \n\nIt is empirically shown that the algorithm is faster than the baselines, and can be used without retraining the classifier when new data is added to the training. It is also shown that the algorithm can be used during finetuning.\n\n=================\n\nUPDATE: I updated my review score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. Theoretical motivation\n\nWe use kernel ridge leverage scores (KRLS) to determine domain weights. KRLS is a well-established tool in data analysis. It quantifies the influence or importance of data points [Alaoui & Mahoney, 2015]. This property is leveraged in machine learning for tasks like density estimation [Pauwels et al., 2018] and novelty detection [Ducharlet et al., 2024; Lasserre & Pauwels, 2019].\n\nThe inverse KRLS is proportional to the Christoffel function value [Pauwels et al., 2018]. This relationship provides additional theoretical justification for our approach. Christoffel functions (Eq. (1) in [Pauwels et al., 2018]) precisely characterize the local density of the data distribution in the feature space, where higher values indicate denser regions. \n\nIn our context, we compute the score $S_\\lambda(D_i)$ of domain $i$ using Eq. (KRLS) on page 4. During pretraining, assigning higher sampling probability to domains with low KRLS (and thus high $S_\\lambda^{-1}$/Christoffel value) upweights high-density data regions, which are most influential on base LMs' performance [1]. LLM finetuning aims to specialize on a novel specific task, requiring the model to learn differential features not fully captured during pretraining, so we instead prioritize the domains with high $S_\\lambda$. Section 3.2 converts either $S_\\lambda^{-1}$ or $S_\\lambda$ into probability distributions $\\alpha$ by appliying softmax normalization.\n\nWe will revise the phrasings in Section 3.2 to explicitly connect the data mixing goal to the mathematical properties of KRLS and inverse KRLS, clarifying its foundation in theoretical principles from statistical learning rather than just empirical heuristics. We will make this discussion self-contained within the main text, drawing upon the analysis currently in Appendix A.\n\n> Q2. Impact of our computational efficiency\n\nThe computational cost associated with determining the domain mixture via proxy training is non-negligible. Table jX7F-1 below reports the required GPU (H100) hours for our experiments in Tab. 2 in the paper. Compared to DoReMi and DoGE, which add over 10% to base model training costs, **we reduce computational overhead to less than 2% of final training cost.** This reduction is crucial for academic labs and smaller-scale training.\n\n**Table jX7F-1:** GPU hours for universal generalization experiments.\n|Method|GPU Hours|\n|-|-|\n|DoReMi|7.4h|\n|DoGE|6.3h| \n|Chameleon|0.8h|\n|684M base model|56h|\n\nEven for larger base models, the computational cost reported is often an optimistic lower bound for the baselines since *DoReMi and DoGE require extensive hyperparameter tuning*. It has been shown that DoReMi's weights are unstable or difficult to reproduce [2; Fan et al., 2024b] and DoGE approximations make it more sensitive to learning rate [Kang et al., 2024b]. We also noticed that DoGE is also extremely sensitive to their Bregman coefficient $\\mu$, as shown in Table jX7F-2, where we report domain weights and validation PPL in the last line. Small variations in $\\mu$ drastically change domain weights and degrade validation PPL, necessitating repeated validation on base models. This sensitivity contradicts the goal of data mixing methods: weights should transfer reliably to large models without costly grid searches.\n\n**Table jX7F-2:** DoGE's weights are highly sensitive to $\\mu$.\n||$\\mathbf{\\mu=0.05}$|$\\mu=0.01$|$\\mu=0.1$|\n|:-:|:-:|:-:|:-:|\n|Arxiv|0.041|0.210|0.222|\n|Book|0.078|0.025|0.069|\n|CC|0.268|0.052|0.068|\n|C4|0.283|0.025|0.050|\n|Github|0.059|0.021|0.378|\n|Stackexchange|0.230|0.649|0.103|\n|Wikipedia|0.041|0.019|0.110|\n|*Avg PPL* of 124M model|24.97|25.45|26.73|\n\nIn contrast, Chameleon is stable across training steps, model sizes, $\\lambda$, and sample counts (Tables 10, 11). This means **our method can produce promising domain weights without repeated validation**, significantly reducing overall costs for users.\n\nAnother key aspect, as the reviewer pointed out, is the cost of incorporating new data sources. Our data-centric approach requires only inference to obtain new embeddings and recompute KRLS, whereas proxy optimization-based methods like DoReMi and DoGE necessitate full retraining and additional tuning.\n\nLastly, we further validate performance improvement by training 1.2B models. Chameleon demonstrates gains in both perplexity and downstream task accuracy (see *Q3 for Reviewer jzfM*).\n\n[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).\n\n[2] Parmar et al. Data, data everywhere: A guide for pretraining dataset construction. ACL 2024.", "claims_and_evidence": "The claims on speed need further explanations, see the section below", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Some of them", "experimental_designs_or_analyses": "Most of them", "relation_to_broader_scientific_literature": "Builds upon existing literature, primarily DoGE", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths:**\n\n-  The method is simple, intuitive, and easy to implement.\n-  The paper is well written, and the experiments are well organized.\n-  The topic is very relevant and timely.\n\n**Weaknesses:**\n\n- The core idea is just a heuristic (in authors words): pretaining needs \"broadly shared semantic structures\" and finetuning needs \"distinct and unique data\" to \"highlight domain-specific characteristics\".\n  To me the statements above are just some vague justifications for what empirically works.\n- In my opinion the performance improvements are virtually non existent--compared to the model DoGE. The main distinction lies in the speed. Authors might argue that when new data is added, the performance improvement is more tangible. But if we put speed aside, and retrain the baseline proxy networks, then again the only distinction becomes speed. In general the proxy model is relatively small, and its training data is a fraction of the entire training data. How many GPU hours are needed to train the proxy networks across the models? Does increasing speed in training this network make any significant difference in energy consumption? How often \"retraining\" the proxy network is needed at all? Is it needed at all?\n\n I would be happy to revise my score if authors give me convincing answers.\n\n**Other comments:**\n\nPlease don't force the reviewers to read your appendix.", "other_comments_or_suggestions": "See above", "questions_for_authors": "See above", "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": "EsUsPTzfsn", "reviewer": "Reviewer_ubpq", "summary": "This paper introduces a new data mixing framework for language pretraining and finetuning wherein the mixing weights for different domains are constructed from a domain affinity matrix generated via kernel functions on domain embeddings. This domain matrix can naturally be transformed into domain weights for pre-training (i.e., emphasizing broader and different domains) and finetuning (emphasizing similar domans). The key advantage of this framework is that it does not rely extensively on training proxy models, and thereby can be seen as a low-cost alternative to existing data mixing frameworks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. Reasoning behind obtaining domain weights from the KRLS\n\nWe use kernel ridge leverage scores (KRLS) to determine domain weights. KRLS is a well-established tool in data analysis. It quantifies the influence or importance of data points [Alaoui & Mahoney, 2015]. This property is leveraged in machine learning for tasks like density estimation [Pauwels et al., 2018] and novelty detection [Ducharlet et al., 2024; Lasserre & Pauwels, 2019].\n\nThe inverse KRLS is proportional to the Christoffel function value [Pauwels et al., 2018]. This relationship provides additional theoretical justification for our approach. Christoffel functions (Eq. (1) in [Pauwels et al., 2018]) precisely characterize the local density of the data distribution in the feature space, where higher values indicate denser regions. \n\nWe compute the score $S_\\lambda(D_i)$ of domain $i$ using Eq. (KRLS) on page 4. During pretraining, assigning higher sampling probability to domains with low KRLS (and thus high $S_\\lambda^{-1}$/Christoffel value) upweights high-density data regions, which are most influential on base LMs' performance [1]. LLM finetuning aims to specialize on a novel specific task, requiring the model to learn differential features not fully captured during pretraining, so we instead prioritize the domains with high $S_\\lambda$. Section 3.2 converts either $S_\\lambda^{-1}$ or $S_\\lambda$ into probability distributions $\\alpha$ by appliying softmax normalization.\n\nWe will revise Section 3.2 to explicitly link the data mixing goal to KRLS and inverse KRLS, grounding it in statistical learning theory.  We will make this discussion self-contained within the main text, incorporating analysis from Appendix A.\n\n> Q2. Computational cost breakdown\n\nChameleon's main computational cost comes from 1) proxy training and 2) embedding extraction, with proxy training being dominant. In our setting, training an 82M proxy model requires $10^{17}$–$10^{18}$ FLOPs, while DoReMi and DoGE take longer to converge, leading to 5-10x higher costs (see line 299, \"Stability and Practicality\"). Regarding our embedding extraction, it requires only $10^{15}$ FLOPs (<1% of proxy training). Importantly, Chameleon avoids proxy retraining when domains change, incurring only embedding extraction costs. In contrast, DoReMi and DoGE induce their proxy retraining FLOPs.\n\nOur method is also significantly cheaper in GPU hours, see *Response to Q1 of Reviewer jzfM* for more details and complexity analysis.\n\nBeyond efficiency, Chameleon is also more stable, making it resource-efficient in practical use. Unlike DoReMi and DoGE, which are sensitive to hyperparameters, Chameleon remains robust, see *Q2 of Reviewer jX7F* for more details.\n\nLastly, we note that Chameleon shows favorable accuracy behaviour on larger models as well, as shown in our additional 1.2B model experiments (see *Q3 of Reviewer jzfM*).\n\n> Q3. Data Mixing Laws\n\nWe first provide discussions and then present empirical comparisons.\n\nData Mixing Laws derive domain weights by leveraging scaling laws of training steps, model sizes, and data mixtures to predict the performance of large models trained on diverse data from small-scale training. This requires training multiple small proxy models with varying domain weights, making it more computationally expensive than ours, which trains just one proxy model.\n\nWe use their reported domain weights to train a 684M model on Slimpajama. Since their weights are optimized with the Pile as the target, they may be suboptimal for SlimPajama. However, given the alignment of their objectives and overlap in data sources, we consider the comparison meaningful.\n\n**Chameleon outperforms Data Mixing Laws in both perplexity and downstream tasks** at a fraction of the cost. Data Mixing Laws' FLOPS is calculated for 4 different proxy sizes and 20 separate mixtures, where **our cost is 2 orders of magnitude lower**.\n\n**Table ubpq-1:** PPL comparison with Data Mixing Laws\n||Data Mixing Laws|Chameleon|\n|-|:-:|:-:|\n|Arxiv|7.55|8.31|\n|Book|45.06|39.23|\n|CC|44.21|40.11|\n|C4|45.79|42.59|\n|Github|4.01|4.20|\n|Stackexchange|7.96|7.94|\n|Wikipedia|16.20|13.90|\n|Avg PPL|24.40|**22.31**|\n|# Domains Over Uniform|4/7|4/7|\n|FLOPs|$5.36×10^{19}$|$1.36×10^{17}$|\n\n**Table ubpq-2:** Downstream accuracy comparison with Data Mixing Laws\n|Task|Data Mixing Laws|Chameleon|\n|-|:-:|:-:|\n|ARC-E|34.5|37.8|\n|COPA|59.0|61.9|\n|HellaSwag|27.4|27.0|\n|Lambada|14.7|15.1|\n|LogiQA|26.0|22.6|\n|MultiRC|57.2|57.2|\n|OpenBook|25.2|14.4|\n|PiQA|58.5|60.5|\n|QQP|36.8|39.2|\n|RACE|26.4|26.5|\n|SciQ|57.2|64.3|\n|Social IQA|36.1|35.7|\n|WinoGrande|48.4|52.1|\n|Average|39.0|**39.6**|\n\n[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).\n\n[2] Parmar et al. Data, data everywhere: A guide for pretraining dataset construction. ACL 2024.", "claims_and_evidence": "The paper is well-supported by rigorous numerical analysis.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "There are some theoretical results in the Appendix, which I briefly reviewed.", "experimental_designs_or_analyses": "The experiments are meaningful and sound.", "relation_to_broader_scientific_literature": "The paper adds to the data mixing literature, with specific focus on compute-efficiency and adaptation to new tasks.", "essential_references_not_discussed": "The paper comprehensively covers the primary literature.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written and timely, addressing rigorous experiments\n\n\nWeaknesses:\n- The reasoning behind obtaining domain weights from the KRLS is unclear. Is there some theoretical relationship or connection that can be derived to show why pretraining weights as designed, or finetuning weights as designed, are appropriate? I appreciate that the authors have provided some intuition but it would be beneficial to get more insight.\n- There doesn't seem to be a significant performance improvement from the mixing law, and indeed, the major selling point is that it achieves a competitive performance at orders of magnitude lower cost. It would be useful to emphasize this point further, especially in numerical results to better break down all cost calculations (e.g., in domain transfer).\n- It would be useful to include Data Mixing Laws as a baseline, at least for experiments on generalization [1]\n\n[1] Ye, Jiasheng, et al. \"Data mixing laws: Optimizing data mixtures by predicting language modeling performance.\" arXiv preprint arXiv:2403.16952 (2024).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Refer to Strengths & Weaknesses", "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": "vfVtEpSwO0", "reviewer": "Reviewer_jzfM", "summary": "This manuscript introduces a flexible data-mixing framework for LLM that uses kernel ridge leverage scores (KRLS) computed from learned domain embeddings using a proxy model. It quantifies each domain’s representativeness and interdependence within the embedding space, and then uses this information to generate a weighted data mixture for both pretraining and finetuning.  The comprehensive experiments and ablation studies demonstrate its practical benefits on both perplexity and downstream tasks by outperforming existing methods such as DoReMi and DoGE with a small computational overhead.\n\n## Post Rebuttal ##\n\nMost of my concerns are properly addressed. I would like to raise my score.\n\n## Post Rebuttal ##", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. Runtime analysis and complexity comparison\n\nObtaining embeddings $x_i, \\, i=1,\\ldots,k$ requires a single forward pass for each $a \\in B_i$ through the proxy $h_{\\theta_p}(a)$; inference is fast as the proxy is a small model. Computing (KRLS) involves inverting a matrix of size $k \\times k$ in $\\mathcal{O}(k^3)$, which is computationally cheap since datasets typically have a small number $k$ of domains. We do not add any overhead in proxy training. In contrast, DoGE requires per-domain gradient computation at each iteration, and DoReMi runs inference of the reference model for perplexity comparisons. \n\nFor a straightforward comparison, we report GPU hours below. **DoReMi and DoGE incur over 10% of the base model training cost, while Chameleon reduces it to under 2%.** These savings are particularly impactful for academic labs with limited computational resources.\n\n**Table jzfM-1**: Runtime comparison.\n|Method|GPU Hours|\n|-|:-:|\n|DoReMi|7.4|\n|DoGE|6.3|\n|Chameleon|0.8|\n|684M base model|56|\n\nSee more details on the FLOPs computations in *Q2 of Reviewer ubpq*.\n\n> Q2. Theoretical insights for the effectiveness of KRLS in domain reweighting\n\nWe use kernel ridge leverage scores (KRLS) to determine domain weights. KRLS is a well-established tool in data analysis. It quantifies the influence or importance of data points [Alaoui & Mahoney, 2015]. This property is leveraged in machine learning for tasks like density estimation [Pauwels et al., 2018] and novelty detection [Ducharlet et al., 2024; Lasserre & Pauwels, 2019].\n\nThe inverse KRLS is proportional to the Christoffel function value [Pauwels et al., 2018]. This relationship provides additional theoretical justification for our approach. Christoffel functions (Eq. (1) in [Pauwels et al., 2018]) precisely characterize the local density of the data distribution in the feature space, where higher values indicate denser regions. \n\nWe compute the score $S_\\lambda(D_i)$ of domain $i$ using Eq. (KRLS) on page 4. During pretraining, assigning higher sampling probability to domains with low KRLS (and thus high $S_\\lambda^{-1}$/Christoffel value) upweights high-density data regions, which are most influential on base LMs' performance [1]. LLM finetuning aims to specialize on a novel specific task, requiring the model to learn differential features not fully captured during pretraining, so we instead prioritize the domains with high $S_\\lambda$. Section 3.2 converts either $S_\\lambda^{-1}$ or $S_\\lambda$ into probability distributions $\\alpha$ by appliying softmax normalization.\n\nWe will revise Section 3.2 to explicitly link the data mixing goal to KRLS and inverse KRLS, grounding it in statistical learning theory.  We will make this discussion self-contained within the main text, incorporating analysis from Appendix A.\n\n> Q3. More detail on how domain weights transfer across datasets and model scales\n\n- **Transfer across model sizes**: Prior works (e.g., DoReMi, DoGE, RegMix) have shown that domain weights transfer well across model scales. To further validate this, we trained 1.2B models on SlimPajama and found that weights from an 82M proxy model effectively transfer to both 684M and 1.2B models. Notably, Chameleon achieves even greater improvements on larger models, highlighting its scalability.\n\n**Table jzfM-1:** PPL with 1.2B model\n|Domain|Uniform|DoReMi|DoGE|Chameleon|RegMix|\n|-|:-:|:-:|:-:|:-:|:-:|\n|Arxiv|6.30|7.09|7.07|6.33|10.61|\n|Book|28.25|32.66|27.83|24.63|27.55|\n|CC|31.19|29.96|28.11|26.95|24.70|\n|C4|34.74|33.05|31.06|29.58|31.94|\n|Github|2.91|3.03|3.07|2.94|4.08|\n|Stackexchange|6.01|6.44|5.80|5.76|9.54|\n|Wikipedia|8.65|7.93|10.88|9.03|20.08|\n|*Average PPL*|16.86|17.17|16.26|**15.03**|18.36|\n\n**Table jzfM-2:** Downstream accuracy with 1.2B model\n|Task|Uniform|DoReMi|DoGE|Chameleon|RegMix|\n|-|:-:|:-:|:-:|:-:|:-:|\n|ARC-E|39.4|41.2|41.9|42.4|43.0|\n|COPA|64.0|66.0|63.0|61.0|66.0|\n|HellaSwag|27.5|27.7|28.2|28.4|27.6|\n|Lambada|17.9|17.3|18.7|21.6|20.7|\n|LogiQA|22.0|24.0|22.0|21.2|20.7|\n|MultiRC|57.2|57.2|57.2|57.2|56.9|\n|OpenBookQA|15.0|13.6|13.8|16.4|17.4|\n|PiQA|61.5|61.9|61.8|63.8|58.7|\n|QQP|36.8|36.8|36.9|36.9|36.8|\n|RACE|26.0|26.7|27.8|29.1|28.4|\n|SciQ|69.7|68.3|69.0|72.6|72.0|\n|SocialIQA|36.2|36.5|35.9|37.2|36.1|\n|WinoGrande|52.8|49.6|48.9|51.5|50.0|\n|*Average*|40.5|40.5|40.4|**41.5**|41.1|\n\n- **Transfer across datasets**: We conduct ablation studies on the Pile (Section 4.2). Specifically, we retrain a proxy model on the Pile for reference. As shown in [Domain weights on the Pile](https://imgur.com/a/2Xnk56t), the weight from a proxy trained on the Pile (the blue column) aligns with the weights transferred from proxies trained on SlimPajama at various sizes (the other columns), confirming Chameleon’s robust transferability.\n\n[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).", "claims_and_evidence": "Overall, most claims are supported by empirical evidence. However, some of the claims are not supported by clear and convincing evidence.\n\n1. The authors claim that the method is computationally efficient. However, the submission does not include a detailed runtime analysis or complexity comparison with baseline methods. \n\n2. The novelty of the manuscript is the ability to quantify each domain’s representativeness and interdependence in the learned embedding. However, the evidence provided does not clearly demonstrate how reliable or robust these kernel ridge leverage scores are across different scenarios. More analysis or visualization of these scores and their correlation with the performance would strengthen this claim.\n\n3. The authors claim that the computed domain weights transfer directly to new data and different model scales without needing to retrain the proxy model. Although the provided experimental results are promising, the evidence could be more convincing if more experiments or detailed ablation studies were provided to thoroughly evaluate this transferability under a wider range of conditions.\n\n4. The paper would be strengthened by an in-depth theoretical analysis explaining why kernel ridge leverage scores are particularly effective for domain reweighting.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria align well with the problem of efficient data mixing for LLM training. The use of kernel ridge leverage scores computed from learned domain embeddings makes sense for quantifying domain representativeness and interdependence. Measuring improvements in perplexity and downstream task performance are standard metrics for evaluating language model training.", "theoretical_claims": "I have checked the theoretical proof in the appendix. The proof of the equivalence between the ridge leverage scores computed in the feature space and those obtained via kernel ridge regression when using a linear kernel. The proof, presented as Lemma A.1, correctly applies standard matrix identities to show that the diagonal elements of the corresponding hat matrices are equal. While some intermediate steps are missing, I did not find any  logic issues.\n\nBeyond this lemma, most theoretical claims are supported by empirical evidence and ablation studies rather than fully formalized proofs. Overall, the proofs that were provided are correct, and no major issues were found.", "experimental_designs_or_analyses": "Yes, the experimental design and analyses are sound. however, I would suggest the author to include a runtime analysis to validate the claim that the proposed method is computationally efficient.", "relation_to_broader_scientific_literature": "The key contributions of the paper are well situated within the broader scientific literature, building upon and extending several established ideas, which include (1) kernel methods and leverage scores; (2) data selection and mixing in LLM Training; and (3) domain adaptation and representation kearning.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "#### Strengths:\n\n1. The paper presents a novel method that leverages the kernel ridge leverage score to reweight data domains, which is a combination of ideas from data mixing, transfer learning, and efficient training of LLMs.\n\n2. The proposed work holds significance in practice as it is importance in scaling LLM training without retraining expensive proxy models.\n\n3. The extensive empirical results indicate improvements in both pretraining and finetuning, which validates the practical values of the proposed method.\n\n4.  writing is clear and observation-driven, which makes it easy for readers to follow along and engage with the authors to approach the problem.\n\n#### Weaknesses:\n\n1. While the paper claims the computational efficiency, it lacks a detailed runtime analysis or formal complexity comparison with baseline methods.\n\n2. The theoretical justification behind using KRLS is largely supported by empirical evidence rather than formal proofs.\n\n3. The clarity regarding the transferability of domain weights across diverse datasets and model sizes could be improved with more detailed ablation studies and analysis.", "other_comments_or_suggestions": "Please see the Weaknesses section.", "questions_for_authors": "1. Could you provide a more detailed runtime analysis and complexity comparison of your method with baseline methods?\n\n2. Can you elaborate on the theoretical insights behinds why KRLS are effective for domain reweighting? \n\n3. Could you discuss in more detail how the computed domain weights transfer across different datasets and model scales?", "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:

Meta-review: The manuscript introduces a data-mixing framework leveraging kernel-based scores for data curation in both pretraining and finetuning. The proposed method demonstrates promising empirical results, often at a significantly reduced computational cost compared to existing methods like DoReMi and DoGE. 

The review process identified several areas that the authors should address to further strengthen the manuscript for its camera-ready version. A primary concern across multiple reviewers was the need for a more detailed exposition of a principled understanding underpinning the approach, specifically regarding why KRLS is effective for domain reweighting and the intuitive connection between KRLS scores and data mixing strategies. I will not detail additional dimensions for improvement (e.g., better visualizations of KRLS scores or various ablations), but overall incorporating reviewers’ feedback will strengthen the grounding and empirical validity of the work and further clarify the advantages of the proposed method.
Decision: Accept (poster)