SFT Scaling Law

"""
Scaling law function for predicting SFT loss based on data size.
"""
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 coefficients differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) as 'sft_loss'.
    """
    # Fitted power-law parameters per group: sft_loss = a * (sft_data_size ** b)
    _params = {
        "('MBZUAI/LaMini-GPT-124M', 'flan')": {'a': 7.465597659576905, 'b': -0.08938336391209896},
        "('MBZUAI/LaMini-GPT-124M', 'gigaword')": {'a': 6.969626540790517, 'b': -0.12488797569428042},
        "('MBZUAI/LaMini-GPT-124M', 'wikiword')": {'a': 3.8453873317979776, 'b': -0.06600249911783039},
        "('MBZUAI/LaMini-GPT-774M', 'flan')": {'a': 5.713040837756609, 'b': -0.07829225824344388},
        "('MBZUAI/LaMini-GPT-774M', 'gigaword')": {'a': 6.293121529736756, 'b': -0.12918481189914824},
        "('MBZUAI/LaMini-GPT-774M', 'wikiword')": {'a': 2.8960123299335456, 'b': -0.05447544270876587},
        "('cerebras/Cerebras-GPT-1.3B', 'flan')": {'a': 3.763368646463403, 'b': -0.05245645350446021},
        "('cerebras/Cerebras-GPT-1.3B', 'gigaword')": {'a': 4.8571354747287865, 'b': -0.09575030044686328},
        "('cerebras/Cerebras-GPT-1.3B', 'wikiword')": {'a': 3.1982974522931946, 'b': -0.05122695523416853},
        "('cerebras/Cerebras-GPT-256M', 'flan')": {'a': 4.582026537858244, 'b': -0.05133628273093655},
        "('cerebras/Cerebras-GPT-256M', 'gigaword')": {'a': 5.787318731177163, 'b': -0.11196737789112267},
        "('cerebras/Cerebras-GPT-256M', 'wikiword')": {'a': 4.508753213335996, 'b': -0.0704142021738492},
        "('facebook/bart-base', 'flan')": {'a': 7.161143632325512, 'b': -0.09164738334291277},
        "('facebook/bart-base', 'gigaword')": {'a': 9.296091035384052, 'b': -0.1581352715189432},
        "('facebook/bart-base', 'wikiword')": {'a': 5.82331151864561, 'b': -0.10583613940732714},
        "('facebook/bart-large', 'flan')": {'a': 5.239565986675214, 'b': -0.07673442669563073},
        "('facebook/bart-large', 'gigaword')": {'a': 9.506911791025203, 'b': -0.16937059578718577},
        "('facebook/bart-large', 'wikiword')": {'a': 3.022916994149884, 'b': -0.06221765532278332},
        "('facebook/opt-1.3b', 'flan')": {'a': 3.242895597483687, 'b': -0.049961389646430265},
        "('facebook/opt-1.3b', 'gigaword')": {'a': 5.693457761695671, 'b': -0.11822788318303022},
        "('facebook/opt-1.3b', 'wikiword')": {'a': 2.352359108086735, 'b': -0.04191838265818277},
        "('facebook/opt-350m', 'flan')": {'a': 4.585856367154473, 'b': -0.06062077349418788},
        "('facebook/opt-350m', 'gigaword')": {'a': 7.476891875498453, 'b': -0.1403389747361835},
        "('facebook/opt-350m', 'wikiword')": {'a': 3.250071982062783, 'b': -0.05575365697201477},
        "('facebook/opt-6.7b', 'flan')": {'a': 2.2344284064674027, 'b': -0.019171706023120263},
        "('facebook/opt-6.7b', 'gigaword')": {'a': 2.180823669815048, 'b': -0.014692781292003678},
        "('facebook/opt-6.7b', 'wikiword')": {'a': 2.0379653707032412, 'b': -0.034542865866457675},
        "('google/mt5-base', 'flan')": {'a': 4.621169816481478, 'b': -0.06486990722366648},
        "('google/mt5-base', 'gigaword')": {'a': 3.4542367430363625, 'b': -0.032132757710130744},
        "('google/mt5-base', 'wikiword')": {'a': 4.880269855685557, 'b': -0.09617770187230981},
        "('google/mt5-large', 'flan')": {'a': 3.5193429654174015, 'b': -0.05367728546091697},
        "('google/mt5-large', 'gigaword')": {'a': 3.5889984641635846, 'b': -0.038154542411887996},
        "('google/mt5-large', 'wikiword')": {'a': 3.9537507864011348, 'b': -0.07896963709680072},
        "('gpt2', 'flan')": {'a': 7.743248668603319, 'b': -0.09034898764189149},
        "('gpt2', 'gigaword')": {'a': 7.448123349328367, 'b': -0.1368351251564653},
        "('gpt2', 'wikiword')": {'a': 4.005381596620439, 'b': -0.06983526356123387},
        "('t5-base', 'flan')": {'a': 3.5927162123671015, 'b': -0.05385301204930485},
        "('t5-base', 'gigaword')": {'a': 1.7460064580179184, 'b': -0.08169520675626497},
        "('t5-base', 'wikiword')": {'a': 2.2734524247637107, 'b': -0.04525490715783883},
        "('t5-small', 'flan')": {'a': 4.10316214565362, 'b': -0.054114550232136904},
        "('t5-small', 'gigaword')": {'a': 1.8635831534568412, 'b': -0.07576270233096472},
        "('t5-small', 'wikiword')": {'a': 2.8159260594072437, 'b': -0.05183604484401076},
    }
    if group not in _params:
        raise ValueError(f"Unknown group '{group}'. Available groups: {list(_params.keys())}")
    a = _params[group]['a']
    b = _params[group]['b']
    results: List[Dict[str, float]] = []
    for entry in input_data:
        n = entry.get('sft_data_size')
        if n is None:
            raise KeyError("Each input_data dict must contain 'sft_data_size'.")
        pred = a * (n ** b)
        results.append({'sft_loss': pred})
    return results

import math

_PARAMS = {
    "('MBZUAI/LaMini-GPT-124M', 'flan')": {"A": 7.465597659576905, "m": -0.08938336391209896},
    "('MBZUAI/LaMini-GPT-124M', 'gigaword')": {"A": 6.969626540790517, "m": -0.12488797569428042},
    "('MBZUAI/LaMini-GPT-124M', 'wikiword')": {"A": 3.8453873317979776, "m": -0.06600249911783039},
    "('MBZUAI/LaMini-GPT-774M', 'flan')": {"A": 5.713040837756609, "m": -0.07829225824344388},
    "('MBZUAI/LaMini-GPT-774M', 'gigaword')": {"A": 6.293121529736756, "m": -0.12918481189914824},
    "('MBZUAI/LaMini-GPT-774M', 'wikiword')": {"A": 2.8960123299335456, "m": -0.05447544270876587},
    "('cerebras/Cerebras-GPT-1.3B', 'flan')": {"A": 3.763368646463403, "m": -0.05245645350446021},
    "('cerebras/Cerebras-GPT-1.3B', 'gigaword')": {"A": 4.8571354747287865, "m": -0.09575030044686328},
    "('cerebras/Cerebras-GPT-1.3B', 'wikiword')": {"A": 3.1982974522931946, "m": -0.05122695523416853},
    "('cerebras/Cerebras-GPT-256M', 'flan')": {"A": 4.582026537858244, "m": -0.05133628273093655},
    "('cerebras/Cerebras-GPT-256M', 'gigaword')": {"A": 5.787318731177163, "m": -0.11196737789112267},
    "('cerebras/Cerebras-GPT-256M', 'wikiword')": {"A": 4.508753213335996, "m": -0.0704142021738492},
    "('facebook/bart-base', 'flan')": {"A": 7.161143632325512, "m": -0.09164738334291277},
    "('facebook/bart-base', 'gigaword')": {"A": 9.296091035384052, "m": -0.1581352715189432},
    "('facebook/bart-base', 'wikiword')": {"A": 5.82331151864561, "m": -0.10583613940732714},
    "('facebook/bart-large', 'flan')": {"A": 5.239565986675214, "m": -0.07673442669563073},
    "('facebook/bart-large', 'gigaword')": {"A": 9.506911791025203, "m": -0.16937059578718577},
    "('facebook/bart-large', 'wikiword')": {"A": 3.022916994149884, "m": -0.06221765532278332},
    "('facebook/opt-1.3b', 'flan')": {"A": 3.242895597483687, "m": -0.049961389646430265},
    "('facebook/opt-1.3b', 'gigaword')": {"A": 5.693457761695671, "m": -0.11822788318303022},
    "('facebook/opt-1.3b', 'wikiword')": {"A": 2.352359108086735, "m": -0.04191838265818277},
    "('facebook/opt-350m', 'flan')": {"A": 4.585856367154473, "m": -0.06062077349418788},
    "('facebook/opt-350m', 'gigaword')": {"A": 7.476891875498453, "m": -0.1403389747361835},
    "('facebook/opt-350m', 'wikiword')": {"A": 3.250071982062783, "m": -0.05575365697201477},
    "('facebook/opt-6.7b', 'flan')": {"A": 2.2344284064674027, "m": -0.019171706023120263},
    "('facebook/opt-6.7b', 'gigaword')": {"A": 2.180823669815048, "m": -0.014692781292003678},
    "('facebook/opt-6.7b', 'wikiword')": {"A": 2.0379653707032412, "m": -0.034542865866457675},
    "('google/mt5-base', 'flan')": {"A": 4.621169816481478, "m": -0.06486990722366648},
    "('google/mt5-base', 'gigaword')": {"A": 3.4542367430363625, "m": -0.032132757710130744},
    "('google/mt5-base', 'wikiword')": {"A": 4.880269855685557, "m": -0.09617770187230981},
    "('google/mt5-large', 'flan')": {"A": 3.5193429654174015, "m": -0.05367728546091697},
    "('google/mt5-large', 'gigaword')": {"A": 3.5889984641635846, "m": -0.038154542411887996},
    "('google/mt5-large', 'wikiword')": {"A": 3.9537507864011348, "m": -0.07896963709680072},
    "('gpt2', 'flan')": {"A": 7.743248668603319, "m": -0.09034898764189149},
    "('gpt2', 'gigaword')": {"A": 7.448123349328367, "m": -0.1368351251564653},
    "('gpt2', 'wikiword')": {"A": 4.005381596620439, "m": -0.06983526356123387},
    "('t5-base', 'flan')": {"A": 3.5927162123671015, "m": -0.05385301204930485},
    "('t5-base', 'gigaword')": {"A": 1.7460064580179184, "m": -0.08169520675626497},
    "('t5-base', 'wikiword')": {"A": 2.2734524247637107, "m": -0.04525490715783883},
    "('t5-small', 'flan')": {"A": 4.10316214565362, "m": -0.054114550232136904},
    "('t5-small', 'gigaword')": {"A": 1.8635831534568412, "m": -0.07576270233096472},
    "('t5-small', 'wikiword')": {"A": 2.8159260594072437, "m": -0.05183604484401076},
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts sft_loss based on sft_data_size using a power-law scaling:   
        sft_loss = A * (sft_data_size ** m)   
    with parameters A and m specific to each group.

    Args:
        input_data: List of dicts with key 'sft_data_size'.
        group: Experimental group identifier matching dataset groups.

    Returns:
        List of dicts with predicted 'sft_loss'.
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group: {group}")
    A = _PARAMS[group]["A"]
    m = _PARAMS[group]["m"]
    results = []
    for datapoint in input_data:
        x = datapoint.get('sft_data_size')
        if x is None:
            raise KeyError("Input data must contain 'sft_data_size'.")
        y_pred = A * (x ** m)
        results.append({"sft_loss": y_pred})
    return results

"""
Scaling law predictor for Supervised Fine-Tuning (SFT) loss.
"""
# Discovered power-law form: loss = a * (data_size ** b), per group
_params = {
    "('MBZUAI/LaMini-GPT-124M', 'flan')": {"a": 7.465598, "b": -0.089383},
    "('MBZUAI/LaMini-GPT-124M', 'gigaword')": {"a": 6.969627, "b": -0.124888},
    "('MBZUAI/LaMini-GPT-124M', 'wikiword')": {"a": 3.845387, "b": -0.066002},
    "('MBZUAI/LaMini-GPT-774M', 'flan')": {"a": 5.713041, "b": -0.078292},
    "('MBZUAI/LaMini-GPT-774M', 'gigaword')": {"a": 6.293122, "b": -0.129185},
    "('MBZUAI/LaMini-GPT-774M', 'wikiword')": {"a": 2.896012, "b": -0.054475},
    "('cerebras/Cerebras-GPT-1.3B', 'flan')": {"a": 3.763369, "b": -0.052456},
    "('cerebras/Cerebras-GPT-1.3B', 'gigaword')": {"a": 4.857135, "b": -0.095750},
    "('cerebras/Cerebras-GPT-1.3B', 'wikiword')": {"a": 3.198297, "b": -0.051227},
    "('cerebras/Cerebras-GPT-256M', 'flan')": {"a": 4.582027, "b": -0.051336},
    "('cerebras/Cerebras-GPT-256M', 'gigaword')": {"a": 5.787319, "b": -0.111967},
    "('cerebras/Cerebras-GPT-256M', 'wikiword')": {"a": 4.508753, "b": -0.070414},
    "('facebook/bart-base', 'flan')": {"a": 7.161144, "b": -0.091647},
    "('facebook/bart-base', 'gigaword')": {"a": 9.296091, "b": -0.158135},
    "('facebook/bart-base', 'wikiword')": {"a": 5.823312, "b": -0.105836},
    "('facebook/bart-large', 'flan')": {"a": 5.239566, "b": -0.076734},
    "('facebook/bart-large', 'gigaword')": {"a": 9.506912, "b": -0.169371},
    "('facebook/bart-large', 'wikiword')": {"a": 3.022917, "b": -0.062218},
    "('facebook/opt-1.3b', 'flan')": {"a": 3.242896, "b": -0.049961},
    "('facebook/opt-1.3b', 'gigaword')": {"a": 5.693458, "b": -0.118228},
    "('facebook/opt-1.3b', 'wikiword')": {"a": 2.352359, "b": -0.041918},
    "('facebook/opt-350m', 'flan')": {"a": 4.585856, "b": -0.060621},
    "('facebook/opt-350m', 'gigaword')": {"a": 7.476892, "b": -0.140339},
    "('facebook/opt-350m', 'wikiword')": {"a": 3.250072, "b": -0.055754},
    "('facebook/opt-6.7b', 'flan')": {"a": 2.234428, "b": -0.019172},
    "('facebook/opt-6.7b', 'gigaword')": {"a": 2.180824, "b": -0.014693},
    "('facebook/opt-6.7b', 'wikiword')": {"a": 2.037965, "b": -0.034543},
    "('google/mt5-base', 'flan')": {"a": 4.621170, "b": -0.064870},
    "('google/mt5-base', 'gigaword')": {"a": 3.454237, "b": -0.032133},
    "('google/mt5-base', 'wikiword')": {"a": 4.880270, "b": -0.096178},
    "('google/mt5-large', 'flan')": {"a": 3.519343, "b": -0.053677},
    "('google/mt5-large', 'gigaword')": {"a": 3.588998, "b": -0.038155},
    "('google/mt5-large', 'wikiword')": {"a": 3.953751, "b": -0.078970},
    "('gpt2', 'flan')": {"a": 7.743249, "b": -0.090349},
    "('gpt2', 'gigaword')": {"a": 7.448123, "b": -0.136835},
    "('gpt2', 'wikiword')": {"a": 4.005382, "b": -0.069835},
    "('t5-base', 'flan')": {"a": 3.592716, "b": -0.053853},
    "('t5-base', 'gigaword')": {"a": 1.746006, "b": -0.081695},
    "('t5-base', 'wikiword')": {"a": 2.273452, "b": -0.045255},
    "('t5-small', 'flan')": {"a": 4.103162, "b": -0.054115},
    "('t5-small', 'gigaword')": {"a": 1.863583, "b": -0.075763},
    "('t5-small', 'wikiword')": {"a": 2.815926, "b": -0.051836},
}

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts SFT loss based on data size using a power-law scaling law.

    Args:
        input_data: List of data points, each with key 'sft_data_size'.
        group: Experimental group specifying which coefficients to use.

    Returns:
        List of dictionaries with predicted 'sft_loss'.
    """
    if group not in _params:
        raise ValueError(f"Unknown group: {group}")
    coeffs = _params[group]
    a = coeffs["a"]
    b = coeffs["b"]
    predictions: list[dict[str, float]] = []
    for point in input_data:
        n = point.get("sft_data_size")
        loss = a * (n ** b)
        predictions.append({"sft_loss": loss})
    return predictions

