MoE Scaling Law

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, each containing 'num_experts' and
                    'dense_parameter_count' values.
        group: The experimental group name for parameter lookup.

    Returns:
        A list of dictionaries with predicted 'loss_validation' values.
    """
    # Scaling law form:
    #   loss_validation = a * num_experts^alpha * dense_parameter_count^beta
    GROUP_PARAMS = {
        'all_data': {
            'a': 10.069179203081301,
            'alpha': -0.027401922527553482,
            'beta': -0.07051248607961777,
        },
        # Additional groups can be added here
    }
    if group not in GROUP_PARAMS:
        raise ValueError(f"Unknown group: {group}")
    params = GROUP_PARAMS[group]
    a = params['a']
    alpha = params['alpha']
    beta = params['beta']

    results: List[Dict[str, float]] = []
    for d in input_data:
        E = d.get('num_experts', 0.0)
        P = d.get('dense_parameter_count', 0.0)
        # Compute predicted validation loss using the power-law
        loss_pred = a * (E ** alpha) * (P ** beta)
        results.append({'loss_validation': loss_pred})
    return results

"""
Module implementing the discovered scaling law for Mixture-of-Experts (MoE) models.
"""
import math

# Coefficients for each experimental group
_COEFFS = {
    "all_data": {
        # loss_validation = A * dense_parameter_count^alpha * num_experts^beta
        "A": 10.069179203081301,
        "alpha": -0.07051248607961783,
        "beta": -0.027401922527553444,
    },
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts validation loss based on a power-law scaling with dense and expert parameters.

    Args:
        input_data: List of dicts with keys 'dense_parameter_count' and 'num_experts'.
        group: Experimental group name. Must match one of the keys in _COEFFS.

    Returns:
        List of dicts with key 'loss_validation' for each input point.
    """
    if group not in _COEFFS:
        raise ValueError(f"Unknown group: {group}")
    params = _COEFFS[group]
    A = params["A"]
    alpha = params["alpha"]
    beta = params["beta"]
    outputs: list[dict[str, float]] = []
    for entry in input_data:
        dense = entry.get("dense_parameter_count")
        experts = entry.get("num_experts")
        # compute prediction via power law
        pred = A * (dense ** alpha) * (experts ** beta)
        outputs.append({"loss_validation": pred})
    return outputs

import math
from typing import List, Dict

_COEFS: Dict[str, Dict[str, float]] = {
    # coefficients for each experimental group
    'all_data': {
        # scale factor a, exponent p for num_experts, exponent q for dense parameters
        'a': 10.069179203081296,
        'p': -0.0274019225275534,
        'q': -0.0705124860796179,
    },
}

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) under key
        'loss_validation'.
    """
    if group not in _COEFS:
        raise ValueError(f"Unknown group '{group}'. Available groups: {list(_COEFS.keys())}")
    params = _COEFS[group]
    a = params['a']
    p = params['p']
    q = params['q']
    predictions: List[Dict[str, float]] = []
    for datum in input_data:
        E = datum.get('num_experts')
        D = datum.get('dense_parameter_count')
        # scaling law: loss = a * E^p * D^q
        loss_pred = a * (E ** p) * (D ** q)
        predictions.append({'loss_validation': 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, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) ("loss_validation").
    """
    # Coefficients for each group: A * E^a * D^b
    coeffs = {
        'all_data': {
            'A': 10.06917920308131,
            'a': -0.02740192252755355,
            'b': -0.07051248607961802,
        }
    }
    if group not in coeffs:
        raise ValueError(f"Unknown group: {group}")
    params = coeffs[group]
    results: List[Dict[str, float]] = []
    for datum in input_data:
        E = datum['num_experts']
        D = datum['dense_parameter_count']
        # Scaling law: loss = A * E^a * D^b
        loss_pred = params['A'] * (E ** params['a']) * (D ** params['b'])
        results.append({'loss_validation': loss_pred})
    return results