SLD - Domain Mixture Scaling Law

All Runs (sorted by R²)

Best Run 1 R² = 0.902224

▼

Python

from math import log

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."""
    EPS = 1e-06
    COEFFS = {'160M': {'loss_domain_1': {'a': 2.3900239026095846, 'b': -0.0495482846415819},
          'loss_domain_2': {'a': 3.3089973473750915, 'b': -0.011847031339016876},
          'loss_domain_3': {'a': 2.7750612243262176, 'b': -0.03703323688495387},
          'loss_domain_4': {'a': 1.3585325213280746, 'b': -0.04416553521395201},
          'loss_domain_5': {'a': 3.1416357474005943, 'b': -0.03640943279009211}},
 '305M': {'loss_domain_1': {'a': 2.2445101169537405, 'b': -0.048262539632469084},
          'loss_domain_2': {'a': 3.1516644411817394, 'b': -0.011228102438934438},
          'loss_domain_3': {'a': 2.627226581493613, 'b': -0.03831519342481427},
          'loss_domain_4': {'a': 1.254081110356706, 'b': -0.042287102605490595},
          'loss_domain_5': {'a': 2.974256274101213, 'b': -0.03681502204362127}},
 '410M': {'loss_domain_1': {'a': 2.183986016421472, 'b': -0.04779882382417481},
          'loss_domain_2': {'a': 3.0802842179137, 'b': -0.010883486046165937},
          'loss_domain_3': {'a': 2.5599133716532463, 'b': -0.039034580115764335},
          'loss_domain_4': {'a': 1.216103723749701, 'b': -0.04115598457974984},
          'loss_domain_5': {'a': 2.898019537149145, 'b': -0.038035573443952336}},
 '70M': {'loss_domain_1': {'a': 2.6991877691852117, 'b': -0.05301830256954046},
         'loss_domain_2': {'a': 3.6412456318395994, 'b': -0.012884469474280064},
         'loss_domain_3': {'a': 3.064818062764666, 'b': -0.038877352813920324},
         'loss_domain_4': {'a': 1.5883184714718916, 'b': -0.04953747817222867},
         'loss_domain_5': {'a': 3.491442851657023, 'b': -0.035415726158268684}}}
    # Determine coeff set for group; fallback to first available group if not found
    group_coeffs = COEFFS.get(group)
    if group_coeffs is None:
        if COEFFS:
            group_coeffs = COEFFS[sorted(COEFFS.keys())[0]]
        else:
            group_coeffs = {}
    outputs = []
    for row in input_data:
        out = {}
        # For each loss key we know, compute using corresponding proportion
        for loss_key, ab in group_coeffs.items():
            # Infer proportion key by replacing loss_ with proportion_ in the key name
            prop_key = loss_key.replace('loss_', 'proportion_')
            p = float(row.get(prop_key, 0.0))
            a = float(ab.get('a', 0.0))
            b = float(ab.get('b', 0.0))
            pred = a + b * log(p + EPS)
            out[loss_key] = float(pred)
        outputs.append(out)
    return outputs

#2 Run 2 R² = 0.880748

▼

Python

from __future__ import annotations
from typing import List, Dict
EPS = 1e-06
PARAMS = {
  "70M": {
    "domain_1": {
      "a": 2.0220104987927257,
      "b": 0.7051338858670713,
      "alpha": 0.05
    },
    "domain_2": {
      "a": 3.478057239579727,
      "b": 0.17092206236157978,
      "alpha": 0.05
    },
    "domain_3": {
      "a": 2.59254285900522,
      "b": 0.5058051486225027,
      "alpha": 0.05
    },
    "domain_4": {
      "a": 0.9657272567112907,
      "b": 0.6545542978911262,
      "alpha": 0.05
    },
    "domain_5": {
      "a": 3.024670907672095,
      "b": 0.47441444506610303,
      "alpha": 0.05
    }
  },
  "160M": {
    "domain_1": {
      "a": 1.7563840140531553,
      "b": 0.6595784179449191,
      "alpha": 0.05
    },
    "domain_2": {
      "a": 3.1588187107783474,
      "b": 0.15725443050159807,
      "alpha": 0.05
    },
    "domain_3": {
      "a": 2.3251342372559862,
      "b": 0.48184980806156213,
      "alpha": 0.05
    },
    "domain_4": {
      "a": 0.8034472891804151,
      "b": 0.583579912689387,
      "alpha": 0.05
    },
    "domain_5": {
      "a": 2.661628446598834,
      "b": 0.487853772515834,
      "alpha": 0.05
    }
  },
  "305M": {
    "domain_1": {
      "a": 1.6270228147234236,
      "b": 0.6426830441112606,
      "alpha": 0.05
    },
    "domain_2": {
      "a": 3.0093884970464932,
      "b": 0.14899746540873754,
      "alpha": 0.05
    },
    "domain_3": {
      "a": 2.1615746228783017,
      "b": 0.4986333285709002,
      "alpha": 0.05
    },
    "domain_4": {
      "a": 0.7225593285828501,
      "b": 0.5587932312372937,
      "alpha": 0.05
    },
    "domain_5": {
      "a": 2.4882874593817577,
      "b": 0.49385627236422996,
      "alpha": 0.05
    }
  },
  "410M": {
    "domain_1": {
      "a": 1.5720931564397607,
      "b": 0.6367651459833956,
      "alpha": 0.05
    },
    "domain_2": {
      "a": 2.9422498586244026,
      "b": 0.14451569098864436,
      "alpha": 0.05
    },
    "domain_3": {
      "a": 2.0855052306255972,
      "b": 0.5080046213375579,
      "alpha": 0.05
    },
    "domain_4": {
      "a": 0.6988734728738473,
      "b": 0.5437907496465303,
      "alpha": 0.05
    },
    "domain_5": {
      "a": 2.395591013031327,
      "b": 0.5105511796136086,
      "alpha": 0.05
    }
  }
}

def _predict_one(sample: Dict[str, float], group: str) -> Dict[str, float]:
    if group not in PARAMS:
        raise ValueError(f"Unknown group: {group}")
    out: Dict[str, float] = {}
    gparams = PARAMS[group]
    for key, coeffs in gparams.items():
        # key looks like 'domain_i' from fitting stage
        d_idx = ''.join(ch for ch in key if ch.isdigit())
        p_key = f'proportion_domain_{d_idx}'
        y_key = f'loss_domain_{d_idx}'
        p = float(sample.get(p_key, 0.0))
        a = coeffs["a"]; b = coeffs["b"]; alpha = coeffs["alpha"]
        x = (p + EPS) ** (-alpha) if alpha != 0.0 else p
        y = a + b * x
        out[y_key] = float(y)
    return out

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: List of dicts, each containing at least proportion_domain_i keys.
        group: Name of the experimental group for which to make predictions.
    Returns:
        List of dicts with predicted loss_domain_i for i=1..5 (as available).
    """
    return [_predict_one(sample, group) for sample in input_data]

