PaperlayerView exact revision on arXiv

Source-linked AI summary

Aioli: A Unified Optimization Framework for Language Model Data Mixing

Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, Christopher Ré

arXiv:2411.05735v2cs.LGcs.AIcs.CLstat.ML

TL;DR

Language-model data mixing requires efficient selection of proportions across heterogeneous data groups, yet existing methods do not consistently outperform stratified sampling. The paper unifies them through mixing laws, diagnoses inaccurate parameter estimation, and develops AIOLI to estimate those parameters online. AIOLI outperforms stratified sampling across all six datasets and improves shorter-run proportion transfer in most evaluated cases.

  • Problem

    Existing data-mixing methods do not consistently outperform a simple stratified-sampling baseline, making their shared assumptions and effectiveness unclear.

  • Method

    The paper introduces LMO to express mixing methods as constrained loss minimization and develops AIOLI, which estimates mixing-law parameters from current-run histories to adjust proportions online without extra training runs.

  • Results

    AIOLI outperforms stratified sampling on all 6 datasets by an average of 0.274 test-perplexity points, while improving shorter-run learned proportions by an average of 1.202 points in 28 of 30 cases.

  • Takeaways & Limitations

    Mixing-law parameter fidelity provides a basis for analyzing existing methods and designing more reliable dynamic data-mixing algorithms.

  • Takeaways & Limitations

    AIOLI incurs extra inference cost from repeated evaluations, although evaluating validation loss on subsets and using each parameter estimate longer can reduce it.

Abstract

from arXiv · show

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. To understand this inconsistency, we unify existing methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law -- an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of a method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we observe. Using this insight, we derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.27 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.012 test perplexity points.

1 Introduction

Data mixing methods aim to learn language-model training proportions efficiently, but existing approaches do not consistently beat stratified sampling. The paper unifies these methods through implicit mixing laws, diagnoses inaccurate parameter estimates, and introduces AIOLI to estimate them online.

  • Motivation: Existing data-mixing methods do not consistently outperform uniform stratified sampling in average test perplexity across data-group sets.This motivates analyzing their shared assumptions rather than treating each method independently.
  • Unified framework: The LMO framework expresses many methods as minimizing average group loss subject to a method-dependent mixing law relating losses to proportions.Existing methods use linear or log-linear parameterizations, with offline methods using static laws and online methods using dynamic laws.
  • Fidelity analysis: Both linear-dynamic and log-linear-static parameterizations fit observed loss–proportion relationships, but existing methods often set their mixing-law parameters inaccurately.Across datasets, these parameterizations achieve 0.0005 average MSE and 0.969 R2; parameter deviations correlate with performance relative to stratified sampling.
  • AIOLI: AIOLI estimates mixing-law parameters from the current run’s loss and proportion history, enabling dynamic adjustment without extra training runs.This directly targets the parameter-estimation problem identified by the LMO analysis.
  • Results: AIOLI beats stratified sampling on all 6 evaluated datasets by an average of 0.274 test-perplexity points, while existing methods can be up to 6.9 points worse.In shorter-run proportion-learning settings, AIOLI improves performance by an average of 1.202 perplexity points in 28 of 30 cases.

2 Problem Setup

The paper formulates data mixing as choosing static or dynamic proportions over data groups during language-model training. The objective is to minimize total test loss across groups, but without additional structure this requires expensive brute-force searches.

  • Data and mixing: Data mixing divides training examples among m data groups using proportions that may remain static or change across T training rounds.Each round uses a probability-simplex vector, while static mixtures use T = 1.
  • Model and loss: The model’s group-specific training, validation, and test losses depend on the sequence of proportions used before and during training.Losses are evaluated separately for each data group at the relevant training stage.
  • Optimization objective: The data-mixing problem chooses proportions for S-step training with N samples and T rounds to minimize total test loss across groups.This objective is intended to produce a model that performs well across many data groups.
  • Computational challenge: Without structural assumptions about how proportions affect future losses, solving the objective requires brute-force training of many models.The LMO framework adds a mixing-law constraint to make more efficient methods possible.

3 A Unified Optimization Framework for Data Mixing

