Data-Constrained Scaling Law

"""
Defines the scaling law function for predicting validation loss
based on model size, training tokens, and dataset unique tokens.
"""

# Coefficients for each experimental group, fitted via log-linear regression
_COEFS = {
    'all_data': {
        'A': 89.03635820053546,
        'alpha': -0.06713156032896106,
        'beta': -0.05741837292779783,
        'gamma': -0.02821632111651308,
    },
}

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, each containing:
            - 'params': model parameter count
            - 'tokens': total number of training tokens
            - 'unique_tokens': number of unique tokens in dataset
        group: Experimental group name; must match a key in the fitted coefficients.

    Returns:
        A list of dictionaries with key 'loss' for each predicted validation loss.
    """
    if group not in _COEFS:
        raise ValueError(f"Unknown group '{group}'. Available groups: {list(_COEFS.keys())}")
    coeffs = _COEFS[group]
    predictions = []
    for entry in input_data:
        p = entry['params']
        t = entry['tokens']
        u = entry['unique_tokens']
        # power-law formula: loss = A * params^alpha * tokens^beta * unique_tokens^gamma
        loss_pred = (
            coeffs['A']
            * (p ** coeffs['alpha'])
            * (t ** coeffs['beta'])
            * (u ** coeffs['gamma'])
        )
        predictions.append({'loss': loss_pred})
    return predictions

from typing import List, Dict

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

    Returns:
        A list of dictionaries with each dictionary containing the predicted
        output variable 'loss' for the corresponding input data point.
    """
    # Coefficients for each experimental group
    coeffs = {
        'all_data': {
            'A': 89.03635820053537,
            'alpha': 0.06713156032896134,
            'beta': 0.057418372927797716,
            'gamma': -0.02821632111651312,
        }
    }
    if group not in coeffs:
        raise ValueError(f"Unknown group: {group}")
    c = coeffs[group]
    results: List[Dict[str, float]] = []
    for x in input_data:
        params = x['params']
        tokens = x['tokens']
        unique_tokens = x['unique_tokens']
        # Scaling law: loss = A * params^{-alpha} * tokens^{-beta} * unique_tokens^{gamma}
        loss = (
            c['A']
            * (params ** (-c['alpha']))
            * (tokens ** (-c['beta']))
            * (unique_tokens ** (c['gamma']))
        )
        results.append({'loss': loss})
    return results

import math

# Coefficients per experimental group for the power-law model:
_COEFFS = {
    'all_data': {
        # log-space intercept and exponents
        'b0': 4.489044805418067,
        'b1': -0.06713156032896134,
        'b2': -0.057418372927797716,
        'b3': -0.02821632111651312,
    }
}

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, each containing 'params', 'tokens', and 'unique_tokens'.
        group: The experimental group name for which to apply group-specific coefficients.

    Returns:
        A list of dictionaries with the predicted 'loss' for each input point.
    """
    if group not in _COEFFS:
        raise ValueError(f"Unknown group: {group}")
    coeffs = _COEFFS[group]
    b0, b1, b2, b3 = coeffs['b0'], coeffs['b1'], coeffs['b2'], coeffs['b3']
    results = []
    for point in input_data:
        p = point['params']
        t = point['tokens']
        u = point['unique_tokens']
        # loss = exp(b0) * params**b1 * tokens**b2 * unique_tokens**b3
        loss = math.exp(b0) * (p ** b1) * (t ** b2) * (u ** b3)
        results.append({'loss': loss})
    return results

"""Scaling law function for model validation loss."""
from math import exp

_COEFFS = {
    'all_data': {
        'a': 89.03635820053537,
        'alpha': 0.06713156032896134,
        'beta': 0.057418372927797716,
        'gamma': 0.02821632111651312,
    }
}

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, each with keys 'params', 'tokens', 'unique_tokens'.
        group: The name of the experimental group for which to make predictions.

    Returns:
        A list of dictionaries with the predicted 'loss'.
    """
    if group not in _COEFFS:
        raise ValueError(f"Unknown group: {group}")
    coeffs = _COEFFS[group]
    a = coeffs['a']
    alpha = coeffs['alpha']
    beta = coeffs['beta']
    gamma = coeffs['gamma']

    results = []
    for d in input_data:
        p = d['params']
        t = d['tokens']
        u = d['unique_tokens']
        loss_pred = a * (p ** (-alpha)) * (t ** (-beta)) * (u ** (-gamma))
        results.append({'loss': loss_pred})
    return results