# PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts

arXiv: https://arxiv.org/abs/2502.06244 · HTML: https://mlsanity.com/paper/2502.06244/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 59 of its 76 references that are in the map.
- First posted on arXiv 10 Feb 2025.

## Builds on (most cited first, 30 of 59)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge](https://mlsanity.com/paper/1811.00937/) (2018)
- [mT5: A massively multilingual pre-trained text-to-text transformer](https://mlsanity.com/paper/2010.11934/) (2020)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Deduplicating Training Data Makes Language Models Better](https://mlsanity.com/paper/2107.06499/) (2021)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [ByT5: Towards a token-free future with pre-trained byte-to-byte models](https://mlsanity.com/paper/2105.13626/) (2021)
- [Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research](https://mlsanity.com/paper/2402.00159/) (2024)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Measuring the Intrinsic Dimension of Objective Landscapes](https://mlsanity.com/paper/1804.08838/) (2018)

## Built on by (most cited first, 8 of 8)

- [More Data or Better Data? A Critical Analysis of Data Selection and Synthesis for Mathematical Reasoning](https://mlsanity.com/paper/2510.07169/) (2025)
- [ADAPT: Learning Task Mixtures for Budget-Constrained Instruction Tuning](https://mlsanity.com/paper/2512.04555/) (2025)
- [AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining](https://mlsanity.com/paper/2505.23878/) (2025)
- [Explaining Data Mixing Scaling Laws](https://mlsanity.com/paper/2606.08167/) (2026)
- [Sampling and Loss Weights in Multi-Domain Training](https://mlsanity.com/paper/2511.06913/) (2025)
- [Harnessing the Potential of Optimizing Data Mixtures via Bayesian Domain Reweighting](https://mlsanity.com/paper/2607.27928/) (2026)
- [Scaling Near-Optimal SFT-RL Annotation Budget Allocation from Small to Large LLMs](https://mlsanity.com/paper/2609.01573/) (2026)
- [Paired Multimodal Scaling Laws](https://mlsanity.com/paper/2609.36263/) (2026)

## Abstract (arXiv)

Modern foundation models are trained on diverse datasets to enhance generalization across tasks and domains A central challenge in this process is determining how to effectively mix and sample data from multiple sources This naturally leads to a multitask learning (MTL) perspective While prior work in MTL has emphasized mitigating gradient conflicts we observe that largescale pretraining scenariossuch as multilingual or multidomain trainingoften exhibit little to no gradient conflict Motivated by this observation we propose PiKE (Positive gradient interaction-based K-task weights Estimator) an adaptive data mixing algorithm that dynamically adjusts sampling weights during training PiKE exploits nonconflicting gradient interactions to minimize a neartight upper bound on the average loss decrease at each step while incurring negligible computational overhead We provide theoretical convergence guarantees and show that PiKE outperforms static and nonadaptive mixing baselines Furthermore we extend PiKE to promote balanced learning across tasks Extensive experiments on largescale language model pretraining confirm that PiKE achieves faster convergence and improved downstream performance compared to existing approaches