BiMix: A Bivariate Data Mixing Law for Language Model Pretraining

• Existing data-mixing methods rely on heuristics or computationally expensive optimization and lack a general scaling framework.This limits efficient resource allocation and performance prediction across data compositions.Source: §1 · pp. 1–2
• BIMIX jointly models domain proportions and training volume to predict and optimize mixed-domain pretraining outcomes.The framework mathematically relates these variables to model performance.Source: §1 · p. 2
• Experiments on two large-scale datasets show accurate performance prediction across varied mixtures and scaling scenarios.The experiments validate BIMIX across diverse datasets and training settings.Source: §1 · p. 2
• BIMIX enables domain-proportion optimization that outperforms existing high-cost methods in convergence speed and downstream task performance.Source: §1 · p. 2
• Entropy measures provide effective proxies for efficient mixture optimization and practical data-mixing guidance.Source: §1 · p. 2

• Manual mixture rules lack universal standards and portability, while iterative methods such as DoReMi and DoGE require considerable computational cost.These approaches optimize domain weights through reference or proxy-model training.Source: §2 · p. 2
• Entropy proxies offer a training-free alternative that can produce mixtures comparable or superior to existing methods.Source: §2 · p. 2
• Existing scaling laws primarily model averaged outcomes against individual scaling factors, overlooking domain-specific behavior in heterogeneous corpora.Source: §2 · pp. 2–3
• BIMIX jointly models domain proportion and data volume with a scalable formulation demonstrated for up to 22 domains.This distinguishes it from approaches whose complexity grows substantially with the number of domains.Source: §2 · p. 3

3 THE PROPOSED BIMIX

• Existing scaling laws overlook domain-specific predictability, while a domain’s training amount depends jointly on total data volume and its allocated proportion.Changing one proportion under fixed total volume also changes other domains’ allocations.Source: §3 · p. 3

3.1 FORMULATION

• BIMIX represents validation losses across m domains as functions of each domain proportion and training steps.The formulation uses five fitting coefficients per domain, giving linear complexity in the number of domains rather than quadratic complexity.Source: §3.1 · p. 3

3.2 OBSERVING SCALING BEHAVIORS BY DISENTANGLING VARIABLES

• Figure 1 compares observed validation losses with Eq. (3) fits as training steps increase at fixed domain proportions, using logarithmic axes.The reported curves closely fit observations and show consistent downward curvature.Source: §3.2 · p. 4, Figure 1 · p. 4
• At fixed training steps, validation loss varies approximately linearly with domain proportion in logarithmic space, supporting a power-law dependence.Figure 2 visualizes observed points and fitted results for this relationship.Source: §3.2 · p. 5, Figure 2 · p. 5
• Disentangling the two variables reveals separable scaling effects, which are combined into a mutually modulated bivariate law.Multiplicative shifts across curves motivate the integrated, interpretable formulation.Source: §3.2 · pp. 5–6

• Experiments use the 22-domain Pile and seven-domain SlimPajama datasets, with samples packed and chunked into 1024-token sequences.Source: §4 · p. 6
• The models are decoder-only transformers based on DoReMi, with 280M-scale and 1B-scale configurations.Both use the GPT-NeoX tokenizer.Source: §4 · p. 6
• Training runs span up to 200,000 update steps, approximately 100 billion tokens, with evaluations every 5 billion tokens.A single DoReMi round required 670 GPU hours on the stated infrastructure.Source: §4 · p. 6
• BIMIX fitting uses tuples of domain proportion, training steps, and domain loss, with up to m × n × k observations fitted by Trust Region Reflective optimization.Source: §4 · p. 6
• Evaluation reports per-domain validation loss and, when needed, one-shot Exact Match on WebQuestions, LAMBADA, and TriviaQA.Source: §4 · pp. 6–7
• Candidate mixtures include dataset baselines, DoReMi proportions, and entropy-derived alternatives.Source: §4 · p. 7
• Entropy-based mixtures normalize Shannon, conditional, joint, and von Neumann entropy across domains to produce proportions summing to one.Higher-entropy domains receive greater emphasis, and measurement adds negligible overhead during tokenization.Source: §4 · p. 7

