SLD - Domain Mixture Scaling Law

All Runs (sorted by R²)

Best Run 1 R² = 0.974767

▼

Python

from typing import List, Dict

# Parameters for scaling law: loss(p) = c - a * p**b
_PARAMS: Dict[str, Dict[int, Dict[str, float]]] = {
    '70M': {
        1: {'a': 0.9228, 'b': 0.2453, 'c': 3.4149},
        2: {'a': 0.3726, 'b': 0.5065, 'c': 3.8184},
        3: {'a': 0.7930, 'b': 0.2212, 'c': 3.6006},
        4: {'a': 0.9436, 'b': 0.2406, 'c': 2.2663},
        5: {'a': 0.5175, 'b': 0.3754, 'c': 3.9317},
    },
    '160M': {
        1: {'a': 0.8432, 'b': 0.2285, 'c': 3.0604},
        2: {'a': 0.3059, 'b': 0.4616, 'c': 3.4721},
        3: {'a': 0.7277, 'b': 0.2081, 'c': 3.2856},
        4: {'a': 0.8371, 'b': 0.2382, 'c': 1.9631},
        5: {'a': 0.5291, 'b': 0.3623, 'c': 3.5949},
    },
    '305M': {
        1: {'a': 0.8159, 'b': 0.2234, 'c': 2.8980},
        2: {'a': 0.4262, 'b': 0.6940, 'c': 3.3062},
        3: {'a': 0.7023, 'b': 0.1831, 'c': 3.1556},
        4: {'a': 0.7988, 'b': 0.2365, 'c': 1.8330},
        5: {'a': 0.5343, 'b': 0.3516, 'c': 3.4344},
    },
    '410M': {
        1: {'a': 0.7997, 'b': 0.2158, 'c': 2.8319},
        2: {'a': 0.3518, 'b': 0.6247, 'c': 3.2303},
        3: {'a': 0.7099, 'b': 0.1805, 'c': 3.0983},
        4: {'a': 0.7849, 'b': 0.2413, 'c': 1.7794},
        5: {'a': 0.5501, 'b': 0.3404, 'c': 3.3746},
    },
}

def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
               The functional form of the law is the same for all groups,
               but the constant parameters/coefficients differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) under keys
        'loss_domain_i'.
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group: {group}")
    group_params = _PARAMS[group]
    results: List[Dict[str, float]] = []
    # Compute prediction for each data point
    for entry in input_data:
        preds: Dict[str, float] = {}
        for i in range(1, 6):
            p = entry.get(f'proportion_domain_{i}')
            if p is None:
                raise KeyError(f"Missing proportion_domain_{i} in input data")
            a = group_params[i]['a']
            b = group_params[i]['b']
            c = group_params[i]['c']
            # scaling law: loss = c - a * p**b
            preds[f'loss_domain_{i}'] = c - a * (p ** b)
        results.append(preds)
    return results

#2 Run 2 R² = 0.974745

▼

Python

from typing import List, Dict

