mixle.task.distill_soft module

Soft-label distillation from a teacher probability distribution.

mixle.task.distill distills hard teacher labels into a local student. This module uses the richer case where the teacher exposes a probability or top-k log-probability vector for each example. Matching that distribution preserves runner-up class information and confidence structure that hard labels discard.

This is the frontier-label analogue of temperature-softened Hinton distillation in mixle.task.distill_methods, without requiring a torch teacher that exposes logits. The teacher is any callable returning a per-example probability vector. The student is the compact hashed-n-gram MLP used by mixle.task.distill, trained against soft targets with temperature-scaled KL and optionally mixed with hard-label loss. The result is a TaskModel whose proba_batch approximates the teacher’s calibrated distribution and can be calibrated or routed like any other student.

distill_from_soft_labels(texts, teacher_probs, *, labels, temperature=2.0, hard_weight=0.0, n=3, dim=256, hidden=(64,), epochs=300, lr=1e-2, seed=0, task='', device='cpu')[source]

Fit a student to per-example teacher probabilities over labels.

teacher_probs is (N, C) with rows summing to 1 (renormalized if not), column j the teacher’s probability of labels[j]. The student minimizes the temperature-softened T^2 * KL(teacher || student) (Hinton’s scaling, so the soft gradients keep magnitude as T grows), optionally mixed with hard_weight times the hard cross-entropy against the teacher’s argmax. temperature > 1 softens both sides so runner-up structure influences the fit. The result is deterministic given seed and returns a TaskModel whose proba_batch approximates the teacher’s full distribution.

Parameters:
Return type:

TaskModel

distill_soft(teacher_proba, texts, *, labels, **kwargs)[source]

Query a probability-returning teacher once over texts and soft-distill it (see distill_from_soft_labels()). teacher_proba(texts) -> (N, C) returns each example’s class distribution over labels (e.g. an LLM’s normalized top-k logprobs).

Parameters:
Return type:

TaskModel

soft_agreement(student, teacher_probs, texts)[source]

Mean KL divergence KL(teacher || student) over texts – how faithfully the student matches the teacher’s full soft distribution (0 = identical), the soft-distillation analog of mixle.task.distill.agreement(). Lower is better; use it to compare soft vs hard students.

Parameters:
Return type:

float