mixle.task.edge module

Edge distillation: jointly optimize the student’s structure and its training process under a hard device budget – driven by a model that designs models.

mixle.task.tune_recipe() tunes one family’s knobs with a soft cost penalty. An edge deployment is a different problem: the budget is hard (bytes of flash, ops per inference, “no torch on the device”), and the biggest wins come from choosing the right kind of student – a hashed-feature MLP versus a structured probabilistic classifier (a learned Bayesian network: kilobytes, torch-free, exactly calibrated posteriors) – not from tweaking widths inside one kind.

This module makes both choices searchable, and makes the search itself a model:

  • DeviceSpec / EdgeFootprint – the deployment budget and a student’s measured cost (serialized bytes, per-inference op count, torch dependence).

  • EdgeSpace – one design space spanning family (structure) and each family’s training recipe (process), decoded from a unit cube so any optimizer can drive it.

  • DesignModel – the model that writes the model: GP surrogates over design -> (quality, budget violations), proposing the next design by feasibility-weighted expected improvement (mixle.doe.propose_next_constrained()). It persists (to_json/from_json) and keys every row on a task_fingerprint(), so design knowledge accumulates across tasks: the surrogate learns task-similarity through the fingerprint coords, and a warm-started search – even on a different task – skips the designs that never work anywhere near it.

  • distill_for_edge() – the front door: screen candidates at reduced fidelity, promote the promising ones to full training, return the best student that fits the device, with the Pareto front over (bytes, agreement) and the updated DesignModel.

  • distill_designer() – the recursion made useful: distill the accumulated design ledger into a tiny torch-free structured student that predicts whether a design is worth training – the design model, compressed by the very machinery it steers.

Teachers are consulted once per dataset (labels are cached), never per candidate.

class DeviceSpec(max_bytes=None, max_ops=None, torch_free=False)[source]

Bases: object

A hard deployment budget. None leaves an axis unconstrained.

max_bytes: model size on flash/disk. max_ops: per-inference op budget (a latency proxy – calibrate ops/sec once per device to turn a latency target into this number). torch_free: the device cannot run torch, so only pure-mixle students qualify.

Parameters:

  • max_bytes (int | None)

  • max_ops (int | None)

  • torch_free (bool)

max_bytes: int | None = None
max_ops: int | None = None
torch_free: bool = False
classmethod for_latency(max_ms, ops_per_second, *, max_bytes=None, torch_free=False)[source]

A budget from a latency target: max_ops = ops_per_second * max_ms / 1000.

ops_per_second must come from a probe run on the target device for the student kind you deploy (measure_ops_per_second() measures it for a representative student) – throughput differs by orders of magnitude across devices and student kinds, so there is no honest built-in constant.

Parameters:
Return type:

DeviceSpec

violations(fp)[source]

Normalized constraint values, feasible when <= 0 (the form constrained BO consumes).

Parameters:

fp (EdgeFootprint)

Return type:

list[float]

feasible(fp)[source]
Parameters:

fp (EdgeFootprint)

Return type:

bool

class EdgeFootprint(bytes, ops, torch_free)[source]

Bases: object

A student’s measured deployment cost: serialized bytes, per-inference ops (multiply- accumulates for an MLP; factor evaluations for a structured classifier), and torch_free.

Parameters:
bytes: int
ops: int
torch_free: bool
class EdgeSpace(families=('mlp', 'structured'), dim_choices=(64, 128, 256, 512), hidden_range=(4, 96), epochs_range=(40, 320), log10_lr_range=(-3.0, -1.0), bits_choices=(32, 8), ngram=3, components_range=(1, 4), bins_range=(2, 8), max_its_range=(10, 60), min_gain_range=(0.0, 5.0))[source]

Bases: object

One unit-cube design space over family (structure) and each family’s recipe (process).

Coordinate 0 selects the family; 1..4 decode family-specifically. families defaults to what the input kind and device allow: hashed-MLP students need torch; structured students need fixed-schema records. Decode is deterministic, so a design point is a reproducible recipe.

Parameters:
families: tuple[str, ...] = ('mlp', 'structured')
dim_choices: Sequence[int] = (64, 128, 256, 512)
hidden_range: tuple[int, int] = (4, 96)
epochs_range: tuple[int, int] = (40, 320)
log10_lr_range: tuple[float, float] = (-3.0, -1.0)
bits_choices: Sequence[int] = (32, 8)
ngram: int = 3
components_range: tuple[int, int] = (1, 4)
bins_range: tuple[int, int] = (2, 8)
max_its_range: tuple[int, int] = (10, 60)
min_gain_range: tuple[float, float] = (0.0, 5.0)
dims()[source]
Return type:

int

bounds()[source]
Return type:

list[tuple[float, float]]

signature()[source]

Fingerprint of the space so persisted design knowledge is only reused where it applies.

Return type:

str

decode(point)[source]

Unit-cube point -> (family, recipe kwargs); the mlp recipe carries a bits precision.

Parameters:

point (ndarray)

Return type:

tuple[str, dict[str, Any]]

class DesignModel(signature, n_constraints, n_fingerprint=0)[source]

Bases: object

A probabilistic model of the design space itself: design point -> (quality, budget violations).

