Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Instruction Tuning Training › Data Mixing Strategies

Michael BrenndoerferDecember 19, 202516 min read

Language AI HandbookAll chapters

Part 36 of 67Instruction Tuning

Chapter 270Instruction Tuning Training › Data Mixing Strategies

Reading section 2 of 8 · Data Mixing Strategies · 16 minSection 2 of 8 · 16 min

Section 2 of 8 · 16 min read

Instruction Tuning Training › Data Mixing StrategiesOne section · Sections in this chapter2 / 8 ↓

Sections in this chapter8

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Instruction tuning datasets are not collections of a single task type. They aggregate examples from many domains. These may include question answering and summarization alongside code generation or mathematical reasoning. Other examples cover creative writing, translation, or classification. Each task type requires different capabilities. Question answering relies heavily on factual recall and information retrieval. Mathematical reasoning requires multi-step logical inference. Code generation demands precise syntax adherence and algorithmic thinking. Creative writing calls for coherent narrative construction and stylistic flexibility.

The distribution of these task types during training shapes what skills the model develops and how well it generalizes across the variety of requests it will encounter at deployment. Leave the distribution entirely to chance and you will get a model whose capabilities mirror your data collection process rather than your actual goals. A thoughtful data mixing strategy is one of the most impactful levers you control during instruction tuning, and it deserves as much attention as model architecture or learning rate.

Think of data mixing like designing a curriculum for a student. A music student who spends 90% of practice time on scales and only 10% on actual pieces will become technically proficient but may struggle with performance. The curriculum must reflect the actual skills the student needs, in proportions that ensure sufficient practice for each competency. Instruction tuning requires the same deliberate curation.

The Sampling ProblemLink Copied

Training on all available instruction data without adjustment causes task imbalance. Instruction datasets are rarely balanced because some tasks are far easier to collect than others. Question-answering pairs can be scraped at scale from documentation, FAQ pages, or discussion forums. High-quality mathematical reasoning examples require expert annotation, careful verification, and significant human effort. If your dataset contains 120,000 question-answering examples but only 8,000 math examples, the model sees QA tasks 15 times more frequently during training. This imbalance causes the model to perform well on common tasks while struggling with rare ones, even if both are equally important for the application you are building.

The imbalance compounds across training steps. After seeing one million examples during training, a model trained with proportional sampling has received gradient updates from 345,000 QA examples but only 23,000 math examples. The model's parameters are shaped far more by what helps it answer questions than by what helps it solve mathematical problems. The math skill does not simply fail to develop; it may regress if the QA optimization subtly interferes with mathematical reasoning patterns learned during pre-training.

Consider a hypothetical instruction dataset with the following distribution, which shows the kind of imbalance commonly found in real instruction tuning corpora:

In[3]:

Code

# Simulated task distribution in a typical instruction dataset
tasks = [
    "QA",
    "Summarization",
    "Code",
    "Math",
    "Creative",
    "Translation",
    "Classification",
]
raw_counts = [120000, 80000, 15000, 8000, 25000, 40000, 60000]

# Calculate proportions
total = sum(raw_counts)
proportions = [c / total for c in raw_counts]

Out[4]:

Console

Task Distribution in Raw Dataset:
----------------------------------------
QA               120,000 examples (34.5%)
Summarization     80,000 examples (23.0%)
Code              15,000 examples ( 4.3%)
Math               8,000 examples ( 2.3%)
Creative          25,000 examples ( 7.2%)
Translation       40,000 examples (11.5%)
Classification    60,000 examples (17.2%)
Total            348,000 examples

In proportional sampling, where examples are drawn according to their natural frequency, the model sees the QA category roughly 15 times more often than math. This imbalance biases the training signal toward the majority task in ways that go beyond simple frequency. The gradient updates from majority tasks dominate the loss landscape, pushing the model's parameters toward a region that is good for QA but potentially suboptimal for reasoning-heavy tasks.

Sampling StrategiesLink Copied

Three main approaches address task imbalance, each with different trade-offs that make them suitable for different situations and goals. Understanding when each approach is appropriate requires thinking carefully about your target deployment scenarios and the nature of the tasks you are balancing.

Proportional sampling uses the natural data distribution without modification. Examples appear according to their frequency in the dataset, meaning that if 34% of your data is question-answering, then 34% of your training batches will contain question-answering examples on average. This works well when the natural distribution matches your goals. A general assistant that users primarily interact with through simple questions might legitimately benefit from heavy QA training. The main drawback is underfitting on minority tasks. The model may not see enough math examples to develop reliable reasoning, even if that skill is important for the use case you care about.

Equal sampling takes the opposite approach, assigning each task equal probability regardless of dataset size. If you have seven task categories, each receives roughly 14.3% of the training examples. This strategy ensures the model sees rare tasks frequently enough to learn from every category. However, it introduces its own problems. It may waste model capacity by over-training on already well-represented tasks. More seriously, it can cause severe overfitting when small task datasets must be repeated many times to reach equal representation. If your math dataset has only 8,000 examples but needs to match 120,000 QA examples, each math example must be seen 15 times across a single epoch, risking memorization of specific problem instances rather than generalization to new problems.

