LR & Batch Size Scaling Law

from typing import List, Dict

# Group-specific scaling law coefficients
GROUP_COEFFS: Dict[str, Dict[str, float]] = {
    'all_data': {
        # lm_loss = a * lr^beta_lr * bsz^beta_bsz * data_size^beta_data * non_embedding_param_size^beta_param
        'a': 21.7693221886775,
        'beta_lr': 0.00863691905384939,
        'beta_bsz': -0.0005162836622543786,
        'beta_data': -0.04700957690670226,
        'beta_param': -0.05174150134631458,
    },
}

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,
               with different coefficients per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) under key 'lm_loss'.
    """
    if group not in GROUP_COEFFS:
        raise ValueError(f"Unknown group '{group}'")
    coeffs = GROUP_COEFFS[group]
    results: List[Dict[str, float]] = []
    for x in input_data:
        # compute predicted language modeling loss
        pred = (
            coeffs['a']
            * x['lr'] ** coeffs['beta_lr']
            * x['bsz'] ** coeffs['beta_bsz']
            * x['data_size'] ** coeffs['beta_data']
            * x['non_embedding_param_size'] ** coeffs['beta_param']
        )
        results.append({'lm_loss': pred})
    return results

"""
Scaling law function for predicting language modeling loss based on hyperparameters.
"""
import math

# Coefficients for each experimental group
COEFFICIENTS = {
    'all_data': {
        'intercept': 3.0805017396527683,
        'beta_lr': 0.00863691905384939,
        'beta_bsz': -0.0005162836622543786,
        'beta_data_size': -0.04700957690670226,
        'beta_param_size': -0.05174150134631458,
    },
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts the language modeling loss ('lm_loss') based on a power-law scaling
    relationship with hyperparameters: learning rate (lr), batch size (bsz),
    dataset size (data_size), and model size (non_embedding_param_size).

    Args:
        input_data: A list of dictionaries, each containing the keys
                    'lr', 'bsz', 'data_size', and 'non_embedding_param_size'.
        group: The experimental group name determining which coefficients to use.

    Returns:
        A list of dictionaries, each with a single key 'lm_loss' and its predicted value.
    """
    coeffs = COEFFICIENTS.get(group)
    if coeffs is None:
        raise ValueError(f"Unknown group: {group}")

    results = []
    for params in input_data:
        lr = params['lr']
        bsz = params['bsz']
        data_size = params['data_size']
        param_size = params['non_embedding_param_size']

        # Compute prediction in log-space
        log_pred = (
            coeffs['intercept']
            + coeffs['beta_lr'] * math.log(lr)
            + coeffs['beta_bsz'] * math.log(bsz)
            + coeffs['beta_data_size'] * math.log(data_size)
            + coeffs['beta_param_size'] * math.log(param_size)
        )
        # Convert back to linear space
        pred = math.exp(log_pred)
        results.append({'lm_loss': pred})

    return results

import math

_COEFFS = {
    # Group-specific coefficients: [intercept, lr, bsz, data_size, non_embedding_param_size]
    "all_data": {
        "intercept": 3.0805017396527683,
        "lr": 0.00863691905384939,
        "bsz": -0.0005162836622543786,
        "data_size": -0.04700957690670226,
        "non_embedding_param_size": -0.05174150134631458,
    },
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts lm_loss based on a power-law scaling of input hyperparameters.

    Args:
        input_data: A list of dicts with keys 'lr', 'bsz', 'data_size', and 'non_embedding_param_size'.
        group: Name of the experimental group. Must match one of the keys in _COEFFS.

    Returns:
        A list of dicts with key 'lm_loss' containing the predicted loss.
    """
    if group not in _COEFFS:
        raise ValueError(f"Unknown group '{group}'")
    params = _COEFFS[group]
    preds: list[dict[str, float]] = []
    for x in input_data:
        # Extract inputs
        lr = x["lr"]
        bsz = x["bsz"]
        data_size = x["data_size"]
        param_size = x["non_embedding_param_size"]
        # Compute log-scale prediction
        log_pred = params["intercept"]
        log_pred += params["lr"] * math.log(lr)
        log_pred += params["bsz"] * math.log(bsz)
        log_pred += params["data_size"] * math.log(data_size)
        log_pred += params["non_embedding_param_size"] * math.log(param_size)
        # Exponentiate to get back to original scale
        loss = math.exp(log_pred)
        preds.append({"lm_loss": loss})
    return preds

"""
Implements the discovered scaling law for language modeling loss.
"""
import math

# Pre-fitted coefficients for each experimental group
_COEFS = {
    'all_data': {
        # constant multiplier C
        'C': 21.76932218867723,
        # exponents for each input variable
        'exponents': {
            'lr': 0.00863692,
            'bsz': -0.00051628,
            'data_size': -0.04700958,
            'non_embedding_param_size': -0.0517415,
        }
    }
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts lm_loss according to the scaling law:
        lm_loss = C * lr^alpha * bsz^beta * data_size^gamma * non_embedding_param_size^delta

    Args:
        input_data: list of dicts with keys 'lr', 'bsz', 'data_size', 'non_embedding_param_size'.
        group: experimental group name. Must be one of the pre-fitted groups.

    Returns:
        List of dicts with key 'lm_loss' for each input.
    """
    params = _COEFS.get(group)
    if params is None:
        raise ValueError(f"Unknown group: {group}")
    C = params['C']
    exps = params['exponents']
    outputs: list[dict[str, float]] = []
    for point in input_data:
        # compute scaled loss
        y = C
        for var, exp in exps.items():
            val = point.get(var)
            if val is None:
                raise KeyError(f"Missing input variable: {var}")
            y *= val ** exp
        outputs.append({'lm_loss': y})
    return outputs