Every evaluated design is a row; GP surrogates fitted on the rows drive propose() (the next design worth training, by feasibility-weighted expected improvement) and predict() (mean, sd, and probability-of-fitting-the-device for untrained designs). It serializes, so what was learned designing students for one task warm-starts the next – the design knowledge is itself a model artifact. (And distill_designer() compresses it into a student – models all the way down, each level a real artifact.)

Parameters:

  • signature (str)

  • n_constraints (int)

  • n_fingerprint (int)

X: list[list[float]]
quality: list[float]
violations: list[list[float]]
tags: list[dict[str, Any]]
add(point, quality, violations, *, fingerprint=None, **tag)[source]
Parameters:
Return type:

None

propose(bounds, *, seed=None, n_candidates=256, prefilter=None, max_tries=8, fingerprint=None)[source]

The next design worth training: feasibility-weighted EI over everything seen so far.

prefilter closes the designer loop: pass a design judge – typically the tiny student from distill_designer(), called as prefilter(point_tuple) -> label – and any proposal it labels "weak" is vetoed and re-drawn (fresh acquisition seed), up to max_tries. The distilled design knowledge thus skips known-bad designs before a single training run is spent; if every retry is vetoed the last proposal is returned anyway (the judge advises, the surrogate decides). fingerprint conditions the proposal on the current task (see _fingerprint_bounds()); the returned point has design coords only.

Parameters:
Return type:

ndarray

predict(points, *, fingerprint=None)[source]

For untrained designs: predicted quality (mean, sd) and P(fits the device).

points carry design coords only; fingerprint (required when the ledger is fingerprinted) selects which task’s slice the prediction conditions on.

Parameters:
Return type:

dict[str, ndarray]

to_json()[source]
Return type:

dict[str, Any]

classmethod from_json(d)[source]
Parameters:

d (dict[str, Any])

Return type:

DesignModel

class EdgeDistillResult(model, family, recipe, agreement, footprint, feasible, pareto, design, trials=<factory>)[source]

Bases: object

Outcome of an edge search: the winning student (with measured footprint), whether it truly fits the device, the (bytes, agreement) Pareto front over everything trained, and the updated DesignModel carrying the accumulated design knowledge.

Parameters:
model: TaskModel
family: str
recipe: dict[str, Any]
agreement: float
footprint: EdgeFootprint
feasible: bool
pareto: list[dict[str, Any]]
design: DesignModel
trials: list[dict[str, Any]]
distill_for_edge(teacher, train_data, val_data, device, *, labels=None, train_labels=None, val_labels=None, space=None, design=None, designer=None, n_init=4, n_iter=6, screen_fidelity=0.3, promote=2, seed=0, task='')[source]

Search structure x training-process for the best student that fits device.

The teacher labels train_data/val_data once (cached) – or pass train_labels/ val_labels when the labels already exist (a harvested dataset, an upstream solve split) and the teacher is then never called (it may be None). Candidates proposed by the DesignModel are trained at screen_fidelity (cheap), scored by held-out agreement, and measured (footprint()); the top promote feasible screens are re-trained at full fidelity and the best feasible one wins (ties -> smaller). Pass a previous search’s design (same space + device shape) to warm-start: the surrogate already knows which regions blow the budget. Pass designer (the tiny judge from distill_designer()) to veto known-weak proposals before any training is spent. If nothing fits the device, the least-infeasible student is returned with feasible=False – inspect result.pareto for the real trade-off frontier.

Parameters:
Return type:

EdgeDistillResult

distill_designer(design, *, quality_quantile=0.5, seed=0)[source]

Distill the design ledger into a tiny torch-free student that judges designs: point -> good/weak.

The DesignModel’s rows are records (the design coordinates); a row is labeled good when it was feasible on the device and its quality reached the ledger’s quality_quantile. distill_structured_from_labels() – the same machinery the design model tunes – compresses that knowledge into a kilobyte Bayesian-network student usable as a zero-cost pre-filter for future searches. Teacher trains student; a model designs the students; the designer is distilled into a student: each level of the tower is a real artifact.

Parameters:

  • design (DesignModel)

  • quality_quantile (float)

  • seed (int)

Return type:

TaskModel

footprint(student)[source]

Measure a student’s deployment cost. Bytes are real (fp32 weights, int8 arrays for a quantized student, or serialized JSON for a structured one); ops are the closed-form per-inference count for the student’s kind.

Parameters:

student (TaskModel)

Return type:

EdgeFootprint

task_fingerprint(data, labels)[source]

A fixed small vector describing which task this is: the coords cross-task warm start keys on.

(log10 #examples, #labels, #fields, fraction of categorical fields, normalized label entropy) – cheap invariants of the dataset, O(1)-scaled so the design surrogate’s default lengthscale treats similar tasks as informative neighbors and dissimilar ones as weakly coupled.

Parameters:
Return type:

list[float]

measure_inference_seconds(student, inputs, *, repeats=5)[source]

Median measured wall-clock seconds per single-input inference for student on this host.

Runs student.batch(inputs) repeats times (after one untimed warm-up) and reports the median per-item time. Measured, not modeled – run it on the machine whose latency you care about (the deploy device, not the dev laptop) for a number that means anything there.

Parameters:
Return type:

float

measure_ops_per_second(student, inputs, *, repeats=5)[source]

Measured throughput (footprint ops / measured second) for this student kind on this host.

The calibration constant that turns a latency budget into DeviceSpec max_ops (DeviceSpec.for_latency()): probe once per (device, student kind), reuse across searches.

Parameters:
Return type:

float