LayerGenius | Smaller Models Without int8: Pruning and Distillation in Keras

Most teams reach for quantization when a model is too big or too slow, and they are right to: int8 is one flag, 4x smaller weights, and it usually works. We wrote about that ladder in Quantizing Keras 3 Models In Place. The question this post answers is what to do when int8 is already in place and you are still over budget — or when accuracy at int8 fell further than you can accept.

The two remaining levers are pruning (remove parameters from the model you have) and distillation (train a smaller model to imitate the one you have). They compose with quantization and with each other, and they fail in different ways. Pick by your constraint: file size, memory bandwidth, or latency on specific hardware.

Decide what the budget actually is

Before touching the model, write down the number you have to hit and on what device. These are not the same problem:

Download/OTA size. Weight bytes dominate. Quantization and unstructured sparsity (which compresses well) both help; latency may not change at all.
Peak memory. Weights plus activations. Shrinking the architecture helps; sparsity on weights often does not, because activations dominate in conv nets at high resolution.
Latency. Only reductions your runtime can exploit help. This is where most pruning work disappoints.

Benchmark the current model on the target device first, with the runtime you will ship (LiteRT with XNNPACK, ONNX Runtime, TF Serving on a specific GPU). A 60% parameter reduction that produces a 3% latency change is a common and entirely predictable result when nobody checked what the kernels do with sparse weights.

Pruning: structured beats unstructured for latency

Unstructured magnitude pruning zeroes individual weights below a threshold. The tensor shape is unchanged, so dense kernels do exactly the same work — you get a file that compresses well and, on most runtimes, identical latency. Useful for OTA size, rarely for speed. The exception is hardware with first-class sparsity support (NVIDIA Ampere-and-later 2:4 structured sparsity through TensorRT), which is a specific pattern, not arbitrary sparsity.

Structured pruning removes whole filters, channels, or attention heads. The resulting model is genuinely smaller and genuinely faster on any runtime, because the shapes shrink. It is also a real architecture change, so it costs more accuracy per parameter removed and requires fine-tuning afterwards.

A workable loop in Keras 3, backend-agnostic:

Score channels on the trained model — L1 norm of each conv filter's weights is a crude but serviceable criterion; BatchNorm gamma magnitude is better when BN follows every conv.
Drop the lowest-scoring fraction per layer (start at 20%, not 60%), and rebuild the architecture with the reduced filters counts.
Copy the surviving weights into the new model, keeping the kept-index order consistent across the conv, its bias, and the following BN.
Fine-tune for a few epochs at a low learning rate — roughly what you would use for the last stage of fine-tuning, around 1e-5 to 1e-4 with Adam.
Re-measure accuracy and device latency. Repeat the ratio upward only while both hold.

Iterative beats one-shot: three rounds of 20% with fine-tuning between them consistently recovers more accuracy than one round of 50%. If you want magnitude sparsity with training-time masking instead, the TensorFlow Model Optimization Toolkit still provides prune_low_magnitude for the TF backend; just be clear-eyed that you are buying bytes, not milliseconds.

Distillation: train the small model you actually wanted

Distillation sidesteps the surgery. You pick a small architecture up front — MobileNetV3, a narrower EfficientNet, a shallower temporal CNN — and train it against the large model's outputs instead of (or alongside) hard labels. The soft probabilities carry information the one-hot labels do not: that this image is 0.7 wheat rust, 0.2 nitrogen deficiency, 0.1 healthy. That extra structure is why a distilled student typically beats the same architecture trained from labels alone, often by a few points on small datasets.

In Keras 3 this is a custom train_step on a Model subclass — the same pattern as any custom training step, so it runs on all three backends if you keep the math in keras.ops:

class Distiller(keras.Model):
    def __init__(self, student, teacher, alpha=0.1, temperature=4.0):
        super().__init__()
        self.student, self.teacher = student, teacher
        self.alpha, self.temperature = alpha, temperature
        self.kl = keras.losses.KLDivergence()

    def compute_loss(self, x=None, y=None, y_pred=None, sample_weight=None):
        teacher_logits = self.teacher(x, training=False)
        T = self.temperature
        soft = self.kl(
            keras.ops.softmax(teacher_logits / T, axis=-1),
            keras.ops.softmax(y_pred / T, axis=-1),
        ) * (T ** 2)
        hard = keras.losses.sparse_categorical_crossentropy(
            y, y_pred, from_logits=True
        )
        return self.alpha * keras.ops.mean(hard) + (1 - self.alpha) * soft

    def call(self, x, training=False):
        return self.student(x, training=training)

Three details decide whether this works:

Students need logits, not softmax. Build both models with linear output layers and apply softmax inside the loss. Distilling from already-normalized probabilities throws away most of the signal you came for.
Temperature between 3 and 6 is the usual useful range. Higher temperature flattens the teacher distribution and emphasizes the relationships between wrong classes. Scaling the soft loss by T**2 keeps gradient magnitudes comparable as you change T; omit it and your effective learning rate moves every time you tune temperature.
Teacher in inference mode, always. training=False on the teacher call. A teacher with live BatchNorm updates or active dropout is a noisy, drifting target.

Distillation also works on unlabeled data, which is often the practical win: you have 2,000 labeled images and 200,000 unlabeled ones from the same line or field. Run the teacher over the unlabeled pool and train the student on soft targets. Set alpha=0 for those batches — there is no hard label to weigh.

Compose, in this order

For an edge deployment that has to shrink substantially, the order that wastes the least time is: distill to a smaller architecture, structurally prune the student if it is still over budget, fine-tune, then quantize to int8 with a representative dataset, then measure on-device. Quantization last, because it is the cheapest step to redo and the one most sensitive to whatever the previous steps did to activation ranges.

Keep the original model as the reference throughout. Every stage gets the same evaluation: task metric at the production operating point on a held-out set, plus p50 and p95 latency on target hardware. A compression pipeline without a per-stage record is impossible to debug six months later, when someone asks why the edge model disagrees with the server one.

When to stop

Compression has a knee. The first 2x is usually nearly free, the second costs a point or two of the metric, the third costs real accuracy. If you are three stages in and still over budget, the honest finding is often architectural — a different input resolution, a two-stage cascade with a cheap gate model, or moving inference off the device. We would rather deliver that conclusion with the measurements behind it than ship a student that hits the latency target and misses every hard case.

If you are working a compression budget on a Keras model and want a second pair of eyes on the trade-offs, get in touch — a senior engineer will reply.