LMO unifies data-mixing algorithms as solutions to a common optimization problem constrained by an implicit mixing law. It shows that methods differ in law parameterization, parameter values, and optimization solver, clarifying why inaccurate parameters can degrade performance.

  • LMO formulation: LMO constrains the data-mixing optimization problem with a mixing law that specifies the assumed relationship between group losses and mixture proportions.The law uses parameters such as A_t, b_t, c_t and either identity or exponential transformations.
  • Framework components: Existing methods are characterized by three choices: mixing-law parameterization, parameter values, and the procedure used to solve for proportions.These choices expose each method’s assumptions within a common framework.
  • Dynamic solving: Exponentiated gradient descent converts a linear dynamic mixing law into a greedy update rule whose direction depends on the mixing-law parameters.This establishes a correspondence between how online methods update proportions and the implicit laws they assume.
  • Existing methods: Offline methods such as DML use static log-linear laws fitted from training runs, whereas online methods such as Skill-It, DoReMi, and DoGE use linear dynamic laws with exponentiated-gradient updates.Their parameter constructions differ, including skills-graph interactions for Skill-It and diagonal parameters for DoReMi.
  • Framework insights: LMO reveals that all existing methods use linear or log-linear laws, but differ in parameter values and solvers; optimality depends on these assumptions holding in practice.Offline methods solve directly, while online methods use greedy exponentiated gradient descent.

4 Analyzing Fidelity of Existing Methods with the LMO Framework

The LMO framework evaluates whether existing mixing laws faithfully describe loss–proportion relationships and whether their parameters and optimization strategies support good mixing decisions.

  • 4.2 Mixing law parameterization: Existing log-linear static and linear dynamic mixing laws fit the true loss–proportion relationship well across six settings.The fitted laws achieve average MSEs of 8.9 × 10^-4 and 1.0 × 10^-4, with R2 values of 0.991 and 0.947, respectively.
  • 4.3 Values of mixing law parameters: Existing methods’ mixing-law parameters differ substantially from optimal parameters, and these disparities largely explain performance differences.Optimal parameters are approximated by fitting the laws to observed loss–proportion relationships.
  • 4.3 Values of mixing law parameters: The parameter-similarity metric combines cosine similarity and Spearman rank correlation over normalized column-sum vectors, because update direction depends on their ordering.The metric ranges from -1 to 1, with 1 indicating identical normalized vectors.
  • 4.3 Values of mixing law parameters: 0.491 R2 links parameter similarity to improvement over stratified sampling, indicating that parameter accuracy is critical for online-method performance.The comparison uses each method’s similarity to optimal parameters and its improvement over the stratified baseline across m = 2 and m = 3 settings.
  • 4.4 Solving strategy: Greedy exponentiated-gradient optimization does not significantly reduce performance relative to full dynamic optimization, despite the latter’s exponentially large solution space.Greedy selection recovers optimal dynamic proportions in 2 of 3 data settings.

5 AIOLI: a Method for Improved Data Mixing

AIOLI uses the LMO framework’s validated dynamic law and greedy optimizer while directly estimating accurate mixing-law parameters during training without extra training runs.

  • 5 AIOLI: a Method for Improved Data Mixing: AIOLI combines a linear dynamic mixing law and exponentiated gradient descent, whose assumptions generally hold in practice.Its primary design goal is accurate estimation of the true parameter matrix At⋆.
  • Estimating At: AIOLI estimates At from interleaved mixture sweeps, recording and averaging validation-loss changes without requiring extra training runs.A δ fraction of each round is partitioned into K = m^k intervals, with mixtures interleaved to estimate loss responses.
  • AIOLI: AIOLI normalizes the estimated parameter matrix before updating proportions, preventing early updates from becoming disproportionately large as losses decrease.The remaining portion of each round trains using the updated mixture.
  • AIOLI: AIOLI can initialize from a previously learned static mixture and dynamically adjust it after an initial training period.This supports starting from pinit after Sinit steps or beginning from a uniform mixture.

6 Experimental Results