Temperature-based sampling uses a temperature parameter to produce a smooth interpolation between proportional and equal sampling. This approach gives you continuous control over how aggressively to upweight minority tasks. The key insight is that you rarely want either extreme. You want minority tasks represented frequently enough for real learning, but you do not need or want perfectly uniform sampling if your tasks vary naturally in importance or volume. Temperature-based sampling lets you find the right balance for your specific situation.

Given raw task proportions $p_i$ , the sampling probability becomes:

q_i = \frac{p_i^{1/T}}{\sum_{j} p_j^{1/T}}

where:

$q_i$ : the adjusted sampling probability for task $i$ , representing how likely you are to draw an example from this task during training
$p_i$ : the raw proportion of task $i$ in the original dataset, calculated as the count of examples in task $i$ divided by the total number of examples
$p_i^{1/T}$ : the proportion for task $i$ raised to the power $1/T$ , which compresses the differences between large and small proportions to make the distribution more uniform
$T$ : the temperature parameter that controls how aggressively to flatten the distribution, with larger values creating more uniform sampling
$\sum_{j} p_j^{1/T}$ : the sum of exponentiated proportions across all tasks, which normalizes the adjusted values so they sum to one and form a valid probability distribution

The mathematical mechanism is elegant. When $T = 1$ , the exponent is $1/1 = 1$ , so $p_i^1 = p_i$ , and the formula reduces to $q_i = p_i$ . You recover proportional sampling exactly. As $T$ increases, the exponent $1/T$ approaches zero. Any positive number raised to a power near zero approaches one, so all exponentiated proportions converge toward one regardless of the original $p_i$ values. The normalized distribution therefore approaches the uniform distribution. The temperature parameter continuously interpolates between these two extremes.

The three key operating regimes are:

When $T = 1$ : The formula reduces to proportional sampling, preserving the original natural distribution without any adjustment.
When $T \to \infty$ : All sampling weights converge to the same value, creating perfectly uniform sampling where each task is equally likely regardless of its original size.
When $T \approx 2$ to $3$ : The distribution is flattened but not made uniform. Rare tasks are upweighted enough to appear frequently during training, while common tasks are still sampled more often than rare ones. This shows their greater natural occurrence. This intermediate regime gives the best balance for most practical instruction tuning scenarios.

In[5]:

Code

def temperature_sampling(proportions, temperature):
    """Apply temperature to sampling distribution."""
    adjusted = [p ** (1 / temperature) for p in proportions]
    total = sum(adjusted)
    return [a / total for a in adjusted]


# Compare different temperatures
temperatures = [1.0, 2.0, 5.0, float("inf")]
sampling_distributions = {}

for temp in temperatures:
    if temp == float("inf"):
        # Equal sampling
        sampling_distributions[temp] = [1 / len(tasks)] * len(tasks)
    else:
        sampling_distributions[temp] = temperature_sampling(proportions, temp)

Out[6]:

Visualization

Grouped bar chart showing task sampling probabilities for seven tasks under four temperature settings. At T=1, QA dominates at 34% while Math is near 2%. Higher temperatures flatten the distribution toward equal sampling. — Task sampling probabilities under four temperature regimes ($T=1, 2, 5, \infty$). At $T=1$, proportional sampling causes QA to dominate with over 34% of draws while Math receives only about 2%. As temperature increases, the distribution progressively flattens, boosting minority tasks like Math and Code. At $T=\infty$, all tasks receive exactly equal 14.3% probability, regardless of original dataset size.

Research from instruction tuning papers like FLAN suggests that moderate temperature values of approximately $T = 2$ to $T = 3$ often work well in practice. These values prevent minority tasks from being neglected during training. This keeps the model develops at least basic competency across all task types, while still which shows the natural importance of common tasks by sampling them somewhat more frequently. The exact optimal temperature depends on your specific dataset and use case, but the $T = 2$ to $T = 3$ range gives a reasonable and well-motivated starting point for experimentation.

The visualization below focuses specifically on minority task representation across temperature settings. The contrast between $T = 1$ and higher temperatures reveals precisely why temperature matters for the rare tasks that most benefit from adjustment.

Out[7]:

Visualization

Grouped bar chart comparing sampling probabilities for Math, Code, and Creative tasks at four temperature settings. All three minority tasks show substantially higher probabilities at T=2, T=5, and Equal compared to T=1. — Sampling probabilities for the three minority tasks (Math, Code, Creative) across four temperature settings. At proportional sampling ($T=1$), Math receives only 2.3% of draws, Code 4.3%, and Creative 7.2%. At $T=2$, these rise to approximately 6%, 8%, and 10% respectively, substantially increasing the training signal for skills that would otherwise be neglected.

Out[8]:

Visualization

Line chart showing boost factors versus task size order for four temperature settings. Small tasks have high boost factors at higher temperatures while large tasks are downweighted below 1.0. — Boost factors quantifying how much each task is upweighted relative to proportional sampling, plotted from smallest to largest task by dataset size. At T=1, all boost factors equal 1.0 by definition. At higher temperatures, small tasks like Math and Code receive boost factors of 3 to 7, while large tasks like QA are downweighted below 1.0 to compensate. This asymmetric reweighting is what makes temperature sampling effective for improving minority task coverage.

