Msft Mixture Overfitting

Ranked Findings

Finding #1: Heterogeneous Overfitting Dynamics (Primary)

Discovery: Individual sub-datasets within multi-task mixtures reach peak performance at substantially different training points, contradicting industry standard practice of uniform training duration.

Evidence: Figure 2 demonstrates across all tested configurations that different sub-datasets achieve maximum validation accuracy at different compute levels. This heterogeneity is systematic, not noise—consistent across model scales and task granularities.

Implication: Standard practice of training all tasks for identical duration is suboptimal. Each task has a distinct "peak performance point" determined by:

• Task difficulty (harder tasks may need more compute)
• Dataset size (smaller datasets overfit earlier)
• Domain relatedness (related tasks may peak together)

Decision Impact: Stopping uniformly at any point necessarily leaves some tasks undertrained (below peak) while others are overtrained (past peak). Heterogeneous stopping can recover performance lost to this tradeoff.

Finding #2: Parameter Divergence Problem (Secondary)

Discovery: When excluding even small fractions of training data (1/10), optimal stopping points for remaining datasets shift significantly. Naive approaches that pre-compute stopping points from an initial search then exclude datasets fundamentally fail.

Evidence: Empirical measurements show:

• Excluding 1/10 of data shifts average optimal stopping point by 0.91 epochs
• Shifts are non-uniform (some tasks shift +2 epochs, others -1)
• Shifting parameters in either direction reduces performance on remaining tasks

Root Cause: The optimization landscape couples all datasets through shared parameters. Removing one dataset creates new local optima for remaining tasks—the coupling cannot be separated.

Implication: Cannot pre-compute optimal stopping per-dataset once, then use those points sequentially. Must search iteratively as datasets are modified.

Finding #3: SFT Compute Negligible (Enabling Condition)

Discovery: SFT (supervised fine-tuning) stage accounts for only 0.01% of total pre-training compute, making exhaustive per-dataset optimization computationally justified.

Arithmetic:

• Total pre-training compute: ~100 GPU-years
• SFT stage: ~0.01 GPU-years
• Even 100× search overhead in SFT remains <1% total training

Implication: Computational budget is not a limiting factor—exhaustive search within SFT is free relative to full training. Enables principled optimization without cost penalty.

Implementation: mSFT Algorithm

Core Insight: Iterative roll-out and roll-back search, maintaining coupling between search and training phases.

def msft(datasets, budget_C):
    """
    Iterative per-dataset optimization with parameter coupling maintained.

    Args:
        datasets: list of training datasets D1, D2, ..., Dn
        budget_C: compute budget per roll-out phase
    """
    active_mixture = datasets.copy()
    peaks = {}  # optimal stopping point per dataset

    while active_mixture:
        # Phase 1: Roll-out
        # Train on current mixture for fixed budget C, record peaks
        checkpoint = train(active_mixture, budget=budget_C)
        for dataset in active_mixture:
            peak_epoch = find_peak(dataset, checkpoint)
            peaks[dataset] = peak_epoch

        # Phase 2: Identify early overfitter
        earliest_peak_dataset = min(active_mixture,
                                   key=lambda d: peaks[d])
        earliest_peak_epoch = peaks[earliest_peak_dataset]

        # Phase 3: Roll-back and remove
        # Return to checkpoint where earliest dataset peaked
        rolled_back = checkpoint[earliest_peak_epoch]

        # Extract parameters at this point
        final_weights = rolled_back.model_state

        # Record stopping point for this dataset
        record(earliest_peak_dataset, final_weights, earliest_peak_epoch)

        # Remove dataset that peaked earliest
        active_mixture.remove(earliest_peak_dataset)

        # Repeat with remaining datasets
        # (Don't reset training; continue from rolled-back point)
        checkpoint = continue_training(active_mixture,
                                      initial_weights=final_weights,
                                      budget=budget_C)

    return peaks  # Per-dataset optimal stopping points

Algorithm Details:

The key is maintaining parameter coupling during removal:

1. Train all datasets together for budget C
2. Identify which dataset peaked earliest
3. Roll back to that dataset's peak (not to original point)
4. Exclude that dataset
5. Continue training remaining datasets from the rolled-back checkpoint
6. Repeat until one dataset remains

This differs from naive approaches that:

• Train once, record peaks
• Exclude datasets sequentially without retraining
• Leads to parameter drift (remaining tasks' optima shift)

Decision Checklist

Before using mSFT:

• Have multiple datasets/tasks in SFT mixture (not single dataset)
• Can measure per-dataset performance (validation set per task)
• Willing to tolerate multi-pass training through data (roll-out phases)
• Budget C chosen (fixed compute per roll-out; empirically insensitive)

When implementing:

• Implement peak detection (track best validation accuracy per dataset)
• Save checkpoints during training (need to roll back to exact epoch)
• Maintain exact parameter coupling (continue from rolled-back weights)
• Iterate until single dataset remains
• Extract final weights from last roll-back (optimal for final task)

Hyperparameter:

• budget_C: Compute per roll-out phase
  • Smaller C = more iterations (finer granularity)
  • Larger C = fewer iterations (faster search)
  • Empirically insensitive (robust across wide range)
  • Suggested: 10-20% of full training duration

Applicable Conditions

Model Scales

Robust across:

• 0.5B parameters
• 1B parameters
• 4B parameters
• 8B parameters

Performance gains consistent regardless of scale.

Dataset Configurations

Works with:

• 9K-sample datasets
• 27K-sample datasets
• Variable dataset sizes (heterogeneous mixtures)

Task Granularities

Applicable to:

• 5-task mixtures
• 15-task mixtures
• Single datasets with 21 sub-categories (still has heterogeneous peaks)

Compute Budgets

Achieves gains even with:

• budget=1: +3.4% improvement with reduced FLOPs
• Low compute settings (fractional budgets)
• Full training duration (standard settings)

When NOT to Use

• Single dataset (no mixture heterogeneity to exploit)
• Extremely tight latency requirements (multi-pass training expensive)
• Pre-computed dataset schedules must be frozen (need flexibility to recompute)
• Tasks have hard dependencies (must train in specific order)

Empirical Results Summary

Improvement Profile:

• Average gain: +3-5% across tested configurations
• Consistent across scales and dataset sizes
• Larger improvements on harder mixtures (greater heterogeneity)

Computational Overhead:

• Multi-pass training: ~2-3× data passes
• SFT compute: 0.01% of total training (negligible amortization)
• Per-roll-out: budget C (controlled parameter)

Robustness:

• Hyperparameter insensitive (C doesn't require tuning)
• Works out-of-box for any dataset mixture
• No manual per-task configuration needed

Integration with Training Pipelines

Typical Use:

Pre-training (99.99% compute)
    ↓
mSFT Phase (0.01% compute)
  - Roll-out/roll-back search
  - Per-dataset peak identification
  - Optimal stopping points
    ↓
Final fine-tuning (using found stopping points)
    ↓
Evaluation

Drop-in replacement for standard SFT: maintains all downstream compatibility.