Msft Mixture Overfitting

---
name: msft-mixture-overfitting
title: "mSFT: Ranked Findings on Supervised Fine-Tuning Dataset Mixture Overfitting"
version: 0.0.3
engine: skillxiv-v0.0.3-claude-opus-4.6
license: MIT
url: https://arxiv.org/abs/2603.21606
keywords: [Training, Dataset Mixture, Overfitting, Fine-Tuning, Empiricism]
description: "Identify three ranked findings on multi-task SFT: (1) heterogeneous overfitting—sub-datasets peak at different training points (contradicts uniform duration practice); (2) parameter divergence—excluding 1/10 of data shifts optimal points 0.91 epochs for remaining tasks; (3) SFT compute negligible (0.01% of training). Implement mSFT: iterative roll-out/roll-back search per-dataset. Robust across 0.5B-8B models, 9K-27K samples, 5-15 tasks, achieving +3.4% improvement with reduced FLOPs."
---

## 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.

```python
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.