5.1 EXTRAPOLATING LOSSES ON SCALED-UP DATA

• Mean relative error remains below 0.2% on both Pile and SlimPajama, while the worst domain error stays below 1.0%.The worst case is reported for the DoReMi mixture on the Pile.Source: §5.1 · p. 7

5.2 ESTIMATING DATA MIXTURES WITHOUT TRAINING

• BIMIX can evaluate candidate mixtures prospectively from limited mixture-training observations, helping prioritize configurations before costly large-scale training.The fitted law predicts validation losses for mixtures such as Baseline and DoReMi without directly training each candidate.Source: §5.2 · pp. 8–9
• High mean R2 values indicate that BIMIX generalizes across new data mixtures, with better fit on SlimPajama than on the Pile.The comparison follows the same dataset-level pattern observed in loss extrapolation.Source: §5.2 · p. 9, Table 2 · p. 9

5.3 OPTIMIZING DOMAIN PROPORTIONS FOR IMPROVED PERFORMANCE

• Optimized mixtures derived from smaller models give performance advantages throughout training for larger models, while RegMix achieves comparable performance on the reported comparison.The method also supplies a mathematical model of mixing behavior rather than only optimized proportions.Source: §5.3 · p. 9, Figure 4 · p. 10

5.4 ENTROPY MEASURES AS EFFICIENT MIXING PROXIES

• Entropy-driven mixtures achieve lower log-perplexity than the baseline, and conditional entropy consistently performs best among the tested proxies.These measures are used to collect fitting observations and support streamlined initial dataset construction.Source: §5.3 · p. 9, §5.4 · p. 10, Figure 5 · p. 10

6 DISCUSSION AND FUTURE WORK

• BIMIX’s applicability under extreme conditions is not guaranteed because experiments covered domain proportions from 0.0007 to 0.7256 and up to approximately 100B tokens.The study maintains domain inclusion during training and evaluation, while larger-scale predictability is suggested by prior work rather than established here.Source: §6 · p. 10
• Generalization to new domains depends on the model’s generalization capability and the correlation between new and training domains.The authors identify this extension beyond basic mixing laws as requiring dedicated exploration.Source: §6 · pp. 10–11
• BIMIX findings can assist computational resource allocation and support more economical and environmentally friendly AI development.The discussion also identifies multimodal data mixing as a future opportunity for improving training efficiency and performance.Source: §6 · p. 11

7 CONCLUSION

• BIMIX accurately models the joint scaling behavior of domain proportions and training volume, enabling loss extrapolation and generalization across different mixtures.The conclusion presents this modeling capability as the paper’s central contribution.Source: §7 · p. 11
• BIMIX optimization of domain proportions outperforms existing methods, while entropy-based measures provide efficient proxies for lightweight data mixing.The conclusion frames these results as practical tools for improving large-scale language-model training efficiency.Source: §7 · p. 11

A ENTROPY PROXIES

• Entropy proxies are computed from tokenized, equal-length domain sequences, while von Neumann entropy uses FastText embeddings and pairwise similarities.The preprocessing records unigram and bigram frequencies for the token-based measures.Source: A ENTROPY PROXIES · p. 15
• Shannon entropy measures expected information content from token frequencies and serves as an indicator of corpus diversity.For domain Di, P(x) denotes the probability of observing token x.Source: A ENTROPY PROXIES · p. 15
• Joint entropy measures uncertainty in consecutive token pairs and highlights sequential dependencies within a domain.The joint probability is estimated from observing token x followed by token x′ at adjacent positions.Source: A ENTROPY PROXIES · p. 15
• Conditional entropy measures surprise when predicting the next token given the current token, capturing text predictability and linguistic structure.It uses P(x′ | x) for the probability of observing x′ at position t + 1 given x at position t.Source: A ENTROPY PROXIES · p. 15
• Von Neumann entropy extends entropy to embedding-derived density matrices, with eigenvalues representing probability distributions analogous to quantum states.The density matrix is constructed from a positive semidefinite kernel over FastText embedding vectors, and entropy uses its trace and matrix logarithm.Source: A ENTROPY PROXIES · p. 15

