mixle.models.compress module

One compress() front door unifying sampling KD, non-sampling (data-free), and hybrid compression (roadmap J1).

Three method FAMILIES, one entry point

1. ``”sampling_kd”`` – this codebase’s EXISTING response/hint/attention/relational KD machinery (mixle.task.distill_methods, “already in-tree and load-bearing” per the roadmap context anchors). compress() wraps response_distill() directly: build a fresh, smaller student architecture and train it against the real teacher on REAL calibration data (real forward/backward passes) – this is the expensive-but-accurate full method every other method here is measured against.

2. ``”non_sampling”`` – the data-free compression stack G1-G3 already built: G1’s moment- propagation surrogate (mixle.models.moment_propagation) and G3’s coarsening operator (mixle.models.coarsening.coarsen(), which itself calls G1’s laws and is the direct consumer of G2’s Sigma-weighted projections per that module’s own docstring). No real forward pass over calibration data is used anywhere in this path – only propagated Gaussian LAWS.

3. ``”hybrid”`` – the genuinely novel piece: run non_sampling first, read off G1’s own per-stage closure-error receipts that coarsen() already produces (ScaleReceipt.surrogate_closure_error – how far the Gaussian-surrogate assumption itself is from a real Monte Carlo check at that stage, i.e. “where is the surrogate blind”), rank stages by that receipt via A5’s acquire() (adapted here: the “pool” is compression STAGES rather than A5’s original unlabeled data pool, and the score is the closure-error receipt rather than an EIG/disagreement score over a classifier’s predictions – see _closure_error_strategy()), and spend a REAL but TINY sampling-KD budget (_finetune_stages(), built directly on kd_loss()) ONLY on the worst-closure-error stages’ own parameters – everywhere else keeps the free non_sampling result untouched.

``method=”auto”``: mirrors roadmap I1’s (mixle.models.unified_quantizer) exact pattern one level up the compression stack – run every method once, measure its REAL quality against the teacher, and let mixle.task.bandit.UCB1 (I1’s own picker, not a new one) sweep the arms (here the arms are the three METHOD FAMILIES, “which compression method for this layer/stage”-scale, rather than I1’s four per-tensor quantization primitives) and pick the highest real reward. Every choice – auto or explicit – carries a CompressionReceipt, I1’s QuantizationReceipt translated to this module’s vocabulary: the real measured quality and sample cost of the method that WAS picked, plus (for auto) the same real numbers for every method that was considered and rejected.

Build vs. borrow

Nothing here reimplements G1/G2/G3, A5, or the sampling-KD machinery – compress() is composition: coarsen() for non_sampling, response_distill() / kd_loss() for sampling_kd/the hybrid fine-tune step, acquire() for hybrid’s stage ranking, UCB1 for auto.

class MethodCandidate(method, quality, sample_count, reward)

Bases: object

Real, measured numbers for one method family considered for one compress() call – the raw material every CompressionReceipt (chosen or rejected) is built from, mirroring mixle.models.unified_quantizer.MethodCandidate.

Parameters:

• method (str)

• quality (float)

• sample_count (int)

• reward (float)

class CompressionReceipt(method, auto, quality, sample_count, candidates=<factory>, notes='')

Bases: object

Explains why method was used: its own measured numbers, plus – for auto-pick – the same real numbers for every OTHER method considered and rejected. Mirrors mixle.models.unified_quantizer.QuantizationReceipt.

Parameters:

• method (str)

• auto (bool)

• quality (float)

• sample_count (int)

• candidates (dict[str, MethodCandidate])

• notes (str)

class CompressedModel(model, method, receipt, non_sampling_receipts=<factory>, hybrid_selected_stages=<factory>)

Bases: object

Unified result: the compressed module plus the receipt explaining which method produced it, and (when the chosen/underlying path touched non_sampling) the raw per-stage ScaleReceipt map that path’s closure-error signal came from.

Parameters:

• model (Any)

• method (str)

• receipt (CompressionReceipt)

• non_sampling_receipts (dict[str, ScaleReceipt])

• hybrid_selected_stages (list[str])

compress(model, method='auto', *, calibration_data=None, eval_data=None, sample_budget=None, input_law=None, budget=5.0, trust_region=5.0, n_mc=64, seed=0, kd_epochs=200, kd_lr=1e-2, hybrid_sample_fraction=0.01, hybrid_max_stages=1, hybrid_epochs=40, hybrid_lr=5e-4, target_n_layer=None)

The single compression entry point (roadmap J1): dispatches to one of the three method families, or lets a UCB1 CompressionReceipt-carrying picker choose (method="auto", default).

Parameters:

• model (Any) – a real mixle.models.transformer.CausalLM (or a coarsened one).

• method (str) – "auto" (default), "sampling_kd", "non_sampling", or "hybrid".

• calibration_data (Any) – integer token-context tensor/array (N, L) – required for "sampling_kd", "hybrid", and "auto" (the two methods that touch real data).

• eval_data (Any) – held-out integer token-context tensor/array used to MEASURE quality (teacher-agreement against model); defaults to calibration_data if omitted.

• sample_budget (int | None) – an explicit cap (absolute count) on how many REAL calibration samples "hybrid"/"auto"’s hybrid arm may use; if None, hybrid_sample_fraction of len(calibration_data) is used instead.

• input_law (MultivariateGaussianDistribution | None) – the data-free input law for G1’s propagation; defaults to _default_input_law() (the model’s own token-embedding statistics).

• budget (float) – forwarded to mixle.models.coarsening.coarsen().

• trust_region (float) – forwarded to mixle.models.coarsening.coarsen().

• n_mc (int) – forwarded to mixle.models.coarsening.coarsen().

• kd_epochs (int) – forwarded to the fresh, from-scratch sampling_kd student’s full training loop.

• kd_lr (float) – forwarded to the fresh, from-scratch sampling_kd student’s full training loop.

• hybrid_sample_fraction (float) – fraction of calibration_data hybrid may use when sample_budget is not given (default 1%, matching the acceptance criterion).

• hybrid_max_stages (int) – how many worst-closure-error stages hybrid fine-tunes (default 1).

• hybrid_epochs (int) – forwarded to hybrid’s own micro-calibration fine-tune loop – deliberately separate from kd_epochs/kd_lr: hybrid fine-tunes ALREADY-good non_sampling weights on a HANDFUL of real samples (a light nudge), a completely different regime from training a fresh architecture from scratch on the full calibration set, and reusing the same hyperparameters for both badly overfits hybrid’s tiny sample budget in practice.

• hybrid_lr (float) – forwarded to hybrid’s own micro-calibration fine-tune loop – deliberately separate from kd_epochs/kd_lr: hybrid fine-tunes ALREADY-good non_sampling weights on a HANDFUL of real samples (a light nudge), a completely different regime from training a fresh architecture from scratch on the full calibration set, and reusing the same hyperparameters for both badly overfits hybrid’s tiny sample budget in practice.

• target_n_layer (int | None) – depth of the fresh sampling_kd student; defaults to coarsen’s own natural 2x depth cut (max(1, model.n_layer // 2)) so all three methods are compared at a MATCHED compression ratio.

• seed (int)

Returns:

CompressedModel

Return type:

CompressedModel