🤖 machine learning

Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

This paper introduces Data Mixing Agent, a model-based framework that employs reinforcement learning to automatically learn generalizable heuristics for re-weighting source and target domain data in continual pre-training, thereby effectively balancing performance across domains while mitigating catastrophic forgetting without requiring manual intervention or retraining.

Original authors: Kailai Yang, Xiao Liu, Lei Ji, Hao Li, Xiao Liang, Zhiwei Liu, Yeyun Gong, Peng Cheng, Mao Yang

Published 2026-04-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kailai Yang, Xiao Liu, Lei Ji, Hao Li, Xiao Liang, Zhiwei Liu, Yeyun Gong, Peng Cheng, Mao Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, well-read librarian (the Large Language Model) who knows everything about history, science, and literature. This librarian is your "Base Model."

Now, you want this librarian to become an expert in Math or Coding. You have a huge stack of new math textbooks and coding manuals (the Target Data), but you also want to make sure the librarian doesn't forget how to write a poem or explain history (the Source Data).

If you just force the librarian to read only the new math books for a month, they might get so obsessed with math that they forget how to speak English properly. This is called "Catastrophic Forgetting."

The Old Way: Guessing the Recipe

Previously, researchers tried to solve this by manually mixing the old books and new books together. They would say, "Okay, let's read 70% math and 30% history." Or maybe "Let's start with 50/50 and slowly shift to 90% math."

But this was like trying to bake a perfect cake by guessing the amount of sugar. Sometimes it worked, sometimes it was a disaster. It relied on human intuition and trial-and-error, which is slow and often inefficient.

The New Way: The "Data Mixing Agent"

This paper introduces a new character: The Data Mixing Agent.

Think of the Agent as a super-smart sous-chef who has watched thousands of cooking videos. Instead of guessing the recipe, this sous-chef has learned the rules of cooking through experience.

Here is how it works, step-by-step:

1. The Training Camp (Reinforcement Learning)

Before the Agent helps your librarian, it goes to a training camp.

The Simulation: Researchers create thousands of "what-if" scenarios. They simulate training a tiny version of the librarian on different mixes of books (e.g., 10% math/90% history, then 20% math/80% history, etc.).
The Feedback Loop: After each simulated mix, they check: "Did the tiny librarian get better at math? Did they forget history?"
The Learning: The Agent watches these thousands of scenarios. It learns patterns like: "Hey, if we add too much math too fast, the librarian gets confused. But if we start slow, then ramp up, then add a little history back at the end, they learn both perfectly."

The Agent isn't just memorizing one recipe; it's learning the general principles of how to mix data to get the best results.

2. The Real Job (Continual Pre-training)

Now, the Agent is ready to help your real librarian.

The Process: As the librarian starts reading the new math books, the Agent watches closely.
The Decision: Every few steps, the Agent asks: "What should the mix be right now?"
The Action: Based on what it learned in the training camp, it might say, "Okay, today we read 60% math and 40% history," and then next week, "Let's switch to 70% math."

It does this dynamically. It's not a static recipe; it's a living strategy that adapts to how the librarian is learning in real-time.

Why is this a Big Deal?

1. It's a "Universal Translator" for Data
The coolest part is that this Agent learned the rules of mixing data in the Math field. When the researchers asked it to help with Coding (a totally different subject), it didn't need to be retrained!

Analogy: Imagine a chef who learned how to bake the perfect cake. You then ask them to bake the perfect bread. They don't need to start from scratch; they just apply the same principles of heat, timing, and ingredient balance. The Agent did the same thing: it took the "Math" lessons and applied them to "Coding" instantly.

2. It Saves Time and Money
Because the Agent knows exactly how much "old data" (history) is needed to keep the librarian from forgetting, it doesn't waste time reading unnecessary books.

Analogy: A human might say, "Let's read 100 books to be safe." The Agent says, "Actually, you only need 80 books to get the same result, and here is the exact order to read them." This saves massive amounts of computer power (GPU hours).

3. It Aligns with Human Intuition (but better)
The researchers found that the Agent's decisions matched what humans thought would work (e.g., "Science books help with general knowledge"), but the Agent found better ways to do it that humans hadn't thought of. It's like a GPS that knows the traffic patterns better than a local driver.

The Bottom Line

Data Mixing Agent is like hiring a master conductor for an orchestra.

The Orchestra: Your AI model.
The Instruments: Different types of data (Math, History, Code).
The Conductor: The Agent.

Instead of letting the musicians play randomly or following a rigid sheet music written by a human, the Conductor listens to the music as it's being played and adjusts the volume of each instrument in real-time to ensure the final song sounds perfect. It prevents the "Math" section from drowning out the "History" section, ensuring the AI becomes a true expert in both without losing its soul.

1. Problem Statement

Continual Pre-training is a standard approach for adapting Large Language Models (LLMs) to new, knowledge-intensive target fields (e.g., math reasoning, code generation). However, directly training on target data often leads to catastrophic forgetting, where the model loses its original general capabilities due to significant distribution shifts between source and target domains.

