mixle.task.distill_methods module

The classic knowledge-distillation FAMILIES as composable, verified procedures on torch modules.

mixle.task.distill distills a frontier teacher’s labels into a tiny 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.

  • 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 small 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 small 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:
student: Any
metric: str
before: float
after: float
lower_is_better: bool
extra: dict[str, Any]
history: list[float]
property improved: bool
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)[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

multi_teacher_distill(student, teachers, x, y=None, *, weights=None, temperature=4.0, alpha=0.9, epochs=300, lr=1e-2, seed=0)[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)[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)[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)[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)[source]

Kim-Rush sequence-level KD for a small 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