def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts validation losses for each domain based on mixture proportions.

    Uses a power-law decrease model per domain and group:
        loss = L0 - alpha * proportion**beta

    Args:
        input_data: list of dicts with keys 'proportion_domain_i' for i in 1..5.
        group: experimental group name ('70M', '160M', '305M', '410M').

    Returns:
        List of dicts with keys 'loss_domain_i' predictions.
    """
    # Pre-fitted parameters per group and domain
    params = {
        '70M': {
            'L0':   [3.415, 3.818, 3.601, 2.266, 3.932],
            'alpha':[0.923, 0.373, 0.793, 0.944, 0.518],
            'beta': [0.245, 0.506, 0.221, 0.241, 0.375],
        },
        '160M': {
            'L0':   [3.060, 3.472, 3.286, 1.963, 3.595],
            'alpha':[0.843, 0.306, 0.728, 0.837, 0.529],
            'beta': [0.229, 0.462, 0.208, 0.238, 0.362],
        },
        '305M': {
            'L0':   [2.898, 3.306, 3.156, 1.833, 3.434],
            'alpha':[0.816, 0.426, 0.702, 0.799, 0.534],
            'beta': [0.223, 0.694, 0.183, 0.236, 0.352],
        },
        '410M': {
            'L0':   [2.832, 3.230, 3.098, 1.779, 3.375],
            'alpha':[0.800, 0.352, 0.710, 0.785, 0.550],
            'beta': [0.216, 0.625, 0.181, 0.241, 0.340],
        },
    }
    if group not in params:
        raise ValueError(f"Unknown group '{group}'")
    gp = params[group]
    results: List[Dict[str, float]] = []
    for row in input_data:
        pred: Dict[str, float] = {}
        for i in range(1, 6):
            p = row.get(f'proportion_domain_{i}', 0.0)
            L0 = gp['L0'][i-1]
            a = gp['alpha'][i-1]
            b = gp['beta'][i-1]
            pred[f'loss_domain_{i}'] = L0 - a * (p ** b)
        results.append(pred)
    return results

#3 Run 3 R² = 0.971579

▼

Python

from typing import List, Dict

# Precomputed parameters for the scaling law per group and domain
_PARAMS: Dict[str, Dict[str, Dict[str, float]]] = {
    '70M': {
        'domain_1': {'a': -0.9196, 'b': 0.2428, 'c': 3.4151},
        'domain_2': {'a': -0.2861, 'b': 0.3259, 'c': 3.8184},
        'domain_3': {'a': -0.8845, 'b': 0.2583, 'c': 3.6004},
        'domain_4': {'a': -0.9339, 'b': 0.2354, 'c': 2.2665},
        'domain_5': {'a': -0.5144, 'b': 0.3494, 'c': 3.9349},
    },
    '160M': {
        'domain_1': {'a': -0.8416, 'b': 0.2271, 'c': 3.0605},
        'domain_2': {'a': -0.2291, 'b': 0.2533, 'c': 3.4721},
        'domain_3': {'a': -0.8212, 'b': 0.2491, 'c': 3.2853},
        'domain_4': {'a': -0.8337, 'b': 0.2360, 'c': 1.9631},
        'domain_5': {'a': -0.5281, 'b': 0.3531, 'c': 3.5961},
    },
    '305M': {
        'domain_1': {'a': -0.8126, 'b': 0.2206, 'c': 2.8982},
        'domain_2': {'a': -0.2314, 'b': 0.2867, 'c': 3.3062},
        'domain_3': {'a': -0.7806, 'b': 0.2189, 'c': 3.1554},
        'domain_4': {'a': -0.7939, 'b': 0.2333, 'c': 1.8330},
        'domain_5': {'a': -0.5326, 'b': 0.3372, 'c': 3.4364},
    },
    '410M': {
        'domain_1': {'a': -0.7973, 'b': 0.2137, 'c': 2.8320},
        'domain_2': {'a': -0.1942, 'b': 0.2113, 'c': 3.2303},
        'domain_3': {'a': -0.7895, 'b': 0.2163, 'c': 3.0981},
        'domain_4': {'a': -0.7782, 'b': 0.2369, 'c': 1.7795},
        'domain_5': {'a': -0.5493, 'b': 0.3332, 'c': 3.3756},
    },
}

def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts per-domain validation loss based on mixture proportions using a power-law model.

    The functional form is:
        loss_domain_i = a * (proportion_domain_i ** b) + c

    Args:
        input_data: List of data points, each with keys 'proportion_domain_i'.
        group: Experimental group key, must be one of the precomputed groups.

    Returns:
        List of dictionaries with predicted 'loss_domain_i' for each input point.
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group '{group}', available: {list(_PARAMS.keys())}")
    params_group = _PARAMS[group]
    predictions: List[Dict[str, float]] = []
    for point in input_data:
        pred_point: Dict[str, float] = {}
        for i in range(1, 6):
            key_p = f'proportion_domain_{i}'
            key_l = f'loss_domain_{i}'
            p = point.get(key_p, 0.0)
            a = params_group[f'domain_{i}']['a']
            b = params_group[f'domain_{i}']['b']
            c = params_group[f'domain_{i}']['c']
            pred_point[key_l] = a * (p ** b) + c
        predictions.append(pred_point)
    return predictions

#4 Run 4 R² = 0.845752

▼

Python

from typing import List, Dict

def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts domain losses based on mixture proportions using a linear scaling law.

    Args:
        input_data: List of dicts with keys 'proportion_domain_i'.
        group: Experimental group identifier ('70M', '160M', '305M', '410M').

    Returns:
        List of dicts with predicted 'loss_domain_i' values.
    """
    # Coefficients from linear regression: loss = a + b * proportion
    coefs = {
        '70M': {
            1: (3.1806393062811864, -1.2285149798567838),
            2: (3.801253923348644, -0.8613783672356026),
            3: (3.5566666364458976, -5.9220601237842345),
            4: (2.092666552495471, -2.465721152611366),
            5: (3.7712578339232965, -0.4114998066410033),
        },
        '160M': {
            1: (2.835641724877016, -1.125053454723479),
            2: (3.4544639388150142, -0.7759431938007308),
            3: (3.243292438159674, -5.633806849579989),
            4: (1.8086663359915747, -2.2029709658709042),
            5: (3.430618715766743, -0.4253675130276537),
        },
        '305M': {
            1: (2.6765745376811547, -1.0852461829640883),
            2: (3.290256899590436, -0.7424407706785378),
            3: (3.1100613611088073, -5.787743320347023),
            4: (1.684474243911805, -2.103498009394809),
            5: (3.2619262695047944, -0.4221154025882421),
        },
        '410M': {
            1: (2.609996558045746, -1.0646721511212967),
            2: (3.2130253217275206, -0.7041613598604771),
            3: (3.0516738117194397, -5.89283215213306),
            4: (1.6354678769114128, -2.0519200859166604),
            5: (3.194356949340593, -0.4345763744862935),
        },
    }
    if group not in coefs:
        raise ValueError(f"Unknown group: {group}")
    params = coefs[group]
    predictions: List[Dict[str, float]] = []
    for point in input_data:
        out: Dict[str, float] = {}
        for i in range(1, 6):
            p = point.get(f'proportion_domain_{i}', 0.0)
            a, b = params[i]
            out[f'loss_domain_{i}'] = a + b * p
        predictions.append(out)
    return predictions