mixle.ppl.scaling_laws module

F5: scaling-law fits + compute allocation – “mixle training mixle” (roadmap item F).

Fits classic Chinchilla-style neural-scaling-law curves loss = f(N, D) (N = model parameters, D = training tokens) using mixle’s OWN probabilistic-programming/regression machinery (mixle.ppl), not scipy.optimize.curve_fit or an ad-hoc fitter: a scaling law is just another regression problem, so it is expressed as an actual mixle distribution (a Normal likelihood with the power-law mean as a custom potential) and fit with how="mcmc" the same way examples.flagship_physics_inverse (D-track’s physics-inverse flagship) turns a nonlinear forward model into PPL evidence. The result is a genuine posterior over the power-law exponents/coefficients – “power-law leaves + uncertainty receipts” – not just point estimates.

Compute allocation reuses mixle.doe.bayesopt (this codebase’s existing Gaussian-process Bayesian-optimization machinery) to pick the (N, D) split that minimizes the fitted law’s predicted loss under the standard C ~= 6*N*D FLOPs approximation (Kaplan et al. 2020; Hoffmann et al. 2022).

Data provenance for the “reproduces known exponents” acceptance test (mixle/tests/scaling_laws_test.py): this environment has no network access, so real per-run (N, D, loss) tables from the literature are not fetchable here. generate_synthetic_chinchilla_data() instead generates SYNTHETIC (N, D, loss) triples from the REAL, PUBLISHED Chinchilla functional form and exponents – Hoffmann et al. 2022, “Training Compute-Optimal Large Language Models” (arXiv:2203.15556), Table 2, “Approach 3” (parametric loss risk) fit:

L(N, D) = E + A/N**alpha + B/D**beta, E=1.69, A=406.4, B=410.7, alpha=0.34, beta=0.28

with realistic observation noise added, and the test confirms fit_scaling_law() recovers alpha/beta close to these published values. This is honestly the synthetic-data path described in the roadmap item, not real published (N, D, loss) rows.