B DOWNSTREAM TASK PERFORMANCE

• BIMIX-trained models consistently outperform other mixtures across most downstream tasks on the Pile and SlimPajama evaluations.Figures 6 and 7 report performance for 1B models trained with various data mixtures on the two datasets.Source: B DOWNSTREAM TASK PERFORMANCE · p. 16, Figure 6 · p. 16, Figure 7 · p. 16
• BIMIX mixtures maintain overall superiority on WebQuestions and LAMBADA despite some performance variability.Source: B DOWNSTREAM TASK PERFORMANCE · p. 16
• TriviaQA performance rises steadily with training iterations for all models, while BIMIX mixtures show clear advantages.Source: B DOWNSTREAM TASK PERFORMANCE · p. 16
• The downstream results support the effectiveness of BIMIX-optimized data mixtures across diverse tasks.Source: B DOWNSTREAM TASK PERFORMANCE · p. 16

C COMPLEXITY ANALYSIS

• BiMix incorporates both domain proportions and training steps, while the exponential law models domain proportions without training-step variation.Both laws model validation loss for multidomain language modeling.Source: C COMPLEXITY ANALYSIS · p. 17
• For an individual domain at fixed training steps, BiMix requires two coefficients, compared with m + 2 coefficients for the exponential law.The two BiMix coefficients are a scaling factor and an exponent.Source: C COMPLEXITY ANALYSIS · p. 17
• Across all m domains, BiMix maintains an order of magnitude fewer coefficients than the competing mixing law.Both methods scale their coefficient requirements with the number of domains, but BiMix remains substantially smaller.Source: C COMPLEXITY ANALYSIS · p. 17
• Across multiple training-step targets, the exponential law requires (m^2 + 2m)n coefficients, while BiMix shares coefficients across steps and requires 5m for all domains.BiMix needs five coefficients for one domain across different training steps.Source: C COMPLEXITY ANALYSIS · p. 17
• Fewer coefficients enable BiMix to be fit with potentially two candidate mixtures, whereas the exponential law requires tens of mixtures.The authors connect this reduction to lower computational, economic, and environmental resource use.Source: C COMPLEXITY ANALYSIS · p. 17

D COMPARISION WITH RECENT MIXING LAW

• The comparison included Ye et al.'s optimal SlimPajama mixture, while the exponential law faces scalability challenges on the 22-domain Pile dataset.The challenge is attributed to the law's previously discussed quadratic complexity.Source: D COMPARISION WITH RECENT MIXING LAW · p. 17
• BIMIX achieved Baseline-equivalent log-perplexity in 80,000 steps, compared with 160,000 steps for Ye et al.'s optimized mixture.Both optimized mixtures accelerated convergence relative to the default Baseline mixture.Source: D COMPARISION WITH RECENT MIXING LAW · p. 17, Figure 8 · p. 17
• The BIMIX-optimized mixture showed substantial downstream-task advantages, whereas Ye et al.'s optimized mixture performed comparably to the Baseline.The comparison concerns the 1B model trained on SlimPajama.Source: D COMPARISION WITH RECENT MIXING LAW · p. 18

E MIXTURE RECIPES

• Tables 4 and 5 report the candidate data-mixture compositions used for the Pile and SlimPajama experiments.These candidate mixtures form the basis of the BIMIX evaluation.Source: E MIXTURE RECIPES · p. 18, Table 4 · p. 18, Table 5 · p. 18