MoE Scaling Law

# EVOLVE-BLOCK-START
"""
6‐parameter saturating mixing scaling law for MoE:
  loss ≈ C + (A * (d_norm^α) + B) / (E^γ + d_norm^δ)
where d_norm = dense_parameter_count / 1e8.

Parameters: [A, α, B, γ, δ, C]
"""
import numpy as np
from scipy.optimize import minimize

def scaling_law_func(data_points, params):
    X = np.atleast_2d(np.asarray(data_points))
    E = X[:, 0]                      # num_experts
    D = X[:, 1]                      # dense_parameter_count
    d = D / 1e8                      # normalized dense scale

    p = np.asarray(params)
    if p.ndim == 1:
        p = p[None, :]
    assert p.shape[1] == 6, "Expected 6 params [A, α, B, γ, δ, C]"

    A, alpha, B, gamma, delta, C = p.T
    E_mat = E[:, None]
    d_mat = d[:, None]

    # numerator: A*d^α + B
    num = A[None, :] * (d_mat**alpha[None, :]) + B[None, :]
    # denominator: E^γ + d^δ
    den = (E_mat**gamma[None, :]) + (d_mat**delta[None, :])
    pred = C[None, :] + num / den

    return pred[:, 0] if pred.shape[1] == 1 else pred

def fit_scaling_law(data_points, loss_values):
    X = np.atleast_2d(np.asarray(data_points))
    y = np.asarray(loss_values).ravel()
    E = X[:, 0]
    D = X[:, 1]
    d = D / 1e8

    def objective(raw):
        # raw = [logA, α, logB, logγ, δ, C]
        logA, alpha, logB, logGamma, delta, C = raw
        A     = np.exp(logA)
        B     = np.exp(logB)
        gamma = np.exp(logGamma)
        pred = C + (A * (d**alpha) + B) / (E**gamma + d**delta)
        return np.mean((pred - y)**2)

    # initialize near A≈1, B≈1, γ≈1, δ≈1, α≈1, C≈mean loss
    init = np.array([0.0, 1.0, 0.0, 0.0, 1.0, np.mean(y)])
    bounds = [
        (-10, 10),   # logA
        (-5, 5),     # α
        (-10, 10),   # logB
        (-5, 5),     # logγ
        (-5, 5),     # δ
        (None, None) # C
    ]
    res = minimize(objective, init, bounds=bounds, method='L-BFGS-B')
    raw_opt = res.x if res.success else init

    logA, alpha, logB, logGamma, delta, C = raw_opt
    return np.array([
        np.exp(logA),
        alpha,
        np.exp(logB),
        np.exp(logGamma),
        delta,
        C
    ])
# EVOLVE-BLOCK-END

# EVOLVE-BLOCK-START
"""
6-parameter saturating mixing law for MoE with hybrid MSE+MAE objective:
  loss ≈ C + (A·(D/1e8)^α + B) / (E^γ + (D/1e8)^δ)
Parameters: [A, α, B, γ, δ, C]
"""
import numpy as np
from scipy.optimize import minimize

def scaling_law_func(data_points, params):
    X = np.asarray(data_points, float)
    if X.ndim == 1: X = X[None, :]
    E, D = X[:,0], X[:,1]
    A, alpha, B, gamma, delta, C = params
    d = D / 1e8
    return C + (A * d**alpha + B) / (E**gamma + d**delta)

def fit_scaling_law(data_points, loss_values):
    X = np.asarray(data_points, float)
    if X.ndim == 1: X = X[None, :]
    E, D = X[:,0], X[:,1]
    d = D / 1e8
    y = np.asarray(loss_values, float).ravel()

    # Hybrid loss: MSE + 0.5 * MAE
    def objective(raw):
        logA, alpha, logB, logG, delta, C = raw
        A, B, G = np.exp(logA), np.exp(logB), np.exp(logG)
        pred = C + (A * d**alpha + B) / (E**G + d**delta)
        r = pred - y
        return np.mean(r*r) + 0.5 * np.mean(np.abs(r))

    # bounds & init
    init = [0.0, 1.0, 0.0, 1.0, 1.0, np.mean(y)]
    bounds = [(-10,10), (-5,5), (-10,10), (-5,5), (-5,5), (None,None)]
    res = minimize(objective, init, method='L-BFGS-B', bounds=bounds)

    if res.success:
        logA, alpha, logB, logG, delta, C = res.x
    else:
        logA, alpha, logB, logG, delta, C = init

    return np.array([np.exp(logA), alpha, np.exp(logB), np.exp(logG), delta, C], dtype=float)
