LR & Batch Size Scaling Law

import json
import os


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 must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """

    # Load parameters
    params_path = os.path.join(os.path.dirname(__file__), 'params.json')
    with open(params_path, 'r') as f:
        params_dict = json.load(f)

    # Get parameters for the specified group
    if group not in params_dict:
        raise ValueError(f"Unknown group: {group}")

    params = params_dict[group]
    a = params['a']
    b = params['b']
    c = params['c']
    d = params['d']
    e = params['e']

    # Apply the power law model to each input data point
    results = []
    for data_point in input_data:
        lr = data_point['lr']
        bsz = data_point['bsz']
        data_size = data_point['data_size']
        non_embedding_param_size = data_point['non_embedding_param_size']

        # lm_loss = a * lr^b * bsz^c * data_size^d * non_embedding_param_size^e
        lm_loss = a * (lr ** b) * (bsz ** c) * (data_size ** d) * (non_embedding_param_size ** e)

        results.append({'lm_loss': lm_loss})

    return results

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 must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """

    # Model parameters discovered through log-space linear regression
    # Formula: lm_loss = exp(a) * lr^b_lr * bsz^b_bsz * data_size^b_data_size * non_embedding_param_size^b_non_embedding_param_size

    # Group-specific parameters (currently only one group in the dataset)
    group_params = {
        'all_data': {
            'a': 3.080501739652756,
            'b_lr': 0.008636919053849421,
            'b_bsz': -0.0005162836622544659,
            'b_data_size': -0.04700957690670223,
            'b_non_embedding_param_size': -0.05174150134631419
        }
    }

    # Get parameters for the specified group
    if group not in group_params:
        # If group not found, use 'all_data' as fallback
        params = group_params['all_data']
    else:
        params = group_params[group]

    # Extract parameters
    a = params['a']
    b_lr = params['b_lr']
    b_bsz = params['b_bsz']
    b_data_size = params['b_data_size']
    b_non_embedding_param_size = params['b_non_embedding_param_size']

    # Pre-compute exp(a)
    import math
    exp_a = math.exp(a)

    # Make predictions for each data point
    results = []
    for data_point in input_data:
        lr = data_point['lr']
        bsz = data_point['bsz']
        data_size = data_point['data_size']
        non_embedding_param_size = data_point['non_embedding_param_size']

        # Apply the power-law formula
        lm_loss = exp_a * (lr ** b_lr) * (bsz ** b_bsz) * (data_size ** b_data_size) * (non_embedding_param_size ** b_non_embedding_param_size)

        results.append({'lm_loss': lm_loss})

    return results

import numpy as np
from typing import Callable

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.

    The underlying mathematical model is a power-law relationship in log-space:

    log10(lm_loss) = 1.3378449070
                     + 0.0086369191 * log10(lr)
                     - 0.0005162837 * log10(bsz)
                     - 0.0470095769 * log10(data_size)
                     - 0.0517415013 * log10(param_size)

    Equivalently in linear scale:
    lm_loss = 21.7693221887 * lr^0.0086369191 * bsz^-0.0005162837
              * data_size^-0.0470095769 * param_size^-0.0517415013

    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. Expected keys: 'lr', 'bsz', 'data_size',
                    'non_embedding_param_size'
        group: The name of the experimental group for which to make predictions.
               The coefficients are the same for all groups ('all_data' in the
               discovered dataset), but this parameter is provided for consistency
               with the API specification.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable 'lm_loss'.
    """

    # Model coefficients discovered through linear regression on log-transformed data
    intercept = 1.3378449070
    coef_lr = 0.0086369191
    coef_bsz = -0.0005162837
    coef_data_size = -0.0470095769
    coef_param_size = -0.0517415013

    results = []

    for data_point in input_data:
        # Extract input variables
        lr = data_point['lr']
        bsz = data_point['bsz']
        data_size = data_point['data_size']
        param_size = data_point['non_embedding_param_size']

        # Compute log10(loss) using the linear model in log-space
        log_loss = (
            intercept
            + coef_lr * np.log10(lr)
            + coef_bsz * np.log10(bsz)
            + coef_data_size * np.log10(data_size)
            + coef_param_size * np.log10(param_size)
        )

        # Convert back to linear scale
        predicted_loss = 10 ** log_loss

        # Return prediction as a dictionary
        results.append({'lm_loss': float(predicted_loss)})

    return results

import numpy as np


# Fitted coefficients for each group
# These were determined by fitting a polynomial power law model in log-space
# Model: log(lm_loss) = a*log(lr) + b*log(bsz) + c*log(D) + d*log(N) + e*log(D)² + f*log(N)² + intercept
# where D = data_size, N = non_embedding_param_size

COEFFICIENTS = {
    'all_data': {
        'log_lr': 0.009718,
        'log_bsz': 0.000187,
        'log_data_size': -0.372180,
        'log_param_size': 0.058584,
        'log_data_size_sq': 0.006943,
        'log_param_size_sq': -0.002747,
        'intercept': 5.775274,
    }
}


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 must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    # Get the coefficients for this group
    if group not in COEFFICIENTS:
        # Fallback to 'all_data' if group is not found
        coeff = COEFFICIENTS.get('all_data')
        if coeff is None:
            raise ValueError(f"No coefficients found for group '{group}'")
    else:
        coeff = COEFFICIENTS[group]

    results = []

    for data_point in input_data:
        # Extract input variables
        lr = data_point.get('lr', 0.001)
        bsz = data_point.get('bsz', 256)
        data_size = data_point.get('data_size', 1e11)
        param_size = data_point.get('non_embedding_param_size', 1e8)

        # Take logarithms of input variables
        log_lr = np.log(lr)
        log_bsz = np.log(bsz)
        log_data_size = np.log(data_size)
        log_param_size = np.log(param_size)

        # Apply the polynomial power law formula in log-space
        log_lm_loss = (
            coeff['log_lr'] * log_lr +
            coeff['log_bsz'] * log_bsz +
            coeff['log_data_size'] * log_data_size +
            coeff['log_param_size'] * log_param_size +
            coeff['log_data_size_sq'] * (log_data_size ** 2) +
            coeff['log_param_size_sq'] * (log_param_size ** 2) +
            coeff['intercept']
        )

        # Convert back from log-space
        lm_loss = np.exp(log_lm_loss)

        # Create output dictionary
        results.append({'lm_loss': float(lm_loss)})

    return results