mixle.task.tune module

Search a student recipe with mixle.doe – find a small 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 cheapest recipe that still matches – the “minimize train time” pillar made concrete. 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.

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:
dim_choices: Sequence[int] = (128, 256, 512, 1024)
hidden_range: tuple[int, int] = (16, 128)
epochs_range: tuple[int, int] = (50, 400)
log10_lr_range: tuple[float, float] = (-3.0, -1.0)
n: int = 4
dims()[source]
Return type:

int

decode(point)[source]
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 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:
model: TaskModel
recipe: dict[str, Any]
agreement: float
score: float
cost: float
history: Any = None
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 cheapest 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