# EVOLVE-BLOCK-END

import numpy as np
from scipy.optimize import minimize

# EVOLVE-BLOCK-START
"""
Refined MoE scaling law with a saturating expert‐term to capture diminishing returns.
Model form (6 params):
   loss ≈ A1 * D_norm^α * E_norm^(−γ)
        + A2 * (1 − exp(−β * E_norm))
        + C

where
   D_norm = dense_parameter_count / 1e8
   E_norm = num_experts / 64

Parameters = [A1, α, γ, A2, β, C]
"""
def scaling_law_func(data_points, params):
    X = np.atleast_2d(np.asarray(data_points, dtype=float))
    E = X[:, 0]                     # num_experts
    D = X[:, 1]                     # dense_parameter_count
    # fixed‐scale normalization (no input‐dependent stats)
    Dn = D / 1e8
    En = E / 64.0

    p = np.asarray(params, dtype=float)
    if p.ndim == 1:
        p = p[None, :]             # shape (1,6)
    # unpack columns
    A1, alpha, gamma, A2, beta, C = p.T

    # broadcast to compute (N × T)
    Dm = Dn[:, None]
    Em = En[:, None]

    term1 = A1[None, :] * (Dm ** alpha[None, :]) * (Em ** (-gamma[None, :]))
    term2 = A2[None, :] * (1.0 - np.exp(-beta[None, :] * Em))
    preds = term1 + term2 + C[None, :]

    # if a single parameter vector was passed, return shape (N,)
    return preds[:, 0] if preds.shape[1] == 1 else preds


def fit_scaling_law(data_points, loss_values):
    """
    Fit [A1, α, γ, A2, β, C] by minimizing mean absolute log‐error
    via multi‐start L‐BFGS‐B in a mixed raw/log space.
    """
    X = np.atleast_2d(np.asarray(data_points, dtype=float))
    y = np.asarray(loss_values, dtype=float).ravel()
    E = X[:, 0]
    D = X[:, 1]

    # same fixed normalization
    Dn = D / 1e8
    En = E / 64.0
    eps = 1e-12

    # unpack raw vector → actual params
    def unpack(raw):
        logA1, alpha, logγ, logA2, logβ, C = raw
        return (np.exp(logA1),
                alpha,
                np.exp(logγ),
                np.exp(logA2),
                np.exp(logβ),
                C)

    # objective: mean absolute log‐error
    def objective(raw):
        A1, alpha, γ, A2, β, C = unpack(raw)
        pred = (A1 * (Dn ** alpha) * (En ** (-γ))
                + A2 * (1.0 - np.exp(-β * En))
                + C)
        pred = np.maximum(pred, eps)
        return np.mean(np.abs(np.log(pred) - np.log(y + eps)))

    # bounds for raw parameters:
    #   logA1 ∈ [-5,5], α ∈ [-2,2], logγ ∈ [-5,5],
    #   logA2 ∈ [-5,5], logβ ∈ [-5,5], C ∈ [0,5]
    bounds = [
        (-5, 5),
        (-2, 2),
        (-5, 5),
        (-5, 5),
        (-5, 5),
        (0, 5),
    ]

    # prepare multi‐start initial guesses
    y0 = np.mean(y)
    inits = [
        np.array([ 0.0,  0.5,  0.0,   0.0,  1.0, y0]),
        np.array([-1.0,  1.0,  0.5,  -1.0,  0.5, y0]),
        np.array([ 1.0, -0.5,  1.0,   1.0, -0.5, y0])
    ]
    # add a couple of random perturbations (deterministic seed)
    rng = np.random.RandomState(0)
    for _ in range(2):
        raw = np.array([
            rng.uniform(-2,2),      # logA1
            rng.uniform(-1,1),      # α
            rng.uniform(-1,1),      # logγ
            rng.uniform(-2,2),      # logA2
            rng.uniform( 0,2),      # logβ
            y0
        ])
        inits.append(raw)

    best_fun = np.inf
    best_raw = inits[0]
    for raw0 in inits:
        res = minimize(
            objective,
            raw0,
            method='L-BFGS-B',
            bounds=bounds,
            options={'maxiter': 500, 'ftol': 1e-9}
        )
        if res.success and res.fun < best_fun:
            best_fun, best_raw = res.fun, res.x

    # recover optimal parameters
    A1_opt, alpha_opt, γ_opt, A2_opt, β_opt, C_opt = unpack(best_raw)
    return np.array([A1_opt, alpha_opt, γ_opt, A2_opt, β_opt, C_opt])
# EVOLVE-BLOCK-END