"""
Module implementing the discovered scaling law for supervised fine-tuning loss.
"""

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 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 'sft_loss'.
    """
    # Parameters per group: L_inf + a * N^{-alpha}
    params = {
        "('MBZUAI/LaMini-GPT-124M', 'flan')": {'L_inf': 4.88891522852415e-13, 'a': 6.973889417936588, 'alpha': 0.08134160664016567},
        "('MBZUAI/LaMini-GPT-124M', 'gigaword')": {'L_inf': 1.2549475345759155e-15, 'a': 5.923742908062525, 'alpha': 0.10520859443461071},
        "('MBZUAI/LaMini-GPT-124M', 'wikiword')": {'L_inf': 6.73534819503915e-22, 'a': 3.7815959443424148, 'alpha': 0.06407311292143066},
        "('MBZUAI/LaMini-GPT-774M', 'flan')": {'L_inf': 3.0986536858492428e-15, 'a': 5.416194447728044, 'alpha': 0.07203582303793636},
        "('MBZUAI/LaMini-GPT-774M', 'gigaword')": {'L_inf': 1.5568069794903192e-15, 'a': 5.350693782433369, 'alpha': 0.10952387398189671},
        "('MBZUAI/LaMini-GPT-774M', 'wikiword')": {'L_inf': 1.5817591620080455e-18, 'a': 2.8815067646791404, 'alpha': 0.05390255281176658},
        "('cerebras/Cerebras-GPT-1.3B', 'flan')": {'L_inf': 4.660261124868628e-19, 'a': 3.72668747103668, 'alpha': 0.05133622852901864},
        "('cerebras/Cerebras-GPT-1.3B', 'gigaword')": {'L_inf': 1.4572929400202255e-21, 'a': 4.601802095259647, 'alpha': 0.08937942187812241},
        "('cerebras/Cerebras-GPT-1.3B', 'wikiword')": {'L_inf': 5.166756221644035e-20, 'a': 3.1670522055347803, 'alpha': 0.05010482255812582},
        "('cerebras/Cerebras-GPT-256M', 'flan')": {'L_inf': 5.681581108258014e-16, 'a': 4.497775771102446, 'alpha': 0.049197297920243487},
        "('cerebras/Cerebras-GPT-256M', 'gigaword')": {'L_inf': 5.792914647146667e-14, 'a': 5.237428498166979, 'alpha': 0.10004126991906952},
        "('cerebras/Cerebras-GPT-256M', 'wikiword')": {'L_inf': 1.3546553850129447e-24, 'a': 4.456019626802189, 'alpha': 0.06905741523821077},
        "('facebook/bart-base', 'flan')": {'L_inf': 4.485165927406458e-21, 'a': 6.796594796854344, 'alpha': 0.08550170315119676},
        "('facebook/bart-base', 'gigaword')": {'L_inf': 2.8011504261522693e-13, 'a': 7.594455935790706, 'alpha': 0.13330962892884293},
        "('facebook/bart-base', 'wikiword')": {'L_inf': 0.2505157704866666, 'a': 5.861961168558081, 'alpha': 0.12018562279904178},
        "('facebook/bart-large', 'flan')": {'L_inf': 8.90097148338901e-13, 'a': 5.121314574164879, 'alpha': 0.07407830522184079},
        "('facebook/bart-large', 'gigaword')": {'L_inf': 7.348996681476381e-17, 'a': 7.683922260497716, 'alpha': 0.14310663430953338},
        "('facebook/bart-large', 'wikiword')": {'L_inf': 0.7814633219457406, 'a': 2.620750838514362, 'alpha': 0.11520371495174195},
        "('facebook/opt-1.3b', 'flan')": {'L_inf': 9.451184733271898e-25, 'a': 3.2171035320467642, 'alpha': 0.049050713966180944},
        "('facebook/opt-1.3b', 'gigaword')": {'L_inf': 4.3610823050173863e-20, 'a': 5.229994584305667, 'alpha': 0.10805297945869073},
        "('facebook/opt-1.3b', 'wikiword')": {'L_inf': 2.9031538655738537e-16, 'a': 2.3491291175103908, 'alpha': 0.041762877333812815},
        "('facebook/opt-350m', 'flan')": {'L_inf': 7.937175363375485e-15, 'a': 4.468174390456382, 'alpha': 0.057611663183111395},
        "('facebook/opt-350m', 'gigaword')": {'L_inf': 2.294599340497756e-16, 'a': 6.499202040847061, 'alpha': 0.12331848582705905},
        "('facebook/opt-350m', 'wikiword')": {'L_inf': 9.338143776217542e-10, 'a': 3.2418388198175885, 'alpha': 0.055463997801027023},
        "('facebook/opt-6.7b', 'flan')": {'L_inf': 8.14082711544948e-10, 'a': 2.2339677911390265, 'alpha': 0.019148365255527804},
        "('facebook/opt-6.7b', 'gigaword')": {'L_inf': 6.4850563005254914e-21, 'a': 2.1770799489607238, 'alpha': 0.014498526959777614},
        "('facebook/opt-6.7b', 'wikiword')": {'L_inf': 0.26936689227561367, 'a': 1.788070219787539, 'alpha': 0.04220960709323751},
        "('google/mt5-base', 'flan')": {'L_inf': 2.526590900644073e-12, 'a': 4.571149010287784, 'alpha': 0.06361718985883816},
        "('google/mt5-base', 'gigaword')": {'L_inf': 3.114373774345398e-18, 'a': 3.439603746623132, 'alpha': 0.03165062217213554},
        "('google/mt5-base', 'wikiword')": {'L_inf': 4.10877531839249e-14, 'a': 4.734453353182449, 'alpha': 0.09261262630404997},
        "('google/mt5-large', 'flan')": {'L_inf': 7.672745382307123e-19, 'a': 3.492798565293053, 'alpha': 0.05281181441669174},
        "('google/mt5-large', 'gigaword')": {'L_inf': 5.873110659782422e-19, 'a': 3.549257797168276, 'alpha': 0.03688126478621837},
        "('google/mt5-large', 'wikiword')": {'L_inf': 6.632774773255506e-16, 'a': 3.926297389250783, 'alpha': 0.07816268985536855},
        "('gpt2', 'flan')": {'L_inf': 5.700579218298801e-13, 'a': 7.1983961879890614, 'alpha': 0.08172041645085604},
        "('gpt2', 'gigaword')": {'L_inf': 5.449791195262194e-16, 'a': 6.339048122442959, 'alpha': 0.11725592134070295},
        "('gpt2', 'wikiword')": {'L_inf': 3.929261588051236e-25, 'a': 3.946267218928482, 'alpha': 0.06811657162760475},
        "('t5-base', 'flan')": {'L_inf': 4.841398974031212e-21, 'a': 3.543052166590739, 'alpha': 0.052255117939287866},
        "('t5-base', 'gigaword')": {'L_inf': 0.41674098898230477, 'a': 1.8233793996455916, 'alpha': 0.16745997442401883},
        "('t5-base', 'wikiword')": {'L_inf': 2.264923914746759e-22, 'a': 2.2598720880828873, 'alpha': 0.04457386963843646},
        "('t5-small', 'flan')": {'L_inf': 5.731082823765475e-18, 'a': 4.0524929725844165, 'alpha': 0.05269102972550974},
        "('t5-small', 'gigaword')": {'L_inf': 0.40091815004312686, 'a': 1.7757420094751855, 'alpha': 0.13439805928221435},
        "('t5-small', 'wikiword')": {'L_inf': 1.5765442089801293e-20, 'a': 2.7905737195023983, 'alpha': 0.0507995309961754},
    }
    if group not in params:
        raise ValueError(f"Unknown group: {group}")
    p = params[group]
    results = []
    for d in input_data:
        N = d.get('sft_data_size')
        # Apply scaling law: loss = L_inf + a * N^{-alpha}
        loss = p['L_inf'] + p['a'] * (N ** (-p['alpha']))
        results.append({'sft_loss': loss})
    return results