PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wanyun Xie, Francesco Tonin, Volkan Cevher

arXiv:2505.24844v1cs.LGcs.CL

TL;DR

Choosing domain weights is difficult, and existing approaches can be costly or require proxy retraining when data changes. CHAMELEON derives weights from domain embeddings with leverage scores, matching established pretraining methods at lower cost, adapting as domains change, and improving finetuning perplexity across domains.

  • Problem

    Choosing domain weights is challenging because data quality, diversity, inter-domain overlap, and task complexity matter, while existing methods often address only pretraining and domain changes.

    Source: §1 · p. 1

  • Method

    CHAMELEON constructs domain relationships from learned embeddings and uses KRLS to quantify each domain's contribution to the embedding space and guide reweighting.

    Source: §1 · p. 1

  • Results

    CHAMELEON matches DoReMi and DoGE in pretraining at a fraction of their cost, outperforms baselines at 1% of retraining cost as domains double, and improves perplexity consistently across finetuning domains.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework supports direct application to new data without rerunning the proxy pipeline and extends data reweighting across pretraining and finetuning.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1. Introduction

Domain reweighting matters for LLM generalization, but choosing weights is difficult and existing proxy-based approaches are costly. CHAMELEON uses embeddings and KRLS to derive domain weights, adapt to changing data, and extend reweighting to finetuning.

1. Introduction

  • Existing proxy-model methods for deriving domain weights are computationally expensive, and frequent recalibration becomes impractical as data evolves.Weight selection is also challenging because data quality, diversity, inter-domain overlap, and task complexity vary.Source: §1 · p. 1
  • CHAMELEON derives weights from learned embeddings using KRLS, which measures each domain's contribution to the overall embedding space.The approach constructs a domain affinity matrix and computes leverage scores to guide reweighting.Source: §1 · p. 1
  • CHAMELEON matches DoReMi and DoGE in pretraining performance at a fraction of their cost.The framework uses KRLS to quantify domain representativeness from embedded data.Source: §1 · p. 2
  • It outperforms baselines at 1% of the retraining cost, even as domains double.It adapts to changes in domain composition without retraining proxy models.Source: §1 · p. 2
  • Finetuning weights emphasize domain-specific uniqueness, with consistent perplexity improvements across all domains on natural language and code datasets.This extends domain reweighting beyond pretraining.Source: §1 · p. 2
  • KRLS scores are described as inexpensive, hyperparameter-robust, and quick to converge, enabling direct application to new data without rerunning the proxy pipeline.The framework is presented as a unified approach to data mixing across pretraining and finetuning.Source: §1 · p. 2

2. Related Work

Related work includes proxy-based reweighting, offline mixture search, and scaling-law approaches. The paper distinguishes these approaches by their computational costs, reliance on downstream tasks or scaling strategies, and ability to adapt to domain changes.

2. Related Work

  • Online adaptation strategies for domain reweighting require frequent recalibration and monitoring.Source: §2 · p. 2
  • DoReMi uses a reference and proxy model with Group DRO, while DoGE tracks domain-specific gradients during proxy training.The paper describes both approaches as inefficient; DoReMi depends on reference-model quality and trains two models, while DoGE incurs gradient-tracking costs.Source: §2 · p. 2
  • Offline methods such as Data Mixing Laws and RegMix search for mixtures using multiple proxy models; unlike RegMix, CHAMELEON does not require downstream-task access.These methods reveal that domain weights transfer across model sizes.Source: §2 · p. 2
  • Data- and model-scaling-law studies inform domain weighting, but rely on the scaling strategy and lack an easy way to adapt to domain expansion.Source: §2 · p. 2

3. Data Mixing via CHAMELEON

CHAMELEON treats domain mixing as a data-centric problem, using proxy-model embeddings to capture domain characteristics and relationships. It computes domain affinity and KRLS scores to guide weights, with different KRLS uses described for pretraining and finetuning.

3. Data Mixing via CHAMELEON

  • The framework follows a two-stage setup: a small proxy model infers domain weights, then a large base model trains using those weights.Source: §3 · p. 2
  • Figure 1's pipeline averages embeddings within domains, computes pairwise-inner-product affinities, applies KRLS, and softmax-normalizes scores into resampling weights.Source: Figure 1 · p. 3
  • The method uses λ during pretraining to promote general knowledge learning and KRLS Sλ during finetuning to emphasize task-specific knowledge.Source: §3 · p. 3
  • CHAMELEON focuses on intrinsic data properties rather than deriving weights from a proxy model's optimization process.Source: §3 · p. 3
  • Proxy-model hidden-layer embeddings represent domains and capture inter-domain relationships; their visualization highlights semantic distinctiveness and the central coverage of broad domains.The example describes CC and C4 as occupying dominant shared semantic regions and Arxiv as more distinct.Source: §3 · p. 3

2. How can such properties inform domain reweighting?

