MixAtlas: Uncertainty-aware Data Mixture Optimization for Multimodal LLM Midtraining
Organizations: University of Washington · Work done during an internship at Apple · Apple · Virginia Tech
Abstract
Domain reweighting can improve sample efficiency and downstream generalization, but data-mixture optimization for multimodal midtraining remains largely unexplored. Current multimodal training recipes tune mixtures along a single dimension, typically data format or task type. We introduce MixAtlas, a method that produces benchmark-targeted data recipes that can be inspected, adapted, and transferred to new corpora. MixAtlas decomposes the training corpus along two axes: image concepts (10 visual-domain clusters discovered via CLIP embeddings) and task supervision (5 objective types including captioning, OCR, grounding, detection, and VQA). Using small proxy models (Qwen2-0.5B) paired with a Gaussian-process surrogate and GP-UCB acquisition, MixAtlas searches the resulting mixture space with the same proxy budget as regression-based baselines but finds better-performing mixtures. We evaluate on 10 benchmarks spanning visual understanding, document reasoning, and multimodal reasoning. On Qwen2-7B, optimized mixtures improve average performance by 8.5%-17.6% over the strongest baseline; on Qwen2.5-7B, gains are 1.0%-3.3%. Both settings reach baseline-equivalent training loss in up to 2 times fewer steps. Recipes discovered on 0.5B proxies transfer to 7B-scale training across Qwen model families.
Figures & tables
| Task Axis | Image Concept Axis | |||||||||
| Benchmark | Unif | Reg | MixAtlas | Gain | Unif | Reg | Cham | MixAtlas | Gain | |
| Visual Understanding | ||||||||||
| AI2D | 58.16 | 58.10 | 59.84 | +2.9% | 54.30 | 53.00 | 56.50 | 57.20 | +1.2% | |
| GQA | 18.95 | 20.90 | 30.65 | +46.6% | 25.14 | 23.35 | 23.90 | 36.48 | +45.1% | |
| ScienceQA | 72.53 | 72.30 | 74.81 | +3.1% | 71.60 | 72.90 | 72.30 | 71.50 | -1.9% | |
| Document and Text | ||||||||||
| Task Axis | Image Concept Axis | |||||||||
| Benchmark | Unif | Reg | MixAtlas | Gain | Unif | Reg | Cham | MixAtlas | Gain | |
| Visual Understanding | ||||||||||
| AI2D | 58.65 | 60.80 | 59.97 | -1.4% | 55.31 | 62.21 | 58.16 | 57.19 | -8.1% | |
| GQA | 32.61 | 39.00 | 45.05 | +15.5% | 25.14 | 44.27 | 23.90 | 36.48 | -17.6% | |
| ScienceQA | 74.62 | 73.80 | 73.18 | -1.9% | 74.07 | 73.53 | 73.33 | 71.49 | -3.5% | |
| Document and Text | ||||||||||
Data Mixture Optimization
Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.
See all 40 papers