ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization

• URL: http://arxiv.org/abs/2508.11551v2
• Date: Mon, 18 Aug 2025 06:38:38 GMT
• Title: ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization
• Authors: Shengzhuang Chen, Xu Ouyang, Michael Arthur Leopold Pearce, Thomas Hartvigsen, Jonathan Richard Schwarz,
• Abstract summary: We present results for pre-training and instruction finetuning across models ranging from 1 million to 7 billion parameters.<n>We demonstrate consistently strong results relative to a wide range of baselines, resulting inspeed-ups of over 500%.<n>In addition, we broaden access to research by sharing ADMIRE IFT Runs, a dataset of 460 full training & evaluation runs worth over 13,000 GPU hours.
• Score: 11.433087692377779
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: Determining the optimal data mixture for large language model training remains a challenging problem with an outsized impact on performance. In practice, language model developers continue to rely on heuristic exploration since no learning-based approach has emerged as a reliable solution. In this work, we propose to view the selection of training data mixtures as a black-box hyperparameter optimization problem, for which Bayesian Optimization is a well-established class of appropriate algorithms. Firstly, we cast data mixture learning as a sequential decision-making problem, in which we aim to find a suitable trade-off between the computational cost of training exploratory (proxy-) models and final mixture performance. Secondly, we systematically explore the properties of transferring mixtures learned at a small scale to larger-scale experiments, providing insights and highlighting opportunities for research at a modest scale. By proposing Multi-fidelity Bayesian Optimization as a suitable method in this common scenario, we introduce a natural framework to balance experiment cost with model fit, avoiding the risks of overfitting to smaller scales while minimizing the number of experiments at high cost. We present results for pre-training and instruction finetuning across models ranging from 1 million to 7 billion parameters, varying from simple architectures to state-of-the-art models and benchmarks spanning dozens of datasets. We demonstrate consistently strong results relative to a wide range of baselines, resulting inspeed-ups of over 500% in determining the best data mixture on our largest experiments. In addition, we broaden access to research by sharing ADMIRE IFT Runs, a dataset of 460 full training & evaluation runs worth over 13,000 GPU hours, greatly reducing the cost of conducting research in this area.

Related papers

• Linear Model Merging Unlocks Simple and Scalable Multimodal Data Mixture Optimization [38.78268216433473]
  We study model merging as an efficient strategy for estimating the performance of different data mixtures.<n>We conduct experiments on 14 multimodal benchmarks, and empirically demonstrate that the proxy models exhibit a high rank correlation with models trained on actual data mixtures.
  arXiv  Detail & Related papers  (2026-02-04T16:06:39Z)
• Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models [5.626695215498112]
  We frame data mixing as an optimization problem and introduce a novel method designed to minimize validation loss.<n>Our approach parametrizes the loss by modeling effective data transferred and leveraging scaling laws for fine-tuning.<n>We show that our algorithm achieves excellent overall and individual performance across all domains.
  arXiv  Detail & Related papers  (2025-08-16T07:28:39Z)
• SPaRFT: Self-Paced Reinforcement Fine-Tuning for Large Language Models [51.74498855100541]
  Large language models (LLMs) have shown strong reasoning capabilities when fine-tuned with reinforcement learning (RL)<n>We propose textbfSPaRFT, a self-paced learning framework that enables efficient learning based on the capability of the model being trained.
  arXiv  Detail & Related papers  (2025-08-07T03:50:48Z)
• Combatting Dimensional Collapse in LLM Pre-Training Data via Diversified File Selection [65.96556073745197]
  DiverSified File selection algorithm (DiSF) is proposed to select the most decorrelated text files in the feature space.<n>DiSF saves 98.5% of 590M training files in SlimPajama, outperforming the full-data pre-training within a 50B training budget.
  arXiv  Detail & Related papers  (2025-04-29T11:13:18Z)
• Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework [10.317740844867913]
  We build a simulator based on 472 language model pre-training runs with varying data compositions from the SlimPajama dataset.<n>We observe that even simple acquisition functions can enable principled training decisions across training models from 20M to 1B kernels.
  arXiv  Detail & Related papers  (2025-03-26T22:19:47Z)
• Computation-Aware Gaussian Processes: Model Selection And Linear-Time Inference [55.150117654242706]
  We show that model selection for computation-aware GPs trained on 1.8 million data points can be done within a few hours on a single GPU.<n>As a result of this work, Gaussian processes can be trained on large-scale datasets without significantly compromising their ability to quantify uncertainty.
  arXiv  Detail & Related papers  (2024-11-01T21:11:48Z)
• Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance [55.872926690722714]
  We study the predictability of model performance regarding the mixture proportions in function forms.<n>We propose nested use of the scaling laws of training steps, model sizes, and our data mixing law.<n>Our method effectively optimize the training mixture of a 1B model trained for 100B tokens in RedPajama.
  arXiv  Detail & Related papers  (2024-03-25T17:14:00Z)
• How to Train Data-Efficient LLMs [56.41105687693619]
  We study data-efficient approaches for pre-training language models (LLMs) We find that Ask-LLM and Density sampling are the best methods in their respective categories. In our comparison of 19 samplers, involving hundreds of evaluation tasks and pre-training runs, we find that Ask-LLM and Density are the best methods in their respective categories.
  arXiv  Detail & Related papers  (2024-02-15T02:27:57Z)
• Efficient Online Data Mixing For Language Model Pre-Training [101.45242332613944]
  Existing data selection methods suffer from slow and computationally expensive processes. Data mixing, on the other hand, reduces the complexity of data selection by grouping data points together. We develop an efficient algorithm for Online Data Mixing (ODM) that combines elements from both data selection and data mixing.
  arXiv  Detail & Related papers  (2023-12-05T00:42:35Z)
• Towards Accelerated Model Training via Bayesian Data Selection [45.62338106716745]
  We propose a more reasonable data selection principle by examining the data's impact on the model's generalization loss. Recent work has proposed a more reasonable data selection principle by examining the data's impact on the model's generalization loss. This work solves these problems by leveraging a lightweight Bayesian treatment and incorporating off-the-shelf zero-shot predictors built on large-scale pre-trained models.
  arXiv  Detail & Related papers  (2023-08-21T07:58:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.