Module 5 — Pruning and size reduction

Modules 3 and 4 shrank the weights from 32 bits to 8. Pruning attacks a different axis: it removes weights entirely, on the observation that a well-trained network has a large fraction of parameters that contribute almost nothing to the output. Done well, pruning cuts the file size again and, on the right hardware, cuts latency too. Done badly, it produces a "sparse" file that is exactly the same size as the dense one — because writing a zero costs the same as writing any other float.

The premise: not every weight matters

If you sort the weights of a trained MobileNetV2 by absolute value, the smallest 30 percent contribute so little to the forward pass that setting them to zero changes the output by less than a percent. Push that to 50 percent, and the change is still small — with a short fine-tuning pass to recover it. Push to 70 percent, and you approach the edge where the model starts to lose classes.

Magnitude pruning implements exactly this idea: at scheduled intervals during training, the smallest-magnitude weights are set to zero, and the training continues so the surviving weights compensate. The final model has the same shape (no layer is dropped) but many of its weights are exactly zero.

The API

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

pruning_params = {
    "pruning_schedule": tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.0,
        final_sparsity=0.60,
        begin_step=0,
        end_step=2000,
    ),
}

pruned_model = prune_low_magnitude(baseline_model, **pruning_params)

pruned_model.compile(
    optimizer=keras.optimizers.Adam(1e-4),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

callbacks = [tfmot.sparsity.keras.UpdatePruningStep()]

pruned_model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=5,
    callbacks=callbacks,
)

# Strip the pruning wrappers before export.
final_model = tfmot.sparsity.keras.strip_pruning(pruned_model)

Four things worth calling out.

First, the schedule is what makes pruning work smoothly. PolynomialDecay from 0 to 0.60 over 2000 steps means the model starts at its full capacity, and every training step raises the sparsity threshold a little. The network has time to redistribute the load among the surviving weights, instead of losing 60 percent of its capacity all at once.

Second, begin_step and end_step are training steps, not epochs. Pick end_step to fall around 60 to 80 percent through your training: enough time to reach the target sparsity, and enough remaining training to recover.

Third, the UpdatePruningStep callback is not optional. Without it, the pruning schedule does not advance and the sparsity stays at zero. This is the most common "why is my pruning doing nothing" bug.

Fourth, strip_pruning must be called before conversion. The pruning wrappers are training-time constructs; the exported model should be the plain Keras model with zeros baked into its weights.

Sparsity that produces no file-size gain

Here is the trap that catches most first-time users. A pruned model with 60 percent zeros is the same size as the dense model, because a float32 zero takes the same 4 bytes as any other float. Look at the .tflite file after conversion and you see no shrinkage. The pruning worked; the file did not.

The trick is what happens downstream of pruning. Zeros compress extremely well: a run-length or Huffman-based compression scheme collapses long zero runs to a few bits. TensorFlow Lite does not compress the model file by itself, but every mobile store (Google Play, App Store) applies gzip or zstd when delivering the app. A gzip-compressed pruned model is 30 to 50 percent smaller than a gzip-compressed dense model:

File	On disk	After gzip
Dense int8	2.35 MB	2.20 MB
60 percent pruned int8	2.35 MB	1.30 MB

The number that matters for the user's download is the gzipped one, and that is where pruning pays. On disk, both files look identical; over the network and on the install line item, the pruned one is nearly half.

"The file is the same size" is not a failure

Report both numbers when you evaluate pruning: the raw .tflite size and the gzipped size. Reporting only the raw size hides the whole point of pruning. Reporting only the gzipped size hides the fact that runtime memory still holds the zeros.

Structured versus unstructured pruning

The pruning above is unstructured: individual weights are zeroed out, wherever they happen to fall. The file compresses well, but the runtime still has to load and process the zero weights — the shape of the tensor is unchanged. Latency does not improve, energy does not improve.

Structured pruning (block sparsity, channel pruning) zeroes entire rows, columns or channels. The runtime can skip whole blocks or drop entire channels, which does improve latency and energy. The cost is a bigger accuracy hit for the same sparsity, because structured constraints are more restrictive.

For our leaf classifier, unstructured pruning is the right choice: our bottleneck is file size on disk, not latency. For a real-time video classifier on a phone with a strict energy budget, structured pruning would be worth the accuracy trade.

Weight clustering

A second technique often paired with pruning: weight clustering replaces the weights of each layer with values drawn from a small palette — say 16 distinct values per layer.

cluster_weights = tfmot.clustering.keras.cluster_weights
CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

clustering_params = {
    "number_of_clusters": 16,
    "cluster_centroids_init": CentroidInitialization.LINEAR,
}

clustered_model = cluster_weights(baseline_model, **clustering_params)
clustered_model.compile(optimizer=keras.optimizers.Adam(1e-5), loss="...", metrics=["accuracy"])
clustered_model.fit(train_ds, epochs=3)

final_model = tfmot.clustering.keras.strip_clustering(clustered_model)

Clustering does not touch the number of weights, but it dramatically reduces the entropy of the weight tensor. Where pruning turns many weights into one specific value (zero), clustering turns all weights into one of 16 specific values. Both create the redundancy that compression exploits, and both compose: pruning + clustering + int8 quantization can shrink a .tflite file to a quarter of its post-quantization size after gzip.

The trade-off matrix, with size reduction

Variant	Raw size	Gzipped size	Latency	Top-1 accuracy
QAT int8 (module 4)	2.35 MB	2.20 MB	22 ms	0.939
QAT int8 + 60 percent prune	2.35 MB	1.30 MB	22 ms	0.933
QAT int8 + prune + cluster 16	2.35 MB	0.95 MB	22 ms	0.930

Under 1 MB after gzip, download-wise. On disk the runtime still allocates 2.35 MB, which we already know fits. The accuracy erosion is real but slow — the last row still hits 0.93, comfortably above the 0.90 minimum we would tolerate for the leaf classifier.

When not to prune

If your app is going to be installed on Wi-Fi over a stable connection, and the raw file already fits the install budget, pruning adds training time for no user-visible benefit. It is a technique aimed at the delivery pipeline: over-the-air updates, cellular downloads, low-bandwidth markets. For a Wi-Fi-only enterprise app, skip it.

Key takeaways

• Magnitude pruning zeroes the smallest weights on a scheduled ramp during training, and the surviving weights compensate; the target sparsity should reach 50 to 60 percent for MobileNet-class models.
• Pruning does not shrink the raw file; it creates structure that gzip and zstd exploit, so always report both raw and compressed sizes.
• Unstructured pruning shrinks the file, structured pruning also shrinks latency and energy at a higher accuracy cost.
• Pruning composes with quantization and with weight clustering; on the leaf classifier, the stack pushes the download size under 1 MB with a small extra accuracy cost.

Next module: the interpreter and its hardware delegates, which turn a .tflite file into a running inference on a specific phone.