AIOLI improves data mixing both with unrestricted method budgets and when existing proportions are learned on shortened runs before full training.

  • 6.1 Unrestricted Setting: 0.274 test-perplexity points: AIOLI outperforms stratified sampling across all 6 data settings, while other methods are up to 6.9 points worse.AIOLI uses no extra training steps, whereas competing methods use additional runs to learn proportions.
  • 6.1 Unrestricted Setting: AIOLI’s parameter similarity to At⋆ correlates with performance, and its lowest similarity is higher than that of competing methods.This supports more consistent parameter estimation across datasets.
  • 6.1 Unrestricted Setting: Offline methods may perform poorly at larger m because their proportion-learning budget is limited to 10S training steps.The paper hypothesizes that increasing this budget would eventually improve their performance.
  • 6.2 Restricted Setting: 1.202 average and 12.012 maximum test-perplexity points: adding AIOLI improves shortened-run methods in 28 of 30 restricted settings.AIOLI dynamically adjusts each method’s learned p throughout the full run.
  • 6.2 Restricted Setting: AIOLI can rescue methods that initially underperform stratified sampling, while preserving their advantage when they already outperform it.In the two regressions, the base methods already beat stratified sampling and perplexity worsens by at most 0.025 points.

7 Related Work

Related work frames data mixing through online adaptation, curriculum learning, data selection, and hyperparameter optimization, with AIOLI’s framework connecting these approaches to mixture-law assumptions.

  • Data mixing: Online data mixing has been formulated as a multi-armed bandit problem and as adaptive credit assignment over domain-level loss effects [2] [28].The credit-assignment approach corresponds to a diagonal At matrix in the LMO framework.
  • Data mixing: Other work transfers learned proportions across model scales or simulates mixtures by averaging models trained on different data partitions [24] [37] [46].These approaches address related mixture-transfer or mixture-simulation settings.
  • Curriculum Learning: Online mixing can also be viewed as curriculum learning over data groups, extending model-driven curricula beyond manually designed sample orderings [21] [26] [41] [65].Curriculum-learning work includes skill generalization and emphasizing high-quality data.
  • Data Selection: Data selection methods curate datasets at the sample level through filtering, matching, or condensation rather than group-level proportion mixing [1] [3] [32] [61] [71].The distinction is between selecting individual samples and setting mixture proportions across groups.
  • Hyperparameter Optimization and Truncation Bias: Data mixing is a hyperparameter-optimization problem in p, while mixing laws impose structure to reduce the inefficiency of brute-force search [37] [72].This perspective connects offline mixing methods to truncation and training-budget constraints.

8 Discussion

The LMO framework unifies data-mixing methods around assumptions about loss–proportion relationships and motivates AIOLI’s parameter-estimation strategy. The framework also exposes extra inference cost and unresolved questions about data-group partitions.

  • Discussion: LMO unifies existing data-mixing methods as solutions to a common optimization problem governed by an implicit method-dependent mixing law.The framework identifies mixing-law parameterization, parameter estimation, and optimization as axes for improving future methods.
  • Discussion: AIOLI’s gains are attributed to estimating linear dynamic mixing-law parameters throughout training.
  • Limitations and Future Work: AIOLI incurs extra inference cost because LEARNPARAMS repeatedly evaluates validation losses.The paper suggests reducing this cost by using a validation subset or retaining each parameter estimate for longer.
  • Limitations and Future Work: The effect of data-group partitioning remains unclear because C4 is a subset of CommonCrawl, motivating study of disjoint groups.
  • Ethics Statement: The work focuses on improving language-model training efficiency and performance rather than directly addressing ethical concerns.The ethics statement notes a possible indirect contribution through reduced computational costs and energy consumption.

B.1 Additional existing methods

The appendix relates additional data-mixing methods to the LMO framework and explains how their training and proportion-update procedures map to mixing laws and exponentiated gradient descent. It also contrasts ODM and RegMix with these formulations.

  • Additional existing methods: ODM frames data mixing as a multi-armed bandit, using EXP3 to explore groups and update proportions from group-level rewards.Its exploration and smoothing prevent direct expression in LMO, although its updates can be loosely interpreted as favoring high-loss groups.
  • Additional existing methods: RegMix fits regression models on shorter, smaller-model runs to predict proportions for longer runs on larger models.Its linear regression mixing law has relatively low R2=0.87, so its main approach uses LightGBM’s nonlinear decision-tree ensemble.
  • Exponentiated gradient descent: Exponentiated gradient descent updates proportions on the simplex using gradients of the timestep cost and a normalizing constant.In LMO, the cost is the weighted average validation loss, while the mixing-law constraint determines the relevant gradient.
  • LMO interpretation: Expressing an existing method in LMO means reproducing its training and proportion-setting procedure under assumptions about the loss–proportion relationship, not claiming identical optimization problems.
  • Skill-It: The LMO framework expresses Skill-It using a linear dynamic mixing law and exponentiated gradient descent, with its update coefficients matched to the law’s parameters.The derivation treats the training procedure as sampling according to the selected proportions and maps Skill-It’s update rule into LMO form.
  • DoReMi and DoGE: DoReMi and DoGE proxy-model procedures can also be represented in LMO through linear dynamic mixing laws and exponentiated gradient descent.DoReMi uses a diagonal parameter matrix tied to training-loss differences, whereas DoGE uses training–validation gradient inner products.

