Research map / Label-efficient and robust learning / Domain adaptation and generalization
Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Domain2Vec, which decomposes any dataset into a distribution over 260 meta-domains via a fine-tuned Qwen2-1.5B classifier, and uses these domain vectors for training-free data mixture search (DA^2) or as inputs to prior fitting methods like RegMix.
Problem
Finding optimal data mixtures for LLM pretraining is computationally expensive (proxy training, e.g., DoReMi at 3.7e19 FLOPs) and not scalable when the dataset composition changes, since fitted mixture-performance functions must be refit.
Approach
Represent any dataset as a normalized vector over a vocabulary of k-means meta-domains (classifier-averaged over sampled documents), then assume lower validation loss when training and validation domain vectors align (DA^2), and fit per-meta-domain loss functions in this vector space so they transfer to unseen datasets without retraining.
Key results
- Domain2Vec+DA^2 matches Human's Pile-CC validation loss using 55.38% of training steps and Domain2Vec+RegMix using 51.50%, with RegMix-based version improving downstream average by 2.83% under equal compute.
- On ranking prediction across validation sets, Domain2Vec+DA^2 reaches Spearman 0.6657 vs 0.3543 for kNN and 0.0497 for random.
- Domain2Vec-based mixtures reach performance comparable to DoReMi at about 0.26% of its estimated FLOPs (9.66e16 vs 3.7e19).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Domain2Vec+DA² to rank preset data mixtures instead of random ranking | Data-mixture ranking across validation datasets · Validation-set ranking | Spearman correlation: 0.0497 -> 0.6657 | Helped |
| use Domain2Vec+DA² to rank preset data mixtures instead of kNN | Data-mixture ranking across validation datasets · Validation-set ranking | Spearman correlation: 0.3543 -> 0.6657 | Helped |
| use Domain2Vec+DA² to rank preset data mixtures instead of Domain2Vec+RegMix | Data-mixture ranking across validation datasets · Validation-set ranking | Spearman correlation: 0.4629 -> 0.6657 | Helped |
| use Domain2Vec+DA² rather than Human's original data mixture | LLaMA-like models; average over 12 downstream tasks · Downstream tasks | Average Performance: 0.424 -> 0.432; improves downstream performance by an average of $1.89\%$ | Helped |
| use Domain2Vec+RegMix rather than Human's original data mixture | LLaMA-like models; average over 12 downstream tasks · Downstream tasks | Average Performance: 0.424 -> 0.436; improves downstream performance by an average of $2.83\%$ | Helped |
| replace Huber Loss with JS Distance (ablation) | Domain2Vec+DA² · Validation-set ranking | Spearman correlation: 0.6657 -> 0.5000 | Hurt |
| replace Huber Loss with L1 Distance (ablation) | Domain2Vec+DA² · Validation-set ranking | Spearman correlation: 0.6657 -> 0.5400 | Hurt |
| replace Huber Loss with L2 Distance (ablation) | Domain2Vec+DA² · Validation-set ranking | Spearman correlation: 0.6657 -> 0.6429 | Hurt |
Surprising findings
- Pile-CC-only training gives a 3.54% average accuracy gain over the original Pile mixture, making it a strong baseline when Pile-CC is the validation target.
- Domain2Vec+DA^2 at the 1B scale averages 0.432 on downstream tasks, below DoReMi (0.439) and RegMix (0.441) in the main table.
Limitations
- Meta-domain classifier reaches only 74.73% accuracy on its test set.
- Validation-set optimization risks overfitting to the chosen guide dataset; downstream gains are modest and sometimes below baselines.
- Only C4 and Knowledge Pile are used as training sources in the ranking study, and experiments are limited to English-focused Pile components and small models (up to 1B).
Experimental setup
Datasets: C4 Knowledge Pile The Pile RedPajama Pile-CC bge embeddings corpus (~5.2TB, 100+ sources)
Benchmarks: Social IQA HellaSwag PiQA OpenBookQA Lambada SciQ ARC Easy COPA RACE LogiQA WinoGrande MultiRC
Metrics: Spearman correlation Pearson correlation validation loss average accuracy estimated FLOPs
Baselines: DoReMi RegMix Human (original Pile mixture) Pile-CC Only kNN
Builds on: RegMix DoReMi Data Mixing Laws
Closest papers by meaning
- Scaling Laws for Optimal Data Mixtures (NeurIPS 2025 · similarity 0.89)
- Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning (ICML 2025 · similarity 0.87)
- Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 2025 · similarity 0.87)
- GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining (NeurIPS 2025 · similarity 0.87)
- RegMix: Data Mixture as Regression for Language Model Pre-training (ICLR 2025 · similarity 0.86)
- DataMan: Data Manager for Pre-training Large Language Models (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.