CHAMELEON builds domain embeddings and uses their affinity matrix to compute KRLS-based weights. It applies different weighting directions in pretraining and fine-tuning, and can score new domains without retraining the proxy.

3.1. Quantifying domain characteristics

  • CHAMELEON averages LM embeddings into domain vectors, forms a linear-kernel affinity matrix, and applies KRLS to quantify domain influence.The affinity matrix captures pairwise semantic similarity but does not itself measure each domain's importance.Source: §3.1 · pp. 3–4
  • High KRLS indicates a domain is relatively unique and difficult to approximate from other embeddings; low KRLS indicates broadly shared characteristics.KRLS values are diagonal entries of the KRR hat matrix and depend on domain inputs and regularization, not target values.Source: §3.1 · p. 4

3.2. Incorporating KRLS into LM training

  • Pretraining uses inverse KRLS to upweight high-density regions, while fine-tuning prioritizes high-KRLS domains with more distinct characteristics.The section ties pretraining to general, dense regions and fine-tuning to learning task-specific features.Source: §3.2 · pp. 4–5, §3.1 · p. 4
  • New domain embeddings can be scored directly without proxy retraining, unlike methods that require retraining when domains are added.The proxy supplies embeddings for new data, from which the new KRLS scores are calculated.Source: §3.2 · p. 4
  • Computing domain embeddings requires one proxy forward pass per sample, and KRLS computation costs O(k^3) for k domains.The method adds no overhead to proxy training.Source: §3.2 · p. 5

4. Experiments

Across pretraining, transfer to new domains, and fine-tuning, CHAMELEON reports competitive or improved performance with low additional computation. Its experiments also test whether weights transfer across model scales and whether the method remains stable across settings.

4. Experiments

  • The experiments evaluate universal generalization on SlimPajama, transfer to Pile without proxy retraining, and fine-tuning on language and programming datasets.The paper frames these as tests of performance, scalability, and applicability to fine-tuning.Source: §4 · p. 5, §4.1 · p. 5, §4.2 · p. 8

4.1. Universal Generalization

  • CHAMELEON outperforms Uniform and DoReMi, matches DoGE on per-domain perplexity, and uses 10% of DoReMi’s and 20% of DoGE’s total FLOPs.The comparison is for 684M models; RegMix uses target downstream-task knowledge, unlike the other methods.Source: §4.1 · pp. 5–6, Table 2 · p. 6
  • CHAMELEON is competitive across downstream tasks, and weights from an 82M proxy transfer to 684M and 1.2B models, with larger improvements observed on larger models.The method’s weights are also reported stable across proxy training steps and robust to variations in proxy size, λ, and embedding sample count.Source: §4.1 · pp. 6–7

4.2. Scalable to Pile

  • On Pile, CHAMELEON outperforms DoReMi, DoGE, and RegMix in perplexity and downstream accuracy while using 1% of the cost of retraining a proxy model.The method reuses the SlimPajama-trained proxy and computes embeddings for 4k samples per Pile domain; the reported transfer evaluation uses available Pile domains.Source: §4.2 · p. 8, Table 7 · p. 8

4.3. Finetuning

  • On Wiki40b and Stack fine-tuning, CHAMELEON improves over uniform weights across all evaluated domains, using forward passes to obtain domain embeddings before applying KRLS.The experiments use seven Latin languages from Wiki40b and seven programming languages from Stack-decup; KRLS-based fine-tuning weights outperform using inverse pretraining weights.Source: §4.3 · p. 9, Table 10 · p. 9

5. Conclusion

The conclusion presents CHAMELEON as an efficient KRLS-based framework for domain reweighting across pretraining and fine-tuning. It also identifies online optimization and task-specific weighting as directions for future work.

  • CHAMELEON uses KRLS to quantify domain representativeness, adapts to new domains without proxy retraining, and improves fine-tuning results.The conclusion characterizes the approach as computationally inexpensive and stable.Source: §5 · p. 9
  • Future work will extend CHAMELEON to online dynamic optimization and modify the KRLS identity matrix to emphasize domains for specific downstream tasks.Source: §5 · p. 9

Impact Statement

The impact statement identifies no additional societal consequences, while the leverage-score appendix explains how KRLS measures domain importance and how scores guide mixture weights.

Impact Statement

  • The authors identify no societal consequences requiring specific discussion beyond the general aim of advancing machine learning.This is the stated scope of the Impact Statement.Source: Impact Statement · p. 10

A.1. Details of leverage scores

  • KRLS assigns domain-importance scores from embeddings, with high scores marking domains difficult to approximate from others and low scores marking shared characteristics.The appendix relates KRLS to ridge leverage scores computed in the embedding feature space.Source: §A.1 · pp. 14–15
  • For pretraining, the method ranks lower-KRLS domains as more important, while finetuning upweights higher-KRLS domains for task-specific representations.The distinction follows the appendix’s contrast between common characteristics useful for general-purpose representations and unique characteristics suited to task-specific ones.Source: §A.1 · p. 15