C.1.1 Additional parameterization experiments

Additional experiments largely support the specified mixing laws across checkpoints and data mixtures, while revealing task-dependent and boundary-specific deviations. These diagnostics also expose when matrix structure and time variation matter for recovering optimal proportions.

  • Parameterization across checkpoints: The log-linear static and linear dynamic laws fit later checkpoints and datasets sufficiently well to support their use across training stages.The supplied passages describe checkpoint-based sweeps and comparisons across training stages, but do not provide the associated table values.
  • Parameterization across other sets of data groups: Nonlinearities can occur at simplex boundaries, but they tend to have little effect on average-loss optimization because boundary regions incur high loss.For Books/C4, the fitted optimum [0.176, 0.824] was close to the grid-search optimum [0.2, 0.8].
  • Parameterization on instruction-tuning mixtures: Instruction-tuning mixtures favor the log-linear static law over the linear dynamic law, with average R2 values of 0.888 and 0.419, respectively.Open-ended generation tasks achieve higher dynamic-law R2 than binary and multiple-choice tasks, with an average of 0.73 for the former.
  • Checking for interactions among groups: Residual diagnostics show no strong interaction patterns, and the linear model appears sufficient despite possible small interactions.The residual plots examine predictions for Arxiv loss, while corresponding plots for other domains look similar.
  • Properties of optimal parameters: A full parameter matrix can reverse which group should receive priority, while time variation can change the ordering of optimal proportions.For Github/C4, the diagonal matrix prioritizes Github whereas the full matrix prioritizes C4; the ordering can also shift from Github initially to C4 later.
  • Solving strategy: Greedy dynamic optimization recovers the optimal proportions in 2 out of 3 examined cases, failing for Github/C4.The supplied evidence also notes that exponentiated-gradient descent can avoid extreme simplex-boundary regions where the linear law is misspecified.

D.3 Modification 2: “interleaving” training sweeps

Interleaving replaces checkpointed training sweeps with sequential short intervals, allowing mixing-law parameters to be estimated without extra training. The approach trades estimation accuracy and exploitation time against the duration and allocation of the learning phase.

  • Interleaving procedure: The modified procedure requires no extra training or checkpoint rollback by simulating simultaneous mixture sweeps through high-frequency sequential interleaving.This converts a training sweep into one round and divides the allocated steps into K intervals.
  • Interleaving procedure: Interleaving assigns candidate mixtures to sequential intervals and uses observed validation-loss changes to recover the parameter matrix from linear equations.The intervals are arranged so each mixture is trained during the sweep, and the resulting loss differences provide the estimation signals.
  • Bias-accuracy trade-off: Larger interleaving frequency reduces time bias in the estimated parameter matrix because each mixture is sampled at multiple positions within the sweep.With k = 1, a mixture may occur only at the beginning, middle, or end, making its loss change depend strongly on time.
  • Bias-accuracy trade-off: The shortened interleaved sweep estimates parameters less accurately than the naive approach and divides training between exploration and later exploitation.A larger learning fraction δ can improve parameter accuracy but leaves less time for using the updated proportions, while overly long exploration may be suboptimal.
  • Experimental setting: The experiments use 160M models with 5000-step or 40000-step runs, depending on the number of data groups, and average results across multiple seeds.The restricted and unrestricted settings retain the same final-run lengths, while their initialization budgets differ.

E.2 Data mixing methods