#3 Run 3 R² = 0.873244

▼

Python

from __future__ import annotations
import math

# Single functional form across all groups:
# L_i = A_{g,i} + B_{g,i} * log(p_i) with clipping p_i -> max(p_i, 1e-9)
COEFS = {
  "160M": {
    "0": [
      2.4222245349880027,
      -0.03132683285487464
    ],
    "1": [
      3.317675549031874,
      -0.007472875400866049
    ],
    "2": [
      2.816384638275675,
      -0.02268262544964051
    ],
    "3": [
      1.3940507054184053,
      -0.027665347428548207
    ],
    "4": [
      3.150894901274346,
      -0.023224558453472387
    ]
  },
  "305M": {
    "0": [
      2.2757879996895203,
      -0.03052740882148673
    ],
    "1": [
      3.1599100833879104,
      -0.007079726627925801
    ],
    "2": [
      2.6699102690836165,
      -0.02347551348122601
    ],
    "3": [
      1.2880743105690569,
      -0.02649081104120577
    ],
    "4": [
      2.9835648450347945,
      -0.023515715591365914
    ]
  },
  "410M": {
    "0": [
      2.214856577818689,
      -0.030250619653697765
    ],
    "1": [
      3.0882309384163564,
      -0.006868468175494051
    ],
    "2": [
      2.6033922128748963,
      -0.023916963654028817
    ],
    "3": [
      1.2492138560279822,
      -0.025778353486428242
    ],
    "4": [
      2.907602130692238,
      -0.024316234178981166
    ]
  },
  "70M": {
    "0": [
      2.733906922836672,
      -0.033479990654170164
    ],
    "1": [
      3.6507355235720613,
      -0.008120459606531918
    ],
    "2": [
      3.108227590265197,
      -0.02380902449916114
    ],
    "3": [
      1.6281637129068052,
      -0.03102931954295701
    ],
    "4": [
      3.5004670876544726,
      -0.022579959393990854
    ]
  }
}