Multi-Epoch ConsiderationsLink Copied

When training for multiple epochs, you must carefully manage how sampling and repetition interact. With proportional sampling across multiple epochs, training remains dominated by common tasks. After three epochs, the model may have seen each QA example three times while still finding math examples sparse. Conversely, with equal sampling, rare task examples must be repeated far more frequently than common ones within each epoch to balance the distribution, and across multiple epochs this repetition compounds sharply, increasing the risk of memorizing specific examples rather than learning general patterns.

The risk of memorization is not hypothetical. A model that has seen the same 8,000 math problems dozens of times may produce correct answers by pattern-matching to memorized problem-answer pairs rather than by executing the underlying reasoning procedure. When it encounters a novel problem that is slightly different from the training examples, this memorized approach fails in ways that are hard to detect without careful evaluation. The model appears to perform well during training but fails to generalize.

To prevent this, you can cap how many times any single example appears during training:

In[9]:

Code

import numpy as np


def create_capped_dataset(task_examples, target_size, max_repeats=3):
    """
    Create balanced dataset with capped repetition.

    Args:
        task_examples: Dict mapping task name to list of examples
        target_size: Desired total dataset size
        max_repeats: Maximum times any example can appear
    """
    balanced_data = []
    examples_per_task = target_size // len(task_examples)

    for task, examples in task_examples.items():
        # How many times to repeat the full dataset
        n_examples = len(examples)
        repeats_needed = examples_per_task / n_examples

        if repeats_needed <= 1:
            # Subsample large datasets
            indices = np.random.choice(
                n_examples, examples_per_task, replace=False
            )
            balanced_data.extend([examples[i] for i in indices])
        elif repeats_needed <= max_repeats:
            # Repeat small datasets
            full_repeats = int(repeats_needed)
            remainder = examples_per_task - (full_repeats * n_examples)

            for _ in range(full_repeats):
                balanced_data.extend(examples)
            balanced_data.extend(
                np.random.choice(examples, remainder, replace=False).tolist()
            )
        else:
            # Cap at max_repeats
            capped_total = n_examples * max_repeats
            for _ in range(max_repeats):
                balanced_data.extend(examples)
            print(
                f"Warning: {task} capped at {capped_total} examples (wanted {examples_per_task})"
            )

    return balanced_data

In[10]:

Code

# Demonstrate with example task sizes
task_examples = {
    "QA": list(range(10000)),
    "Math": list(range(500)),
    "Code": list(range(2000)),
}

balanced = create_capped_dataset(
    task_examples, target_size=15000, max_repeats=3
)

Out[11]:

Console

Created balanced dataset with 11500 examples

The final dataset has 11,500 examples. The Math task was capped at 3 repetitions to prevent overfitting. This balances task representation while avoiding memorization. In practice, the warning output tells you which tasks are constrained by the cap, signaling that you might want to collect more data for those tasks or accept the limitation as a practical constraint.

Out[12]:

Visualization

Side-by-side bar chart for QA, Math, and Code showing repetitions needed for equal sampling versus the capped maximum of 3. Math bars show a large gap between the 10x needed and the 3x cap. — Comparison of repetitions required for equal sampling versus the maximum allowed repetition cap of 3, for three sample tasks. QA is so large it needs fewer than 1 repetition (subsampled). Code requires about 2.5 repetitions, fitting within the cap. Math would require 10 repetitions for equal sampling but is capped at 3, preventing memorization at the cost of under-representing this task in the final dataset.

Out[13]:

Visualization

Pie chart showing the final 11,500-example dataset composition after capping: QA 43.5%, Code 43.5%, and Math 13.0%. — Final composition of the 11,500-example dataset after applying repetition capping with a maximum of 3 repeats. QA and Code each contribute 5,000 examples (43.5% apiece), while Math contributes 1,500 examples (13.0%) because its 500 source examples reach the three-repeat cap. This composition is the practical outcome of balancing memorization risk against task coverage.

Key Data Mixing ParametersLink Copied

The two central parameters that govern data mixing are worth understanding in concrete terms:

temperature: Controls the sharpness of the sampling distribution. A higher temperature (for example $T = 2$ ) upweights minority tasks relative to their natural frequency. This keeps they receive more training attention. Lower temperatures preserve more of the original distribution, while very high temperatures approach uniform sampling. In practice, $T = 2$ is a safe default for multi-task instruction tuning.
max_repeats: The maximum number of times a single example can appear during balanced dataset construction. This parameter prevents overfitting on small task categories by capping repetition even when equal sampling would require seeing those examples far more often. A value of 3 is commonly used; values above 5 are generally risky for datasets smaller than a few thousand examples.

These two parameters interact. A high temperature combined with a low max_repeats cap means that rare tasks will be sampled often but can only appear a limited number of times, creating a tension that cannot be resolved without collecting more data. When you hit this limit, the practical answer is to invest in data collection for the under-represented tasks rather than trying to compensate through sampling tricks.