scaling-laws-optimal-data-mixtures

name	scaling-laws-optimal-data-mixtures
title	Scaling Laws for Optimal Data Mixtures
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2507.09404
keywords	["Scaling Laws","Data Mixtures","Multi-Domain Training","Compute Optimization","Foundation Models"]
description	Predict optimal data mixture proportions for multi-domain LLM training using scaling laws that require only 10-20 small experiments. Determine which domains should contribute how much data across model sizes (186M-7B), reducing computational waste in foundation model pretraining.

Scaling Laws for Optimal Data Mixtures: Efficient Multi-Domain Foundation Model Training

Foundation models trained on diverse domains (code, math, text) must decide how much data from each domain to include. Current practice uses trial-and-error, wasting compute on suboptimal mixtures. This work proposes scaling laws that predict model loss as a function of model size (N), training tokens (D), and domain weights (h), enabling practitioners to determine optimal mixtures efficiently.

The key insight is that only 10-20 small-scale training runs are needed to fit accurate scaling laws that extrapolate to much larger models. Two law formulations are provided: additive (mixture-independent scale parameters, predicting scale-independent optimal mixtures) and joint (mixture-dependent scale parameters, predicting compute-budget-dependent optima). Both formulations have been validated across language models, multimodal models, and vision models.

Core Concept

Multi-domain training loss can be modeled as: E + 1/∑(C_i·h_i^γ_i) + A/N^α + B/D^β, where only the bias term E depends on domain weights h. This reveals a crucial property: for additive scaling laws, optimal domain weights are compute-independent—they don't change with model size or total training budget. Only the absolute performance scales.

For joint scaling laws where scale parameters A and B themselves depend on mixture composition, optimal weights become budget-dependent. The fitting procedure uses Basin-hopping with L-BFGS to navigate the high-dimensional parameter space efficiently.

Architecture Overview

• Additive Scaling Law Formulation: Loss depends on domain-weight-dependent bias term; scale parameters (A, B) are mixture-independent
• Joint Scaling Law Formulation: Both scale parameters and bias depend on mixture; accounts for scale-mixture interactions
• Loss Prediction Pipeline: Fit scaling laws to small experiments, extrapolate to large models/compute budgets
• Optimization Framework: Basin-hopping + L-BFGS to solve mixture optimization problem over weight simplex
• Multi-Model Validation: Tests on language models (186M-7B), multimodal (image+text), and vision encoders

Implementation

Additive Scaling Law: Mixture-Independent Scale Parameters

Model loss where domain weights affect bias term but not scaling.

import numpy as np
from scipy.optimize import minimize, basinhopping
import torch

class AdditiveScalingLaw:
    

     ():
        .num_domains = num_domains

        
        .C = np.ones(num_domains)  
        .gamma = np.ones(num_domains)  
        .A =   
        .alpha =   
        .B =   
        .beta =   

     ():
        
        
        bias = 
        weighted_sum = 

         i  (.num_domains):
            weighted_sum += .C[i] * (mixture_weights[i] ** .gamma[i])

        mixture_term =  / (weighted_sum + )
        bias += mixture_term

        
        scale_loss = .A / (model_size ** .alpha) + .B / (tokens ** .beta)

         bias + scale_loss

     ():
        
         ():
            
            C = params[:.num_domains]
            gamma = params[.num_domains:*.num_domains]
            A = params[*.num_domains]
            alpha = params[*.num_domains + ]
            B = params[*.num_domains + ]
            beta = params[*.num_domains + ]

            
            .C = C
            .gamma = gamma
            .A = A
            .alpha = alpha
            .B = B
            .beta = beta

            
            mse = 
             exp  experiments:
                predicted = .predict_loss(
                    exp[],
                    exp[],
                    exp[]
                )
                mse += (predicted - exp[]) ** 

             mse / (experiments)

        
        x0 = np.concatenate([
            .C,
            .gamma,
            [.A, .alpha, .B, .beta]
        ])

        
        minimizer_kwargs = {
            : ,
            : [(, )] * (*.num_domains) + [(, ), (, ), (, ), (, )]
        }

        result = basinhopping(
            objective,
            x0,
            minimizer_kwargs=minimizer_kwargs,
            niter=,
            seed=
        )

         result.fun

     () -> np.ndarray:
        
         ():
            
              np.isclose(weights.(), ):
                 

            
            weighted_sum = 
             i  (.num_domains):
                weighted_sum += .C[i] * (weights[i] ** .gamma[i])

            
            loss =  / (weighted_sum + )

             loss

        
         scipy.optimize  LinearConstraint

        A_constraint = np.ones((, .num_domains))
        constraint = LinearConstraint(A_constraint, [], [])

        
        x0 = np.ones(.num_domains) / .num_domains

        result = minimize(
            objective_mixture,
            x0,
            method=,
            bounds=[(, )  _  (.num_domains)]
        )

        optimal = result.x / result.x.()  

         optimal