def _predict_log(p: float, A: float, B: float) -> float:
    if p <= 0.0:
        p = 1e-9
    return A + B * math.log(p)

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."""
    g = group if group in COEFS else (next(iter(COEFS.keys())) if COEFS else group)
    out_list = []
    for row in input_data:
        out = {}
        for i in range(5):
            pc = f"proportion_domain_{i+1}"
            lc = f"loss_domain_{i+1}"
            p = float(row.get(pc, 0.0))
            A,B = COEFS[g].get(str(i), [0.0, 0.0])
            out[lc] = float(_predict_log(p, A, B))
        out_list.append(out)
    return out_list

#4 Run 4 R² = 0.852129

▼

Python

from typing import List, Dict

_PARAMS = {
  "160M": {
    "1": {
      "a": 0.4777718023660796,
      "b": 1.9651028132012958,
      "alpha": 0.01
    },
    "2": {
      "a": 2.853979648026824,
      "b": 0.46913461080072405,
      "alpha": 0.01
    },
    "3": {
      "a": 1.4416433577788708,
      "b": 1.3992301424068143,
      "alpha": 0.01
    },
    "4": {
      "a": -0.31200145102844623,
      "b": 1.7281252544418098,
      "alpha": 0.01
    },
    "5": {
      "a": 1.6988970818463138,
      "b": 1.458160686737703,
      "alpha": 0.01
    }
  },
  "305M": {
    "1": {
      "a": 0.380355700173621,
      "b": 1.9155046248750955,
      "alpha": 0.01
    },
    "2": {
      "a": 2.720727222497648,
      "b": 0.4443470722501475,
      "alpha": 0.01
    },
    "3": {
      "a": 1.2468028929418933,
      "b": 1.4484164085003954,
      "alpha": 0.01
    },
    "4": {
      "a": -0.34563740615899546,
      "b": 1.6548398159870399,
      "alpha": 0.01
    },
    "5": {
      "a": 1.5118433832803935,
      "b": 1.4779286413154575,
      "alpha": 0.01
    }
  },
  "410M": {
    "1": {
      "a": 0.3358801860263219,
      "b": 1.8988049175912776,
      "alpha": 0.01
    },
    "2": {
      "a": 2.6618934312385854,
      "b": 0.43132199370582547,
      "alpha": 0.01
    },
    "3": {
      "a": 1.153496029332803,
      "b": 1.475677952459072,
      "alpha": 0.01
    },
    "4": {
      "a": -0.3403973481604199,
      "b": 1.6101852758013921,
      "alpha": 0.01
    },
    "5": {
      "a": 1.3848194753051277,
      "b": 1.5291803128611205,
      "alpha": 0.01
    }
  },
  "70M": {
    "1": {
      "a": 0.6576096958255063,
      "b": 2.098519217677623,
      "alpha": 0.01
    },
    "2": {
      "a": 3.14714983348316,
      "b": 0.5095246197489386,
      "alpha": 0.01
    },
    "3": {
      "a": 1.6653425516285103,
      "b": 1.4686043699822906,
      "alpha": 0.01
    },
    "4": {
      "a": -0.2852957436802397,
      "b": 1.9382200618934766,
      "alpha": 0.01
    },
    "5": {
      "a": 2.089249143779692,
      "b": 1.4172201529137438,
      "alpha": 0.01
    }
  }
}

def _get_params(group: str):
    return _PARAMS.get(group, _PARAMS.get(str(group), None)) or next(iter(_PARAMS.values()))

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 parameters can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    params = _get_params(group)
    outputs = []
    for row in input_data:
        out = {}
        for i in range(1, 6):
            p = float(row.get('proportion_domain_%d' % i, 0.0))
            if p <= 0:
                p = 1e-12
            dom = params[str(i)]
            a = dom['a']; b = dom['b']; alpha = dom['alpha']
            out['loss_domain_%d' % i] = float(a + b * (p ** (-alpha)))
        outputs.append(out)
    return outputs

#5 Run 5 R² = -1.000000

▼

Python

from __future__ import annotations
from typing import List, Dict
import math

# Embedded fitted parameters per group and domain.
# Each domain uses either a power form: y = a + k * p**c
# or a log form: y = a + b * log(p)
PARAMS = {'160M': {'1': {'a': 3.358131379887285,
                'c': 0.15000000000000036,
                'form': 'power',
                'k': -1.1267042149203725,
                'r2': 0.9941862658661204},
          '2': {'a': 3.318780734454316,
                'c': -3.0,
                'form': 'power',
                'k': 2.2734784837905097e-05,
                'r2': 0.06062614424441071},
          '3': {'a': 2.931566771454378, 'c': 3.0, 'form': 'power', 'k': -174.62340395796082, 'r2': 0.7061503010752743},
          '4': {'a': 3.8716421496206315,
                'c': 0.050000000000000266,
                'form': 'power',
                'k': -2.6874577135715296,
                'r2': 0.9762826201636258},
          '5': {'a': 3.876397587476255,
                'c': 0.20000000000000018,
                'form': 'power',
                'k': -0.8015051387810923,
                'r2': 0.995953951272856}},
 '305M': {'1': {'a': 4.832122416331462,
                'c': 0.050000000000000266,
                'form': 'power',
                'k': -2.721490588848711,
                'r2': 0.9923678438528458},
          '2': {'a': 3.1584886677226094,
                'c': -2.75,
                'form': 'power',
                'k': 4.513541219029818e-05,
                'r2': 0.066949255913345},
          '3': {'a': 2.7848933137166916, 'c': 3.0, 'form': 'power', 'k': -166.12136907423988, 'r2': 0.6853980994244002},
          '4': {'a': -0.9540330670713865,
                'c': -0.04999999999999982,
                'form': 'power',
                'k': 2.067152378143916,
                'r2': 0.9773900180368574},
          '5': {'a': 5.599492209154898,
                'c': 0.050000000000000266,
                'form': 'power',
                'k': -2.6836119274365684,
                'r2': 0.9937674933260627}},
 '410M': {'1': {'a': 4.670978316018143,
                'c': 0.050000000000000266,
                'form': 'power',
                'k': -2.613131878961905,
                'r2': 0.9885683848277665},
          '2': {'a': 3.091279676924703,
                'c': -3.0,
                'form': 'power',
                'k': 1.7491883130243175e-05,
                'r2': 0.03333320137761886},
          '3': {'a': 2.7201830236984295, 'c': 3.0, 'form': 'power', 'k': -168.02135230322037, 'r2': 0.6489941549529987},
          '4': {'a': 0.5427391497008389,
                'c': -0.1499999999999999,
                'form': 'power',
                'k': 0.5526416189568344,
                'r2': 0.9732551074881751},
          '5': {'a': 3.634036329340057,
                'c': 0.20000000000000018,
                'form': 'power',
                'k': -0.8017750604189218,
                'r2': 0.9939499265653914}},
 '70M': {'1': {'a': 4.297380991024045,
               'c': 0.10000000000000009,
               'form': 'power',
               'k': -1.775827036691083,
               'r2': 0.9952894285760332},
         '2': {'a': 3.646005195575224,
               'c': -2.6,
               'form': 'power',
               'k': 8.526182739288986e-05,
               'r2': 0.12259002392235419},
         '3': {'a': 3.2284514219048646, 'c': 3.0, 'form': 'power', 'k': -180.79876453488362, 'r2': 0.6915505904622448},
         '4': {'a': 1.1444240089156006, 'c': -0.25, 'form': 'power', 'k': 0.3211850398410141, 'r2': 0.9856867978166857},
         '5': {'a': 2.3099069373584133,
               'c': -0.09999999999999964,
               'form': 'power',
               'k': 1.1287948701911399,
               'r2': 0.9930440593027846}}}

# Default fallback params (averaged) if group not present
FALLBACK = {'1': {'a': 4.289653275815233, 'c': 0.08750000000000024, 'form': 'power', 'k': -2.059288429855518},
 '2': {'a': 3.3036385686692133, 'c': -2.8375, 'form': 'power', 'k': 4.265597688783408e-05},
 '3': {'a': 2.9162736326935907, 'c': 3.0, 'form': 'power', 'k': -172.39122246757617},
 '4': {'a': 1.151193060291421, 'c': -0.09999999999999987, 'form': 'power', 'k': 0.06338033084255867},
 '5': {'a': 3.8549582658324058, 'c': 0.08750000000000024, 'form': 'power', 'k': -0.7895243141113606}}

DOMAINS = [1,2,3,4,5]

def _predict_domain(p: float, spec: dict) -> float:
    # Guard against non-positive proportions for log and power
    p = max(p, 1e-12)
    form = spec.get('form')
    if form == 'power':
        return spec['a'] + spec['k'] * (p ** spec['c'])
    elif form == 'log':
        return spec['a'] + spec['b'] * math.log(p)
    else:
        # Should not happen; fall back to identity-like
        return float('nan')

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: list of dicts with keys proportion_domain_1..5
        group: group name. Functional form is same across groups (per-domain chosen from {power, log}); coefficients differ per group.

    Returns:
        list of dicts with keys loss_domain_1..5
    """
    # Select group-specific params or fallback
    spec = PARAMS.get(group, FALLBACK)
    outputs = []
    for row in input_data:
        out = {}
        for i in DOMAINS:
            p = float(row.get(f'proportion_domain_{i}', 0.0))
            out[f'loss_domain_{i}'] = _predict_domain(p, spec[str(i)])
        outputs.append(out)
    return outputs