mixle.task.distill_methods module

Classic knowledge-distillation families as composable, verified procedures on Torch modules.

mixle.task.distill distills a teacher’s labels into a local student (response/label distillation, no gradients through the teacher). This module is its representation-matching sibling: given a trained torch teacher and an untrained torch student, transfer knowledge through the classic KD signals –

  • response_distill() – Hinton dark-knowledge: temperature-softened KL(teacher || student) mixed with the hard-label loss.

  • analytic_response_distill() – fit a student’s final linear head to teacher logits by one regularized least-squares solve, without sampling, autograd, or an iterative optimizer.

  • planned_response_distill() – analytic head projection followed by optional minibatched, automatically routed refinement of the remaining mismatch.

  • multi_teacher_distill() – distill from several teachers via averaged (or weighted) soft targets.

  • hint_distill() – FitNets feature/hint transfer: match an intermediate student feature to a teacher feature (a learned linear regressor bridges a dimensionality gap), read out with forward hooks.

  • attention_transfer() – Zagoruyko-Komodakis: match spatial attention maps sum_c F_c^2 between a teacher and student layer.

  • relational_distill() – RKD: match the pairwise distance (and angle) structure of a batch in feature space, so the student preserves relations rather than absolute activations.

  • sequence_level_distill() – Kim-Rush: for a compact LM, train the student on teacher-generated sequences (the teacher’s argmax continuation stands in for the reference).

Everything is deterministic given seed and needs only torch. Each returns a compact result record with the before/after fidelity numbers the tests assert on.

class DistillResult(student, metric, before, after, lower_is_better, extra=<factory>, history=<factory>)[source]

Bases: object

Outcome of a distillation run: the trained student plus before/after fidelity numbers.

metric names what before/after measure (e.g. "teacher_agreement", "soft_kl", "feature_gap"). improved is the sign-aware verdict – higher-is-better for agreement/accuracy, lower-is-better for a KL/gap. history holds the per-epoch training loss.

Parameters:
property improved: bool

Whether the after metric is better than the before metric.

property gain: float

Signed improvement, always positive when the student got better.

kd_loss(student_logits, teacher_logits, y=None, *, temperature=4.0, alpha=0.9)[source]

The Hinton KD loss: alpha * T^2 * KL(teacher||student) + (1-alpha) * CE(student, y).

With y=None (or alpha=1) it is pure soft-target distillation. Returned as a scalar torch tensor so it composes into any training loop.

Parameters:
Return type:

Any

response_distill(student, teacher, x, y=None, *, temperature=4.0, alpha=0.9, epochs=300, lr=1e-2, seed=0, baseline=True, batch_size=None, optimizer='auto', reset_student=True)[source]

Distill teacher’s soft outputs into student on inputs x (labels y optional for the hard mix).

Trains student in place with kd_loss(). When baseline is set, an identically-initialized copy is trained on the hard labels alone (or, if y is None, left untrained) and the two students’ teacher-agreement is compared – the number the test asserts the KD student wins.

Parameters:
Return type:

DistillResult

analytic_response_distill(student, teacher, x, *, temperature=4.0, ridge=1.0e-6, weights=None, seed=0, reset_student=True)[source]

Project teacher logits into the student’s final linear feature space in one exact solve.

The teacher is evaluated once. The student’s body is held fixed and its final Linear head is solved by weighted ridge regression against row-centered teacher logits (softmax is invariant to a row-wise logit shift). This is the useful symbolic path: recognize a block whose optimum is available analytically, eliminate it, and reserve iterative differentiation for any residual mismatch.

Parameters:
Return type:

DistillResult

planned_response_distill(student, teacher, x, y=None, *, temperature=4.0, alpha=0.9, ridge=1.0e-6, weights=None, refinement_epochs=0, lr=1.0e-2, batch_size=None, optimizer='auto', seed=0)[source]

Run analytic head elimination, then optionally refine all student blocks on cached teacher logits.

Parameters:
Return type:

DistillResult

multi_teacher_distill(student, teachers, x, y=None, *, weights=None, temperature=4.0, alpha=0.9, epochs=300, lr=1e-2, seed=0, optimizer='auto')[source]

Distill an averaged (or weights-weighted) ensemble of teachers into student.

The soft target is the weighted mean of the teachers’ softmax probabilities (in probability space, so the ensemble is a genuine mixture). after is the student’s agreement with the ensemble argmax; extra reports the mean single-teacher agreement the ensemble student should beat.

Parameters:
Return type:

DistillResult

hint_distill(student, teacher, x, *, student_layer, teacher_layer, epochs=300, lr=1e-2, seed=0, optimizer='auto')[source]

FitNets hint transfer: drive student’s student_layer feature toward teacher’s teacher_layer.

A learned linear regressor maps the student feature into the teacher’s feature dimension when they differ; the loss is the MSE between the regressed student feature and the (detached) teacher feature, read out with forward hooks. before/after are the normalized feature gap – lower is better.

Parameters:
Return type:

DistillResult

attention_transfer(student, teacher, x, *, student_layer, teacher_layer, beta=1.0, epochs=300, lr=1e-2, seed=0, optimizer='auto')[source]

Match the spatial attention map of student_layer to teacher_layer (Zagoruyko-Komodakis).

Loss is the MSE between the two L2-normalized attention maps sum_c F_c^2. before/after are the attention-map gap – lower is better.

Parameters:
Return type:

DistillResult

relational_distill(student, teacher, x, *, student_layer, teacher_layer, use_angle=True, dist_weight=25.0, angle_weight=50.0, epochs=300, lr=1e-2, seed=0, optimizer='auto')[source]

RKD: match the pairwise distance (and optionally angle) structure of a batch in feature space.

Instead of matching absolute activations, RKD matches relations between samples – so the student need not live in the teacher’s coordinate frame, only preserve its geometry. before/after are the distance- structure gap – lower is better.

Parameters:
Return type:

DistillResult

sequence_level_distill(student, teacher, prompts, *, gen_length, vocab_size, epochs=300, lr=1e-2, seed=0, baseline=True, optimizer='auto')[source]

Kim-Rush sequence-level KD for a compact autoregressive LM.

teacher(context) -> next-token logits generates a hard continuation (its greedy argmax sequence) for each prompt; the student is then trained by teacher forcing to reproduce that teacher-generated sequence. x are integer prompts of shape (N, L0); student maps (N, L) token ids to (N, L, vocab) logits. before/after are the student’s token match to the teacher’s sequences – higher is better.

Parameters:
Return type:

DistillResult