Research map / Optimization / Optimizers
PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts
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 PiKE, an adaptive data-mixing method for multi-task LLM pretraining that exploits the near-absence of gradient conflicts to update per-task sampling weights, plus a fairness variant Balanced-PiKE.
Problem
Choosing and adapting how to mix and sample data from multiple domains or languages during large-scale LLM pretraining, where static mixtures are suboptimal and existing MTL gradient-conflict methods are too costly at scale.
Approach
Observes that large-scale LM pretraining task gradients are rarely conflicting, derives a bound on per-step loss decrease under the Mix batching strategy, and solves for weights via a cheap periodic mirror-descent update using per-task gradient norms and variances; a tilted-ERM variant (Balanced-PiKE) promotes fairness across tasks.
Key results
- On 750M GLaM pretraining, PiKE improves average downstream accuracy by 3.4% over DoReMi, 6.2% over GLaM, 7.1% over FAMO, and 4.8% over ADO.
- On 1B multilingual pretraining (C4 en/hi/de), PiKE reports 28.7% average HellaSwag accuracy versus 27.5% for Mix, with 1.9x faster convergence to baseline accuracy claimed in the introduction.
- PiKE's overhead is about 1.2% to 2.4% of total training time across model sizes, using per-task gradient estimates every 1,000 steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Adaptively update task sampling weights with PiKE instead of using DoReMi weights | GPT-2 750M pretrained on GLaM · GLaM Q/A tasks | average accuracy (%): +3.4% over DoReMi | Helped |
| Use PiKE instead of static Mix sampling | GPT-2 1B pretrained on mC4 en/hi/de · mC4 en/hi/de validation | average accuracy (%): 27.5 -> 28.7 | Helped |
| Initialize PiKE with GLaM-tuned weights instead of uniform weights (ablation) | PiKE on GPT-2 750M pretrained on GLaM · GLaM Q/A tasks | average accuracy (%): 47.6 -> 48.1 | Helped |
| Use Balanced-PiKE with τ=3 rather than τ=1 (ablation) | GPT-2 1B pretrained on mC4 en/hi/de · mC4 en/hi/de validation | average accuracy (%): 27.9 -> 28.9 | Helped |
| Use Balanced-PiKE with τ=3 rather than τ=1 (ablation) | GPT-2 1B pretrained on mC4 en/hi/de · mC4 en/hi/de validation | average perplexity: 8.29 -> 8.18 | Helped |
| Increase Balanced-PiKE fairness parameter from τ=1 to τ=5 (ablation) | GPT-2 750M pretrained on GLaM · GLaM Q/A tasks | average accuracy (%): 47.5 -> 48.2 | Helped |
| Use Balanced-PiKE with τ=3 rather than τ=1 (ablation) | GPT-2 750M pretrained on GLaM · GLaM Q/A tasks | average accuracy (%): 47.5 -> 47.2 | Hurt |
| Remove the gradient-variance term from PiKE (ablation) | PiKE multi-task pretraining · multi-task performance | balanced performance | Hurt |
Surprising findings
- Validation perplexity does not always track downstream accuracy: Mix attains lower perplexity on 750M GLaM models while PiKE gets higher downstream accuracy.
- PiKE attains the highest perplexity on several C4 subsets (e.g., C4 de 13.87 vs 9.29 for Mix) despite higher average downstream accuracy.
Limitations
- Sampling weights are insensitive to dataset size, as acknowledged in the conclusion.
- Gains are mixed across individual tasks and perplexities; PiKE is worse on some C4 subsets and HellaSwag (de).
- Theory relies on assumptions (L-smooth, bounded variance, low conflict constants) and pretraining experiments are limited to small GPT-2-style models up to 1B.
Experimental setup
Datasets: mC4 (en, hi, de) GLaM six-domain dataset
Benchmarks: HellaSwag (multilingual) ARC-Easy CommonsenseQA PIQA HellaSwag
Metrics: perplexity accuracy
Baselines: Mix (uniform) Round-Robin Random FAMO ADO GLaM DoReMi PCGrad AdaTask
Builds on: FAMO DoReMi GLaM data mixture PCGrad tilted empirical risk minimization
Closest papers by meaning
- Towards Consistent Multi-Task Learning: Unlocking the Potential of Task-Specific Parameters (CVPR 2025 · similarity 0.87)
- Delve into the Applicability of Advanced Optimizers for Multi-Task Learning (NeurIPS 2026 · similarity 0.86)
- GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining (NeurIPS 2025 · similarity 0.86)
- Boosting Multi-Domain Fine-Tuning of Large Language Models through Evolving Interactions between Samples (ICML 2025 · similarity 0.86)
- AdaMerging: Adaptive Model Merging for Multi-Task Learning (ICLR 2024 · similarity 0.86)
- Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 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.