AIOLI estimates mixing-law parameters online and updates proportions throughout training, with hyperparameters controlling update frequency, sweep accuracy, smoothing, and exploration. Its trajectories demonstrate substantial adaptation across datasets and training-budget settings.

  • AIOLI-specific hyperparameters: AIOLI updates proportions using parameters estimated from the current run, rather than relying on extra training runs to initialize them.The method uses an exponential moving average option to decay the influence of earlier parameter estimates.
  • AIOLI-specific hyperparameters: More frequent proportion updates are obtained by increasing the number of rounds T, which is set to 20 in the experiments.The supplied passage defines larger T as producing more frequent updates.
  • AIOLI-specific hyperparameters: Increasing the sweep count k improves parameter accuracy but leaves less of each round for exploiting the resulting proportions.The same trade-off appears in the interleaving design, where more sampling across the sweep reduces time bias but consumes training budget.
  • AIOLI-specific hyperparameters: The smoothing factor ε controls the sweep mixtures between one-hot and uniform proportions, with ε = 0.75 used in all experiments.One-hot sweeps can produce poor learning dynamics, whereas ε = 1 yields only uniform proportions.
  • AIOLI-specific hyperparameters: The exploration fraction δ trades parameter-estimation accuracy against the amount of training that uses the learned proportions.The experiments use δ values of 0.128, 0.288, and 0.007 for m = 2, 3, and 7, respectively.
  • Weight trajectories: AIOLI significantly changes mixture proportions over time in both unrestricted and restricted settings across the reported dataset pairs.The trajectories also vary substantially with the initial proportions, indicating dependence on initialization.

F.1 Downstream Tasks

The paper tests AIOLI across downstream-task alignment, ablations, larger models, and out-of-domain evaluation. These results examine whether its dynamic parameter estimation remains effective across settings and clarify where perplexity-based optimization applies.

  • Downstream Tasks: A perplexity-based objective may misalign with downstream performance: lower perplexity correlated with worse downstream results, with correlation 0.529.The evaluation covered eight downstream tasks after SlimPajama training.
  • Downstream Tasks: The downstream mismatch remains unresolved: AIOLI robustly minimizes average test perplexity, but incorporating downstream evaluations into data mixing is left for future work.The paper attributes the mismatch partly to distribution shift between pretraining and downstream data.
  • Ablations: AIOLI outperforms both ablations in 3 of 6 settings and achieves the lowest average test perplexity, indicating that continual updates and off-diagonal modeling matter.AIOLI-STATIC learns parameters once, whereas AIOLI-DIAGONAL omits cross-group interactions.
  • Results on Larger Models: At 1.4B parameters, the static and dynamic mixing laws achieve average R2 values of 0.989 and 0.929, respectively.These results support the parameterizations as a prerequisite for evaluating AIOLI at larger scale.
  • Results on Larger Models: AIOLI outperforms stratified sampling at both 160M and 1.4B scales, while DoGE continues to underperform stratified sampling at 1.4B.The larger-model comparison used three random seeds and two data settings.
  • Out-of-domain Setting: In the out-of-domain setting, every evaluated method beats stratified sampling, but AIOLI is the only one doing so without additional training runs.The comparison trains on Arxiv, Books, and Github and evaluates on StackExchange; distribution shift can make uniform mixing include irrelevant data.

G Why the method is called AIOLI

The name AIOLI is explained through an analogy between an emulsion and the method’s parameter updates. The surrounding tables document mixture proportions, downstream evaluations, ablations, hyperparameter sweeps, and larger-scale comparisons.

  • Why the method is called AIOLI: AIOLI is named after an emulsion because separate components remain distinct while being combined into one mixture, like separate test-run estimates combined into one update.The analogy specifically refers to the At matrix and the update for pt.
  • Downstream Evaluation: Table 19 reports downstream evaluation metrics for data mixing methods trained on SlimPajama across three random seeds.The table concerns unrestricted-setting downstream evaluation rather than the method-name analogy.
  • Ablations: Table 20 reports ablation differences in average test perplexity relative to stratified sampling, with negative values indicating improvement.It covers AIOLI, AIOLI-STATIC, and AIOLI-DIAGONAL across six data settings.
  • Hyperparameter Robustness: Tables 21–23 vary AIOLI’s η, δ/m, and γ hyperparameters and report differences in average test perplexity relative to stratified sampling.These correspond to step size, the fraction of each round used to learn At, and exponential-moving-average computation.
  • Larger Models: At 1.4B scale, Table 25 reports test-perplexity differences relative to stratified sampling for two data settings and specifies the AIOLI hyperparameters used.The settings are Arxiv/StackExchange and Arxiv/Books/StackExchange.
Loading 2411.05735v2…