# Dynamic Gradient Alignment for Online Data Mixing

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 33 references that are in the map.
- First posted on arXiv 3 Oct 2024.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [A Survey on Hallucination in Large Language Models: Principles, Taxonomy, Challenges, and Open Questions](https://mlsanity.com/paper/2311.05232/) (2023)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [Weak-to-Strong Generalization: Eliciting Strong Capabilities With Weak Supervision](https://mlsanity.com/paper/2312.09390/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [Approximation Methods for Bilevel Programming](https://mlsanity.com/paper/1802.02246/) (2018)
- [A Pretrainer's Guide to Training Data: Measuring the Effects of Data Age, Domain Coverage, Quality, & Toxicity](https://mlsanity.com/paper/2305.13169/) (2023)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Unsupervised Domain Clusters in Pretrained Language Models](https://mlsanity.com/paper/2004.02105/) (2020)
- [A framework for bilevel optimization that enables stochastic and global variance reduction algorithms](https://mlsanity.com/paper/2201.13409/) (2022)
- [Amortized Implicit Differentiation for Stochastic Bilevel Optimization](https://mlsanity.com/paper/2111.14580/) (2021)
- [Optimizing Data Usage via Differentiable Rewards](https://mlsanity.com/paper/1911.10088/) (2019)
- [Adaptive Training Distributions with Scalable Online Bilevel Optimization](https://mlsanity.com/paper/2311.11973/) (2023)
- [Need a Small Specialized Language Model? Plan Early!](https://mlsanity.com/paper/2402.01093/) (2024)

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

- [Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning](https://mlsanity.com/paper/2505.24844/) (2025)
- [Diversity as a Reward: Fine-Tuning LLMs on a Mixture of Domain-Undetermined Data](https://mlsanity.com/paper/2502.04380/) (2025)
- [Fast Data Attribution for Text-to-Image Models](https://mlsanity.com/paper/2511.10721/) (2025)
- [DIDS: Domain Impact-aware Data Sampling for Large Language Model Training](https://mlsanity.com/paper/2504.13227/) (2025)
- [GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining](https://mlsanity.com/paper/2505.20380/) (2025)
- [DynaMiCS: Fine-tuning LLMs with Performance Constraints using Dynamic Mixtures](https://mlsanity.com/paper/2605.10770/) (2026)
- [Soup-of-Experts: Pretraining Specialist Models via Parameters Averaging](https://mlsanity.com/paper/2502.01804/) (2025)
- [R&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training](https://mlsanity.com/paper/2505.00358/) (2025)
- [Can Small Training Runs Reliably Guide Data Curation? Rethinking Proxy-Model Practice](https://mlsanity.com/paper/2512.24503/) (2025)
- [MaD-Mix: Multi-Modal Data Mixtures via Latent Space Coupling for Vision-Language Model Training](https://mlsanity.com/paper/2602.07790/) (2026)

## Abstract (arXiv)

The composition of training data mixtures is critical for effectively training large language models (LLMs), as it directly impacts their performance on downstream tasks. Our goal is to identify an optimal data mixture to specialize an LLM for a specific task with access to only a few examples. Traditional approaches to this problem include ad-hoc reweighting methods, importance sampling, and gradient alignment techniques. This paper focuses on gradient alignment and introduces Dynamic Gradient Alignment (DGA), a scalable online gradient alignment algorithm. DGA dynamically estimates the pre-training data mixture on which the models' gradients align as well as possible with those of the model on the specific task. DGA is the first gradient alignment approach that incurs minimal overhead compared to standard pre-training and outputs a competitive model, eliminating the need for retraining the model. Experimentally, we demonstrate significant improvements over importance sampling in two key scenarios: (i) when the pre-training set is small and importance sampling overfits due to limited data; and (ii) when there is insufficient specialized data, trapping importance sampling on narrow pockets of data. Our findings underscore the effectiveness of gradient alignment methods in optimizing training data mixtures, particularly in data-constrained environments, and offer a practical solution for enhancing LLM performance on specific tasks with limited data availability.