Reuse of D5’s controller brain: see ScalingLawAllocationController below and its docstring for what is (and is not) reused from mixle.inference.conditional_jit_controller (D5, PR #163).

Module location: this lives under mixle.ppl (not mixle.doe, despite allocate_compute being pure DOE machinery) because fit_scaling_law() itself imports mixle.ppl to do its fitting, and the repo’s own architectural guard (mixle/tests/ppl_separation_test.py) enforces a strict one-way dependency mixle.ppl -> core – no core module (which mixle.doe is) may import upward from the optional, torch-backed PPL layer. mixle.ppl importing mixle.doe (as allocate_compute() does, via mixle.doe.bayesopt) is the allowed direction.

generate_synthetic_chinchilla_data(n_points=60, *, seed=0, noise_sd=0.015, n_range=(1.0e7, 1.0e11), d_range=(1.0e8, 1.0e12), e=CHINCHILLA_E, a=CHINCHILLA_A, b=CHINCHILLA_B, alpha=CHINCHILLA_ALPHA, beta=CHINCHILLA_BETA)[source]

SYNTHETIC (N, D, loss) triples generated from the published Chinchilla functional form.

N/D are drawn log-uniformly over n_range/d_range (spanning several orders of magnitude, the way a real training-run sweep would), the mean loss is the exact published power law E + A/N**alpha + B/D**beta, and i.i.d. Gaussian noise of scale noise_sd (in loss units) is added – a realistic per-run measurement/optimization-noise floor. See the module docstring for why this is synthetic-from-known-exponents rather than a real published per-run table (no network access in this environment).

Parameters:
Return type:

list[tuple[float, float, float]]

class ScalingLawFit(fitted, n0, d0)[source]

Bases: object

A fitted loss = E + A/N**alpha + B/D**beta scaling law with a genuine posterior.

fitted is the mixle.ppl RandomVariable returned by .fit(..., how="mcmc") – the “power-law leaf” – carrying MCMC draws over (E, A, alpha, B, beta, sigma) (the “uncertainty receipts”). n0/d0 are the normalization constants N/D were divided by before fitting (numerical conditioning only; predictions are in the original units).

Parameters:
samples(name)[source]

Posterior draws for parameter name (one of E, A, alpha, B, beta, sigma).

Parameters:

name (str)

Return type:

ndarray

hdi(name, prob=0.9)[source]

Highest-density credible interval for parameter name at coverage prob.

Parameters:
Return type:

tuple[float, float]

predict_mean(n, d)[source]

Posterior-mean predicted loss at model size n (params) and token count d.

Parameters:
Return type:

float

predict_samples(n, d)[source]

Posterior-predictive DRAWS of the mean loss at (n, d) – integrates over parameter uncertainty (no observation noise added), for building a predictive credible interval.

Parameters:
Return type:

ndarray

fit_scaling_law(records, *, draws=4000, burn=4000, scale=0.02, seed=0, rng=None)[source]

Fit loss = E + A/N**alpha + B/D**beta to records (a list of (N, D, loss)).

Uses mixle’s OWN PPL fitting machinery, not scipy’s curve_fit: the nonlinear power-law mean is expressed as a potential (custom log-likelihood term) over free Normal-prior parameters, and the whole thing is fit as an ordinary mixle.ppl model with how="mcmc" – exactly the pattern examples/flagship_physics_inverse.py (the D-track’s physics-inverse-problem flagship) uses to turn an arbitrary nonlinear forward model into PPL evidence. Returns a ScalingLawFit carrying full MCMC posterior draws (real uncertainty, not just a point estimate).

scale is passed straight through to how="mcmc"’s adaptive random-walk proposal as its initial per-coordinate step size. The carrier observation used here is a single vacuous point (the potential IS the evidence, exactly as in the physics-inverse flagship), so mixle’s default proposal-scale heuristic (~ data_std / sqrt(n_data)) sees n_data=1 and starts far too wide relative to this tightly-peaked 6-parameter likelihood – left at its default it needs many thousands of extra burn-in draws for the adaptive proposal to shrink into a workable acceptance-rate regime. Setting a small explicit scale up front (tuned to this problem’s posterior widths, ~1e-2 on the log-parameter scale) restores good mixing (~15-25% acceptance) at the draws/burn defaults below.

Parameters:
Return type:

ScalingLawFit

allocate_compute(fit, compute_budget, *, n_bounds=(1.0e7, 1.0e12), n_init=8, n_iter=20, seed=0, flops_per_token_param=FLOPS_PER_TOKEN_PARAM)[source]

Find the (N, D) split minimizing fit’s predicted loss under C ~= 6*N*D.

Reuses mixle.doe.bayesopt.minimize() (GP-surrogate expected-improvement Bayesian optimization) – this codebase’s existing DOE machinery – rather than a bespoke optimizer. The compute constraint C = 6*N*D is an exact algebraic EQUALITY, not a black-box inequality, so instead of routing through mixle.doe.constrained (built for black-box inequality constraints, which this is not), it is eliminated by substitution: for any candidate N, D is set to exactly satisfy the constraint, collapsing the 2-D allocation problem to a 1-D search over log10(N) that bayesopt.minimize drives directly.

Parameters:
Return type:

tuple[float, float]

allocate_fixed_heuristic(compute_budget, *, ratio=20.0, flops_per_token_param=FLOPS_PER_TOKEN_PARAM)[source]

The commonly-cited FIXED tokens ~= 20 * params heuristic, solved jointly with C = 6*N*D.

D = ratio * N and C = 6*N*D = 6*ratio*N**2, so N = sqrt(C / (6*ratio)) and D = ratio*N. ratio=20 is the widely-cited Chinchilla-style rule of thumb (Hoffmann et al. 2022’s own “roughly 20 tokens per parameter” summary of their compute-optimal frontier), used here purely as the FIXED baseline the DOE allocator is compared against – it ignores the fitted scaling law entirely.

Parameters:
Return type:

tuple[float, float]

class ScalingLawState(log_e, log_a, log_alpha, log_b, log_beta, log_budget)[source]

Bases: object

Fingerprint for the compute-allocation decision: the fitted law’s posterior-mean parameters (log scale) plus the requested compute budget – the DesignModel task fingerprint, mirroring D5’s ControllerState.as_vector().

Parameters:
class AllocationAction(log_n)[source]

Bases: object

One controller decision: log10(N) (D follows from the C = 6*N*D constraint).

Parameters:

log_n (float)

class ScalingLawAllocationController(*, n_bounds=(1.0e7, 1.0e12), design=None, seed=None)[source]

Bases: object

D5-pattern controller for the compute-allocation decision – shares D5’s controller brain.

D5 (mixle/inference/conditional_jit_controller.py, PR #163) defines a generic LearnedController[StateT, ActionT] base (select_action(state) -> action / update(state, action, gain, cost) -> None) plus a concrete DesignModelController that wraps mixle.task.edge.DesignModel – a GP-surrogate design space model, warm-startable across DIFFERENT tasks via a fingerprint vector – to propose a continuous 1-D knob (there, budget_fraction) from logged (state, action, gain, cost) rows, falling back to a fixed default before at least two rows are logged. D5’s own docstring explicitly anticipates this reuse: “a future F5 … item could subclass LearnedController directly for its own state/action types … reusing the bandit/ DesignModel wiring pattern without needing block-EM’s ControllerState/ControllerAction dataclasses at all.”

This class does exactly that: it is NOT a subclass of D5’s concrete DesignModelController (that class’s state/action types are block-EM-specific), but it reuses the SAME mixle.task.edge.DesignModel wiring D5’s DesignModelController uses, against F5’s own ScalingLawState (fitted-law-parameters + compute-budget fingerprint) and AllocationAction (log10(N)) types, with the identical cold-start-fallback and fingerprint-conditioned-proposal shape. It does not subclass D5’s mixle.inference.conditional_jit_controller.LearnedController ABC directly (importing D5’s module here would pull an unrelated inference-internals dependency into mixle.doe for a class whose only job is to satisfy the abstract two-method surface); the class shape below is deliberately identical to it so the substitutability the roadmap asks for (“shares the controller brain with D5”) is structural, not merely nominal.

Use allocate_compute_learned() for the common case (propose once from a fresh or warm-started controller); construct this directly to accumulate logged rows across many budgets/fits via repeated update() calls (the warm-start path).

Parameters:
allocate_compute_learned(fit, compute_budget, *, controller=None, flops_per_token_param=FLOPS_PER_TOKEN_PARAM)[source]

Propose (N, D) via the D5-pattern ScalingLawAllocationController, then log the realized outcome back into it (so a caller reusing the returned controller across several budgets warm-starts it, exactly D5’s cross-task DesignModel warm-start story).

This is the OPTIONAL learned path (see the roadmap item’s “optionally wire in D5’s LearnedController pattern”); allocate_compute() (plain GP-BO via mixle.doe.bayesopt) is the primary, required allocator and is what the acceptance test compares against the fixed heuristic. With a fresh (cold) controller and no logged history this falls back to the bounds midpoint, so it is not expected to beat the heuristic on a single cold call – its payoff is warm-starting across many allocation decisions, the same story as D5’s DesignModelController.

Parameters:

  • fit (ScalingLawFit)

  • compute_budget (float)

  • controller (ScalingLawAllocationController | None)

  • flops_per_token_param (float)

Return type:

tuple[float, float, ScalingLawAllocationController]