mixle.task.quantize module

Post-training quantization of distilled MLP students: int8/int4 weights, numpy-only inference.

An fp32 MLP student costs 4 bytes x params and needs torch at inference. Quantizing to int8 (per-tensor symmetric: W ~ round(W / s) * s with one fp32 scale per layer) cuts the weight bytes 4x; int4 packs two weights per byte for 8x – and because the dequantized forward pass is three numpy matmuls, the quantized student needs no torch at all on the device: it joins the structured students in the torch-free deployable class, while keeping the MLP’s shape. Accuracy is whatever it measures after quantization – the edge search (mixle.task.edge.distill_for_edge()) scores the quantized model’s real agreement, so the bits axis trades measured bytes against measured fidelity, never assumed ones.

quantize_mlp(student, bits=8|4) converts a trained torch student in place of retraining; the result stores quantized arrays (payload="arrays", int4 stored nibble-packed), round-trips through the artifact as an .npz, and reports its true byte size.

LNS (log-number-system; mixle.engines.lns) is wired where it is a complete fit: the structured student, whose inference is sums of factor log-densities – lns_classifier() re-executes it on integers (add / max / LUT, no transcendentals above the leaf boundary). Signed MLP matmuls in LNS would need signed-logadd kernels and stay out of scope.

class QuantizedMLP(layers, *, bits=8)

Bases: object

A quantized-weight MLP with a pure-numpy forward pass.

layers is [(W_int (out, in), scale fp32, bias fp32 (out,)), ...] with weights in the symmetric bits range (int8: [-127, 127]; int4: [-7, 7], stored nibble-packed on disk); the forward is x @ (W * s).T + b with ReLU between layers – exactly the dequantized version of the trained torch stack, so its logits match torch-on-dequantized-weights to float tolerance.

Parameters:

• layers (list[tuple[np.ndarray, float, np.ndarray]])

• bits (int)

logits(feats)

Parameters:

feats (ndarray)

Return type:

ndarray

nbytes()

Deployable payload bytes: packed weights (1 B/weight at int8, 1/2 B at int4) + fp32 biases + one fp32 scale per layer.

Return type:

int

macs()

Per-inference multiply-accumulates (integer x fp32 dequant multiplies count the same).

Return type:

int

to_arrays()

Return type:

dict[str, ndarray]

classmethod from_arrays(arrays)

Parameters:

arrays (dict[str, ndarray])

Return type:

QuantizedMLP

class QuantizedClassifierIO(featurizer, labels)

Bases: _ClassifierIO

The classifier IO for quantized students: same featurize -> logits -> label contract, no torch.

Parameters:

• featurizer (Any)

• labels (list[str])

kind = 'quantized_classifier'

logits_batch(model, raw_inputs)

Parameters:

• model (Any)

• raw_inputs (list[Any])

Return type:

ndarray

to_spec()

Return type:

dict[str, Any]

classmethod from_spec(spec)

Parameters:

spec (dict[str, Any])

Return type:

QuantizedClassifierIO

quantize_mlp(student, *, bits=8, clip_percentile=None)

Quantize a trained torch MLP student to an int8/int4, numpy-inference TaskModel.

Per-tensor symmetric weight quantization (scale = max|W| / qmax with qmax 127 for int8, 7 for int4); biases stay fp32 (they are a negligible byte fraction and quantizing them buys nothing). The returned student reuses the same featurizer and label list, reports payload="arrays" (int4 weights nibble-packed on disk: two per byte), and – having no torch dependence at inference – qualifies for torch_free devices. LNS needs LUT matmul kernels (mixle.engines.lns) and is left explicitly unimplemented.

clip_percentile guards heavy-tailed weights. Plain max-scaling lets one outlier set the whole layer’s scale: at int4 (qmax=7) a single weight 30x the rest quantizes everything else to 0, collapsing the layer. When set (e.g. 99.9), the scale is derived from that percentile of |W| instead of the max, and weights above it saturate at +/-qmax – the bulk of the distribution keeps its resolution at the cost of clipping a few outliers. Default None keeps the exact max-scale behavior (bit-identical on well-behaved weights).

Parameters:

• student (TaskModel)

• bits (int)

• clip_percentile (float | None)

Return type:

TaskModel

class LNSStructuredClassifierIO(field_keys, label_index, labels, step=1e-2)

Bases: StructuredClassifierIO

The structured classifier executed in the log-number system: integers above the leaf boundary.

A structured student’s per-label score is a sum of factor log-densities – in log-space that is products of probabilities, which is exactly what LogNumberSystem runs on integers: each factor’s log-density is quantized once at the leaf boundary (k = round(logp / step)), then the per-label accumulation is integer ADDs, mixture components fold with the integer logadd LUT, the classification is an integer argmax, and the posterior is the integer log-softmax of mixle.engines.lns_nn – no exp/log anywhere above the leaves (one exp only if you ask for linear-scale probabilities). The dequantized scores match the float classifier within the engine’s documented bound (~``1.5 * step`` per fold), so step is a dial between integer-width and fidelity.

Categorical factors are pre-quantized to integer tables at first use, so their leaves are pure integer lookups – on an all-discrete schema inference touches no floats at all. Continuous leaves evaluate in float and quantize at the boundary, the same contract as the engine’s SumProductCircuit.

Parameters:

• field_keys (list[str] | None)

• label_index (int)

• labels (list[str])

• step (float)

kind = 'lns_structured_classifier'

int_logits_batch(model, raw_inputs)

Per-label INTEGER log-joint scores (m, K) – the whole combination is integer math.

Parameters:

• model (Any)

• raw_inputs (list[Any])

Return type:

ndarray

logits_batch(model, raw_inputs)

Per-label log-joint log P(features, label) as an (m, K) score matrix (the classifier logits).

Parameters:

• model (Any)

• raw_inputs (list[Any])

Return type:

ndarray

proba_batch(model, raw_inputs)

Posterior via the INTEGER log-softmax (max + LUT); one exp at the very end for linear scale.

The LUT rounds each log-probability to ~``step``, so the raw exp sums to 1 +/- K*step/2; the final float renormalization (free – we already left integer space for the exp) removes that systematic drift without touching the integer pipeline.

Parameters:

• model (Any)

• raw_inputs (list[Any])

Return type:

ndarray

predict_batch(model, raw_inputs)

Parameters:

• model (Any)

• raw_inputs (list[Any])

Return type:

list[str]

to_spec()

Return type:

dict[str, Any]

classmethod from_spec(spec)

Parameters:

spec (dict[str, Any])

Return type:

LNSStructuredClassifierIO

lns_classifier(student, *, step=1e-2)

Re-execute a structured student in integer log-space (the LNS rung for structured students).

The fitted model is unchanged (same factors, same JSON artifact); what changes is how inference runs: factor log-densities are quantized once at the leaf boundary, and everything above – per-label accumulation, mixture folding, the argmax decision, the posterior’s log-softmax – is integer add/max/LUT arithmetic (LNSStructuredClassifierIO). step trades fidelity for integer width; the dequantized scores match the float classifier within ~``1.5 * step`` per fold. This is compute quantization (transcendental-free combination), not weight compression – pair it with the structured student’s already-tiny JSON payload.

Parameters:

• student (TaskModel)

• step (float)

Return type:

TaskModel