mixle.task.tune module

Search a student recipe with mixle.doe – find a compact model that matches the teacher for the least compute.

Distillation has knobs (feature width, hidden size, epochs, learning rate) that trade fidelity against training cost. Rather than grid-search them, tune_recipe() runs GP Bayesian optimization (mixle.doe.minimize) over the recipe space, distilling and scoring a handful of candidates and homing in on the best. The objective is held-out agreement with the teacher, optionally minus a compute penalty (cost_weight) so the search prefers the lowest-cost recipe that still matches. Returns the re-distilled winner as a callable TaskModel plus the full search history.

The recipe space is a few interpretable axes with sensible defaults; override space to widen or pin them.

tune_recipe_for_routing is the routing-ready sibling: it runs the same search, then calibrates the winning recipe into a CalibratedTaskModel on data the search never touched – so a task gets an automatically right-sized model (search picks the complexity) that is also immediately decide()-able for Cascade / Router, with no separate calibration step.

class RecipeSpace(dim_choices=(128, 256, 512, 1024), hidden_range=(16, 128), epochs_range=(50, 400), log10_lr_range=(-3.0, -1.0), n=4)[source]

Bases: object

The tunable axes of a distillation recipe and how a unit-cube point decodes into concrete knobs.

Parameters:
dims()[source]

Return the normalized recipe-search dimensionality.

Return type:

int

decode(point)[source]

Decode a normalized design point into a distillation recipe.

Parameters:

point (ndarray)

Return type:

dict[str, Any]

cost(recipe)[source]

Relative training cost of a recipe in [0, 1] (params x steps, normalized by the space’s max).

Parameters:

recipe (dict[str, Any])

Return type:

float

bounds()[source]

Return normalized DOE bounds for recipe search.

Return type:

list[tuple[float, float]]

class TuneResult(model, recipe, agreement, score, cost, history=None)[source]

Bases: object

The outcome of a recipe search: the winning model, its recipe and scores, and the full BO history.

Parameters:
tune_recipe(teacher, train_texts, val_texts, *, labels=None, space=None, n_init=4, n_iter=8, cost_weight=0.0, seed=0, task='')[source]

Bayesian-optimize the distillation recipe; return the best re-distilled TaskModel.

Maximizes held-out agreement(student, teacher, val_texts) minus cost_weight * relative_train_cost. Set cost_weight > 0 to prefer the lowest-cost recipe that still matches the teacher. teacher is called once per candidate on val_texts (cached across the search) and once per candidate on train_texts.

Parameters:
Return type:

TuneResult

class CalibratedTuneResult(model, recipe, agreement, score, cost, history=None)[source]

Bases: object

The outcome of a routing-ready recipe search: the calibrated winner, its recipe and scores, and history.

Parameters:
tune_recipe_for_routing(teacher, train_texts, val_texts, *, labels=None, space=None, n_init=4, n_iter=8, cost_weight=0.0, calibration_frac=0.3, alpha=0.1, seed=0, task='', density_gate=False, density_gate_alpha=0.05)[source]

Optimize a distillation recipe and calibrate the winning model for routing.

The search holds back a calibration_frac slice of val_texts before evaluating candidate recipes. That slice does not score candidates or influence the search; it is used afterward to calibrate the winning model into a CalibratedTaskModel. The result is a task-specific recipe whose complexity and epoch budget were selected from data and whose model can be passed directly to a Cascade or Router.

Teacher calls are shared through one cache. train_texts are queried once for the whole search rather than once per trial, and validation inputs that appear in both calibration and search slices are not queried twice. Every distinct input is priced once, no matter how many candidate recipes the search evaluates.

density_gate=True wires the same OOD escalation as distill_for_routing(): a gate fit on train_texts, its floor calibrated on the disjoint cal_texts slice.

Parameters:
Return type:

CalibratedTuneResult