Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning

We thank the reviewer for their valuable feedback and address all remaining concerns below:

Q1. Runtime analysis and complexity comparison

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

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

Table jzfM-1: Runtime comparison.

See more details on the FLOPs computations in Q2 of Reviewer ubpq.

Q2. Theoretical insights for the effectiveness of KRLS in domain reweighting

We 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].

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

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.

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

Q3. More detail on how domain weights transfer across datasets and model scales

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

Table jzfM-1: PPL with 1.2B model

Table jzfM-2: Downstream accuracy with 1.2B model

• 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, 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.

[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).

We thank the reviewer for their valuable feedback and address all remaining concerns below:

Q1. Reasoning behind obtaining domain weights from the KRLS

We 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].

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

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.

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

Q2. Computational cost breakdown

Chameleon'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.

Our method is also significantly cheaper in GPU hours, see Response to Q1 of Reviewer jzfM for more details and complexity analysis.

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

Lastly, 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).

Q3. Data Mixing Laws

We first provide discussions and then present empirical comparisons.

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

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

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.

Table ubpq-1: PPL comparison with Data Mixing Laws

Table ubpq-2: Downstream accuracy comparison with Data Mixing Laws

[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).

[2] Parmar et al. Data, data everywhere: A guide for pretraining dataset construction. ACL 2024.

We thank the reviewer for their valuable feedback and address all remaining concerns below:

Q1. Theoretical motivation

We 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].

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

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

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

Q2. Impact of our computational efficiency

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

Table jX7F-1: GPU hours for universal generalization experiments.

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

Table jX7F-2: DoGE's weights are highly sensitive to $\mu$.

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

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

Lastly, 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).

[1] Mallen et al. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. ACL (2023).

[2] Parmar et al. Data, data everywhere: A guide for pretraining dataset construction. ACL 2024.