The common solution is Domain Reweighting: curating a data mixture of source and target fields to balance performance. Existing methods rely on:

Manual Heuristics: Human intuition or empirical rules (e.g., "Science data helps MMLU").
Static/Dynamic Algorithms: Methods like RegMix (regression-based) or DBL (loss-based dynamic loading) that require extensive proxy model training or specific evaluation sets for every domain.
Limitations: These approaches are often labor-intensive, lack generalizability to unseen domains/models, and fail to capture complex, non-linear heuristics regarding data mixing trajectories.

2. Methodology: Data Mixing Agent

The authors propose Data Mixing Agent, the first model-based, end-to-end framework that learns to re-weight domains using Reinforcement Learning (RL).

A. Formulation as an MDP

The domain reweighting process is formalized as a Markov Decision Process (MDP):

State ( $S$ ): The history of data distributions from previous reweighting steps and the corresponding environment feedback (performance metrics).
Action ( $A$ ): The probability distribution over the domain space for the current training step (i.e., the ratio of data sampled from each domain).
Reward ( $R$ ): Feedback from a lightweight evaluation environment assessing the model's performance on target capabilities (e.g., MMLU for general, MATH for math reasoning).

B. Training Pipeline

Trajectory Sampling: The authors randomly sample thousands of data mixing trajectories. To ensure diversity, they use an inductive scoring algorithm that penalizes distributions deviating too far from the previous step or failing to align with the target distribution, while encouraging diversity across trajectories.
Proxy Model Training: Small proxy models (50M parameters) are trained from scratch on these sampled trajectories. At each step, they are evaluated in a lightweight environment to generate feedback vectors.
Agent Optimization (RL):
- Architecture: A lightweight Transformer decoder (2.1M parameters) that takes the trajectory history and feedback as input to predict the next domain distribution.
- Algorithm: The agent is first warm-up via Supervised Fine-Tuning (SFT) on high-quality trajectories. It is then optimized using Conservative Q-Learning (CQL), an offline RL algorithm. CQL is chosen to prevent overestimation of Q-values for out-of-distribution actions, ensuring the agent learns robust heuristics from both successful and failed trajectories.

C. Inference

During actual continual pre-training of the target model, the agent operates on-the-fly. It observes the current state (training history and performance feedback) and predicts the optimal domain distribution for the next step without requiring retraining.

3. Key Contributions

Novel Framework: Introduction of the Data Mixing Agent, a model-based approach that learns generalizable heuristics for domain reweighting, replacing manual or heuristic-based strategies.
Generalization: The agent demonstrates strong zero-shot generalization across:
- Unseen Source Fields: Works with different pre-training data distributions (e.g., DCLM, Fineweb-Edu).
- Unseen Target Models: Effective on models of varying sizes and architectures (3B LLaMA variants, Pythia).
- Unseen Target Domains: Successfully transfers from math reasoning to code generation without retraining.
- Domain Spaces: Generalizes across different domain definitions (2D source-based vs. 52D topic-based).
Efficiency: Achieves superior performance with fewer source-field tokens, reducing the total training budget and computational cost compared to static baselines.

4. Experimental Results

Experiments were conducted on Math Reasoning and Code Generation tasks using 3B parameter LLaMA models and Pythia models.

Performance:
- DataAgentRL outperformed state-of-the-art baselines (RegMix, DBL, Naive Training).
- On Math Reasoning, it achieved an average improvement of 3.02% over RegMix and 8.88% over the base model across 12 benchmarks.
- It significantly reduced performance variance, ensuring balanced capabilities across general and target fields.
Generalization:
- When applied to Code Generation (a domain not seen during training), the agent still outperformed RegMix and DBL by 1.45% and 2.78% respectively, proving the transferability of learned heuristics.
Efficiency:
- The agent achieved early stopping (converging faster) and required ~2.14B fewer tokens from the source field compared to RegMix to reach superior performance.
- Training the agent itself is computationally cheap (approx. 10 minutes for RL optimization after SFT), and the inference overhead is negligible.
Heuristic Alignment:
- Analysis showed the agent's learned strategies aligned with human intuition (e.g., increasing "Science" and "Health" domains for MMLU, decreasing "Fashion") while discovering new nuances (e.g., specific handling of "Pets & Animals").

5. Significance

Automation of Data Curation: Moves domain reweighting from a manual, heuristic-driven process to an automated, learned capability.
Scalability: The ability to generalize across domains and models without retraining makes it highly scalable for adapting LLMs to new verticals (e.g., medicine, law) where data mixing recipes are unknown.
Cost Reduction: By optimizing the data mixture to minimize catastrophic forgetting while maximizing target performance, it reduces the computational cost of continual pre-training by lowering the required token budget.
Theoretical Insight: Validates that a rich, model-agnostic heuristic space exists for data mixing, which can be effectively parameterized by a small agent model using offline RL.

In conclusion, the Data Mixing Agent represents a significant step forward in making continual pre-training more robust, efficient, and adaptable, solving the critical trade-off between acquiring new skills and retaining existing knowledge.