Aioli: A Unified Optimization Framework for Language Model Data Mixing

• URL: http://arxiv.org/abs/2411.05735v1
• Date: Fri, 08 Nov 2024 17:50:24 GMT
• Title: Aioli: A Unified Optimization Framework for Language Model Data Mixing
• Authors: Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, Christopher Ré,
• Abstract summary: We show that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity per group. We derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions.
• Score: 74.50480703834508
• License:
• Abstract: Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over training runs to dynamically updating proportions throughout training. Surprisingly, we find that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity per group. In this paper, we study the cause of this inconsistency by unifying existing methods into a standard optimization framework. We show that all methods set proportions to minimize total loss, subject to a method-specific mixing law -- an assumption on how loss is a function of mixture proportions. We find that existing parameterizations of mixing laws can express the true loss-proportion relationship empirically, but the methods themselves often set the mixing law parameters inaccurately, resulting in poor and inconsistent performance. Finally, we leverage the insights from our framework to derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Empirically, Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.28 test perplexity points, whereas existing methods fail to consistently beat stratified sampling, doing up to 6.9 points worse. Moreover, in a practical setting where proportions are learned on shorter runs due to computational constraints, Aioli can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.01 test perplexity points.

Related papers

• Parallel Split Learning with Global Sampling [9.57839529462706]
  parallel split learning has emerged as a promising derivative of split learning well suited for distributed learning on resource-constrained devices. These challenges include large effective batch sizes, non-independent and identically distributed data, and the straggler effect. We introduce a new method called uniform global sampling to decouple the effective batch size from the number of clients and reduce the mini-batch deviation. Our simulations reveal that our proposed methods enhance model accuracy by up to 34.1% in non-independent and identically distributed settings and reduce the training time in the presence of stragglers by up to 62%.
  arXiv  Detail & Related papers  (2024-07-22T15:41:23Z)
• Optimal Baseline Corrections for Off-Policy Contextual Bandits [61.740094604552475]
  We aim to learn decision policies that optimize an unbiased offline estimate of an online reward metric. We propose a single framework built on their equivalence in learning scenarios. Our framework enables us to characterize the variance-optimal unbiased estimator and provide a closed-form solution for it.
  arXiv  Detail & Related papers  (2024-05-09T12:52:22Z)
• 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. We propose nested use of the scaling laws of training steps, model sizes, and our data mixing law. 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)
• Rethinking Classifier Re-Training in Long-Tailed Recognition: A Simple Logits Retargeting Approach [102.0769560460338]
  We develop a simple logits approach (LORT) without the requirement of prior knowledge of the number of samples per class. Our method achieves state-of-the-art performance on various imbalanced datasets, including CIFAR100-LT, ImageNet-LT, and iNaturalist 2018.
  arXiv  Detail & Related papers  (2024-03-01T03:27:08Z)
• Tailoring Mixup to Data for Calibration [12.050401897136501]
  Mixup is a technique for improving calibration and predictive uncertainty. In this work, we argue that the likelihood of manifold intrusion increases with the distance between data to mix. We propose to dynamically change the underlying distributions of coefficients depending on the similarity between samples to mix.
  arXiv  Detail & Related papers  (2023-11-02T17:48:28Z)
• Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
  Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios. We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
  arXiv  Detail & Related papers  (2022-04-13T12:43:12Z)
• FORML: Learning to Reweight Data for Fairness [2.105564340986074]
  We introduce Fairness Optimized Reweighting via Meta-Learning (FORML) FORML balances fairness constraints and accuracy by jointly optimizing training sample weights and a neural network's parameters. We show that FORML improves equality of opportunity fairness criteria over existing state-of-the-art reweighting methods by approximately 1% on image classification tasks and by approximately 5% on a face prediction task.
  arXiv  Detail & Related papers  (2022-02-03T17:36:07Z)
• Network Estimation by Mixing: Adaptivity and More [2.3478438171452014]
  We propose a mixing strategy that leverages available arbitrary models to improve their individual performances. The proposed method is computationally efficient and almost tuning-free. We show that the proposed method performs equally well as the oracle estimate when the true model is included as individual candidates.
  arXiv  Detail & Related papers  (2021-06-05T05:17:04Z)
• Scalable Personalised Item Ranking through Parametric Density Estimation [53.44830012414444]
  Learning from implicit feedback is challenging because of the difficult nature of the one-class problem. Most conventional methods use a pairwise ranking approach and negative samplers to cope with the one-class problem. We propose a learning-to-rank approach, which achieves convergence speed comparable to the pointwise counterpart.
  arXiv  Detail & Related papers  (2021-05-11T03:38:16Z)

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.