# EVOLVE-BLOCK-START
"""
6-parameter MoE scaling law with combined expert-log damping:
  loss ≈ A·(D/1e8)^α / [(E·log₁₊E)^γ]  +  B / (log₁₊E)^β  +  C

– First term captures joint scaling of dense params and experts,
  damping experts by both E and log(1+E).
– Second term captures long-tail expert-only diminishing returns.
– Fits in three stages: 
    1) Exponents α,γ,β via L-BFGS-B with closed-form A,B,C.
    2) Closed-form solve for A,B,C.
    3) Joint L-BFGS-B refine all six in log-space for stability.
"""
import numpy as np
from scipy.optimize import minimize

def scaling_law_func(data_points, params):
    """
    data_points: (N,2) array [num_experts E, dense_param_count D]
    params: [A, α, γ, B, β, C]
    returns: (N,) predicted validation losses
    """
    X = np.atleast_2d(np.asarray(data_points, float))
    E = X[:, 0]
    D_norm = X[:, 1] / 1e8
    p = np.asarray(params, float)
    if p.ndim == 1:
        p = p[None, :]
    assert p.shape[1] == 6, "Expected 6 params [A,α,γ,B,β,C]"
    A, alpha, gamma, B, beta, C = p.T
    L = np.log1p(E)
    pred = A * (D_norm ** alpha) / ((E * L) ** gamma) + B / (L ** beta) + C
    return pred.ravel()

def fit_scaling_law(data_points, loss_values):
    """
    Fits the 6-parameter model by:
      Stage 1: Optimize exponents [α,γ,β] via L-BFGS-B, solving [A,B,C] by LS.
      Stage 2: Solve [A,B,C] in closed form with chosen exponents.
      Stage 3: Jointly refine [logA,α,logγ,logB,β,C] via L-BFGS-B.
    Returns optimized [A, α, γ, B, β, C].
    """
    X = np.atleast_2d(np.asarray(data_points, float))
    E = X[:, 0]
    D_norm = X[:, 1] / 1e8
    y = np.asarray(loss_values, float).ravel()
    L = np.log1p(E)

    # Stage 1: optimize exponents α,γ,β
    def obj_exps(raw):
        alpha, gamma, beta = raw
        x1 = D_norm**alpha / ((E * L)**gamma)
        x2 = 1.0 / (L**beta)
        M = np.vstack([x1, x2, np.ones_like(x1)]).T
        coeffs, *_ = np.linalg.lstsq(M, y, rcond=None)
        pred = M.dot(coeffs)
        return np.mean((pred - y) ** 2)

    init_exps = np.array([0.5, 0.5, 1.0])
    bounds_exps = [(-2, 2), (-5, 5), (0, 5)]
    res_exps = minimize(obj_exps, init_exps, bounds=bounds_exps, method='L-BFGS-B')
    alpha0, gamma0, beta0 = res_exps.x if res_exps.success else init_exps

    # Stage 2: closed-form solve for A, B, C
    x1 = D_norm**alpha0 / ((E * L)**gamma0)
    x2 = 1.0 / (L**beta0)
    M = np.vstack([x1, x2, np.ones_like(x1)]).T
    coeff0, *_ = np.linalg.lstsq(M, y, rcond=None)
    A0, B0, C0 = coeff0

    # Stage 3: joint refine all six parameters
    def obj_full(raw):
        logA, alpha, logG, logB, beta, C = raw
        A = np.exp(logA)
        gamma = np.exp(logG)
        B = np.exp(logB)
        pred = A * (D_norm ** alpha) / ((E * L) ** gamma) + B / (L ** beta) + C
        return np.mean((pred - y) ** 2)

    init_raw = np.array([
        np.log(max(A0, 1e-8)),
        alpha0,
        np.log(max(gamma0, 1e-8)),
        np.log(max(B0, 1e-8)),
        beta0,
        C0
    ])
    bounds_full = [
        (-10, 10),  # logA
        (-2, 2),    # α
        (-10, 10),  # logγ
        (-10, 10),  # logB
        (0, 5),     # β
        (None, None)# C
    ]
    res_full = minimize(obj_full, init_raw, bounds=bounds_full, method='L-BFGS-B')
    raw_opt = res_full.x if res_full.success else init_raw

    logA, alpha, logG, logB, beta, C = raw_opt
    A_opt = np.exp(logA)
    gamma_opt = np.exp(logG)
    B_opt = np.exp(logB)

    return np.array([A_opt, alpha, gamma_opt, B_opt, beta, C])
# EVOLVE-BLOCK-END