A.2. Discussions

The discussion describes regularization and cross-model stability, connects KRLS to Christoffel functions, and illustrates clustered domain embeddings.

A.2. Discussions

  • Regularization suppresses small kernel-matrix components; the authors report that computed weights are robust to λ, with larger values not significantly changing them.A large λ emphasizes dominant kernel-matrix eigenvectors, while small regularization is sufficient in practice.Source: §A.2 · p. 15
  • The method’s domain weights are reported to remain robust across model sizes because pairwise embedding similarities—and thus KRLS scores—are more consistent than absolute embeddings.The discussion gives semantic proximity between GitHub and StackExchange as an example and points to Table 16 for empirical evidence.Source: §A.2 · p. 15
  • The Christoffel function is proportional to inverse ridge leverage scores: high values correspond to denser, more representable regions, and the authors link this focus to improved generalization and transfer.The regularized kernelized Christoffel function at a sample is inversely proportional to its KRLS when η = 1.Source: §A.2 · pp. 15–16

B.1. Additional figures for Figure 1

  • The embedding visualization shows semantically similar SlimPajama domains occupying nearby regions and forming high-density clusters.Source: Figure 4 · p. 16

B.2. Experimental setup.

The appendix specifies model and training settings for the experiments, the temperature ranges used to normalize weights, and the proxy-training budgets for SlimPajama domain weights.

B.2. Experimental setup.

  • The experiments use an 82M auxiliary decoder-only Transformer and a 684M pretraining base model, with batch size 128 for SlimPajama, Wiki40b, and Stack and 512 for the noisier, more domain-rich Pile.The reported common settings include batch size 128, cosine learning-rate scheduling, weight decay 0.01, and gradient clipping 1.0.Source: §B.2 · p. 16
  • CHAMELEON applies a temperature factor in softmax normalization, typically using τ_PT ∈ [5, 10] for pretraining and τ_FT ∈ [0.2, 0.5] for finetuning.Source: §B.2 · p. 16
  • The CHAMELEON ablation evaluates proxy models at 60M, 124M, and 210M; SlimPajama RegMix trains 200 1M proxy models for 1k steps to fit a regression model.Source: §B.2 · p. 16

B.3. Domain weights on Slimpajama

  • For SlimPajama weights, DoReMi and DoGE use 10k-step proxy training, CHAMELEON uses 2k steps, and RegMix uses 200 1M proxy models with “CC” as the target domain.The final domain weights are reported in Table 12; its caption provides no weight values in the supplied passage.Source: §B.3 · p. 17, Table 12 · p. 17

B.4. Temporal cost

CHAMELEON reduces the reported cost of obtaining domain weights and is described as converging faster and remaining stable across several settings.

B.4. Temporal cost

  • Less than 2% of final training cost is added by CHAMELEON’s weight computation, versus over 10% for DoReMi and DoGE.The comparison is reported for the experiments in Table 2; the appendix notes that baseline cost estimates may be optimistic lower bounds because of hyperparameter tuning.Source: Table 13 · p. 17, §B.4 · p. 17

B.5. Stable domain weights of CHAMELEON

  • CHAMELEON’s domain weights are reported to be more stable and converge faster than DoReMi’s and DoGE’s; the 2k-step proxy comparison also shows the latter methods converge more slowly.Tables 14 and 15 report the domain weights at different steps and perplexity using weights from 2k-step proxies.Source: §B.5 · p. 17, Table 14 · p. 18, Table 15 · p. 18
  • CHAMELEON’s domain weights are reported as stable across model sizes, λ values, and the number of embedding-computation samples.Table 16 reports the weights across these settings.Source: §B.5 · p. 18, Table 16 · p. 18

B.6. Additional baseline: Data Mixing Laws

The appendix compares Chameleon with Data Mixing Laws, reports domain weights for the Pile and finetuning datasets, and compares finetuning perplexity for KRLS-based weights and their inverse.

B.6. Additional baseline: Data Mixing Laws

  • CHAMELEON outperforms Data Mixing Laws in perplexity and downstream tasks at 2 orders of magnitude lower cost.Data Mixing Laws derives weights using multiple small proxy models; the comparison applies its Pile-optimized weights to Slimpajama, where they may be suboptimal.Source: §B.6 · pp. 18–19, Table 17 · p. 19, Table 18 · p. 19

B.7. Domain weights on Pile

B.8. Domain weights for finetuning

B.9. PPL of finetuning.

  • Fine-tuning with KRLS-based domain weights performs better than using the inverse of those weights in the reported perplexity comparisons.Tables 22 and 23 compare per-domain perplexity with the uniform baseline for 684M-parameter models on seven Wiki40b and seven Stack languages.Source: §B.9 · p. 19, Table 22 · p. 20, Table 23 · p. 21
Loading 2505.24844v1…