Vocabulary Scaling Law

import math

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, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s), here:
        'unigram_normalized_loss'.
    """
    # Coefficients for each group
    # Currently only one group is supported
    if group != 'all_data':
        raise ValueError(f"Unknown group: {group}")
    # Fitted parameters from log-linear regression
    intercept = 6.380591236628991
    beta_vocab = 0.06340183374111474
    beta_non_vocab = 0.016411064426657424
    beta_num_chars = -0.5017006627222854

    results: list[dict[str, float]] = []
    for item in input_data:
        # Natural log transform of inputs
        ln_vocab = math.log(item['vocab_size'])
        ln_non_vocab = math.log(item['non_vocab_parameters'])
        ln_num_chars = math.log(item['num_characters'])
        # Linear model in log-space
        pred = (
            intercept
            + beta_vocab * ln_vocab
            + beta_non_vocab * ln_non_vocab
            + beta_num_chars * ln_num_chars
        )
        results.append({'unigram_normalized_loss': pred})
    return results

"""
Module implementing the discovered scaling law for unigram-normalized loss.
"""
import math

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts unigram_normalized_loss based on input variables using a
    log-linear scaling law: loss = A + B1*ln(vocab_size)
    + B2*ln(non_vocab_parameters) + B3*ln(num_characters).

    Args:
        input_data: List of dicts with keys 'vocab_size',
                    'non_vocab_parameters', 'num_characters'.
        group: Experimental group name. Coefficients are group-specific.

    Returns:
        List of dicts with key 'unigram_normalized_loss' and predicted value.
    """
    # Pre-fitted coefficients per group
    params = {
        'all_data': {
            'intercept': 6.380591236628991,
            'coef_vocab': 0.06340183374111474,
            'coef_non_vocab': 0.016411064426657424,
            'coef_chars': -0.5017006627222854,
        }
    }
    if group not in params:
        raise ValueError(f"Unknown group: {group}")
    p = params[group]
    results = []
    for d in input_data:
        v = d.get('vocab_size')
        nv = d.get('non_vocab_parameters')
        ch = d.get('num_characters')
        loss = (
            p['intercept']
            + p['coef_vocab'] * math.log(v)
            + p['coef_non_vocab'] * math.log(nv)
            + p['coef_chars'] * math.log(ch)
        )
        results.append({'unigram_normalized_loss': loss})
    return results

import math

__all__ = ['law']

# Precomputed parameters for each experimental group
_PARAMS = {
    'all_data': {
        'theta0': 6.380591236628991,
        'alpha': 0.06340183374111474,
        'beta': 0.016411064426657424,
        'gamma': -0.5017006627222854,
    },
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts unigram_normalized_loss based on vocabulary size,
    non-vocabulary parameters, and number of characters using
    a log-linear scaling law:
        loss = theta0 + alpha*ln(vocab_size)
               + beta*ln(non_vocab_parameters)
               + gamma*ln(num_characters)

    Args:
        input_data: List of data points, each with keys:
                    'vocab_size', 'non_vocab_parameters', 'num_characters'.
        group: Experimental group name; must be one of _PARAMS keys.

    Returns:
        List of dicts with key 'unigram_normalized_loss' for each input.
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group '{group}'. Available groups: {list(_PARAMS.keys())}")
    params = _PARAMS[group]
    results = []
    for d in input_data:
        vs = d['vocab_size']
        nv = d['non_vocab_parameters']
        nc = d['num_characters']
        loss = (
            params['theta0']
            + params['alpha'] * math.log(vs)
            + params['beta'] * math.log(nv)
            + params['gamma'] * math.log(nc)
        )
        results.append({'unigram_normalized_loss': loss})
    return results

import math

_COEFFICIENTS = {
    'all_data': {
        'intercept': 6.380591236628994,
        'log_chars': -0.5017006627222831,
        'log_vocab': 0.06340183374111423,
        'log_params': 0.016411064426654527,
    }
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts unigram_normalized_loss based on scaling law.

    Args:
        input_data: List of dicts with keys 'vocab_size',
                    'non_vocab_parameters', 'num_characters'.
        group: Experimental group name.

    Returns:
        List of dicts with key 'unigram_normalized_loss'.
    """
    if group not in _COEFFICIENTS:
        raise ValueError(f"Unknown group '{group}'")
    coeffs = _COEFFICIENTS[group]
    results = []
    for entry in input_data:
        nc = float(entry['num_characters'])
        vs = float(entry['vocab_size'])
        pv = float(entry['non_vocab_parameters'])
        pred = (
            coeffs['intercept']
            + coeffs['log_chars'] * math.log(nc)
            + coeffs['log_vocab'] * math.log(vs)
            + coeffs['log_params'] * math.log(pv)
        )
        results.append({'unigram_normalized_loss': pred})
    return results