Every team that ships a model eventually hits the same wall: the model that scores best is too expensive to serve. A 9B-parameter Gemma answers support tickets beautifully and costs more per request than the ticket is worth. A 300M-parameter vision backbone hits 94% accuracy and blows the latency budget on a mid-range Android phone.
Quantization and pruning — which we covered in Shrinking TensorFlow Models — squeeze an existing architecture. Knowledge distillation does something different: it trains a smaller architecture to imitate the bigger one. The student never sees only the hard labels; it sees the teacher's full output distribution, which carries far more information per example ("this image is 0.7 husky, 0.2 malamute, 0.01 cat" tells the student something about the world that the one-hot label husky does not).
Distillation has become the default cost-reduction move because open-weight teachers are now good enough to be worth copying, and because a distilled student stacks cleanly with quantization afterwards. This tutorial shows both flavours on the TensorFlow/Keras 3 stack: classic logit distillation for a classifier, and sequence-level distillation of an LLM teacher into a small student.
The mechanics: soft targets and temperature
Given teacher logits z_t and student logits z_s, the distillation loss is the KL divergence between their temperature-softened distributions:
L = alpha * CE(y_true, softmax(z_s)) + (1 - alpha) * T^2 * KL(softmax(z_t / T) || softmax(z_s / T))
Two details matter and are routinely got wrong:
- Temperature
T> 1 flattens both distributions, exposing the teacher's relative ranking of the wrong classes.Tbetween 2 and 5 is the usual range; tune it, it matters more than the architecture choice. - The
T^2factor rescales the distillation gradient, which otherwise shrinks as1/T^2and lets the hard-label term dominate. Leave it out and your ablation will "prove" distillation doesn't work.
alpha weights the ground-truth term. With a strong teacher and plenty of unlabeled data, alpha near 0.1 is common; with a noisy teacher, push it up.
A Keras 3 Distiller model
Keras 3 makes this a subclassed Model with a custom train_step. The train_step below is the TensorFlow-backend version (see our multi-backend tutorial for why that distinction matters).
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import tensorflow as tf
class Distiller(keras.Model):
def __init__(self, student, teacher, alpha=0.1, temperature=4.0):
super().__init__()
self.student = student
self.teacher = teacher
self.teacher.trainable = False
self.alpha = alpha
self.temperature = temperature
self.hard_loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
self.soft_loss = keras.losses.KLDivergence()
def call(self, x, training=False):
return self.student(x, training=training)
def train_step(self, data):
x, y = data
teacher_logits = self.teacher(x, training=False)
with tf.GradientTape() as tape:
student_logits = self.student(x, training=True)
hard = self.hard_loss(y, student_logits)
T = self.temperature
soft = self.soft_loss(
keras.ops.softmax(teacher_logits / T, axis=-1),
keras.ops.softmax(student_logits / T, axis=-1),
) * (T ** 2)
loss = self.alpha * hard + (1.0 - self.alpha) * soft
grads = tape.gradient(loss, self.student.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.student.trainable_variables))
for metric in self.metrics:
metric.update_state(y, student_logits)
out = {m.name: m.result() for m in self.metrics}
out.update({"hard_loss": hard, "soft_loss": soft})
return out
def test_step(self, data):
x, y = data
student_logits = self.student(x, training=False)
for metric in self.metrics:
metric.update_state(y, student_logits)
return {m.name: m.result() for m in self.metrics}
Train it like any other model:
distiller = Distiller(student=student, teacher=teacher, alpha=0.1, temperature=4.0)
distiller.compile(
optimizer=keras.optimizers.AdamW(1e-3, weight_decay=1e-4),
metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
distiller.fit(train_ds, validation_data=val_ds, epochs=30)
student.save("student.keras") # save the student, not the wrapper
Both models must output logits, not softmax probabilities. If your teacher was exported with a softmax head, strip it (keras.Model(teacher.input, teacher.layers[-2].output)) or you will be softening an already-normalized distribution and quietly destroying the signal.
Precompute teacher logits when you can
Running the teacher forward pass every step is the single biggest cost in naive distillation, and it caps your effective batch size. If your augmentation pipeline is deterministic (or absent), precompute once:
@tf.function
def add_teacher_logits(x, y):
return x, y, teacher(x, training=False)
cached = (
base_ds.batch(256)
.map(add_teacher_logits, num_parallel_calls=tf.data.AUTOTUNE)
.unbatch()
.cache("/mnt/fast/teacher_cache")
)
Then the train_step takes (x, y, teacher_logits) and skips the teacher entirely — typically a 2–4x speedup. If you use random augmentation, the teacher must see the same augmented view as the student, so either precompute over a fixed number of augmented epochs or accept the online cost. The tf.data patterns we published earlier apply directly; the teacher forward pass is just another expensive map.
Distillation needs unlabeled data, not more labels
The under-appreciated part: the distillation term needs no ground truth. Any in-domain input your teacher can score is a training example. Teams with 20k labeled rows and 2M unlabeled rows should be training the student on all 2M, with the hard-label term weighted only on the labeled subset:
labeled = labeled_ds.map(lambda x, y: (x, y, 1.0)) # sample weight on the hard term
unlabeled = unlabeled_ds.map(lambda x: (x, -1, 0.0))
train_ds = tf.data.Dataset.sample_from_datasets(
[labeled, unlabeled], weights=[0.3, 0.7]
).batch(128).prefetch(tf.data.AUTOTUNE)
In our experience this is where distillation earns its keep on client engagements — it converts an unlabeled log archive into model quality without a labeling budget.
Distilling a generative teacher
For LLMs, logit-level distillation across a very large vocabulary is possible but memory-hungry. The pragmatic path most teams take is sequence-level distillation: sample high-quality outputs from the teacher, then fine-tune the student on those outputs as ordinary supervised data.
import keras_hub
teacher = keras_hub.models.GemmaCausalLM.from_preset("gemma2_9b_en")
prompts = load_domain_prompts() # real user queries, ideally
synthetic = []
for batch in chunks(prompts, 16):
outs = teacher.generate(batch, max_length=512)
synthetic.extend(zip(batch, outs))
Then filter — this step is not optional. Drop responses that fail a schema check, drop near-duplicates, and score the rest with an eval harness of the kind described in Evaluating and Guardrailing Generative AI. A student trained on unfiltered teacher output inherits every teacher hallucination and amplifies it.
Finally, fine-tune the small student on the surviving pairs much as in our Gemma LoRA tutorial — except that for distillation you usually want a full fine-tune rather than LoRA, since you are trying to move the student's whole distribution, not add one skill.
Two cautions before you start: check that the teacher's licence permits training other models on its outputs (open-weight terms often do, with conditions; many hosted commercial APIs explicitly do not), and record the provenance of every synthetic row in your model card.
Choosing the student architecture
- Same family, fewer layers. The easiest win: take the teacher, drop every other transformer block, initialize the student from the surviving weights. Layer-dropping initialization consistently beats random initialization.
- Match the hidden width where you can. If widths match, you can add an intermediate feature distillation loss (MSE between hidden activations), which speeds convergence noticeably on vision backbones.
- Don't go below roughly a tenth of the teacher's parameters in one hop. If you need a 50x reduction, distill in two stages; the intermediate model is a better teacher for the tiny student than the original.
Evaluate the student honestly
Report three numbers, not one:
| Metric | Why |
|---|---|
| Task accuracy on a held-out set | The obvious one |
| Agreement rate with the teacher | Catches students that are accurate but behave differently on edge cases |
| p95 latency and cost per 1k requests at production batch size | The reason you did this |
Also check calibration. Distilled students are frequently better calibrated than their teachers on head classes and worse on the tail, and a shifted confidence distribution will break any downstream threshold that has been tuned against the teacher.
Stack it with quantization
Distillation and post-training quantization are multiplicative, not redundant. The usual production recipe is: distill to a smaller architecture, then apply QAT or int8 post-training quantization, then export to LiteRT or TF Serving. A 7B teacher to a 2B student to int8 weights is roughly an order-of-magnitude memory reduction end to end, and the accuracy loss is concentrated in the distillation step, not the quantization step — which is why you should evaluate after each stage separately. Our LiteRT migration notes cover the export side, and TF Serving covers the server side.
When not to distill
If the teacher is barely above your accuracy floor, distillation will push you under it — you cannot distil headroom you don't have. If your traffic is low enough that the teacher's serving cost is not a line item, distillation is engineering effort spent on nothing. And if your task depends on rare tail behaviour (long-tail entity extraction, safety-critical classes), test that tail specifically: it is the first thing a student loses.
Distillation is one of the highest-return optimizations available on the TensorFlow stack, but it is a training project, not a config change. If you want help sizing the student, building the synthetic-data pipeline, or validating the accuracy and cost trade-off before committing, get in touch — our TensorFlow consultants do this work on client infrastructure.