With the help of DeepSeek, this post helps me understand three model compression techniques for mobile AI: Knowledge Distillation, Quantization, and Pruning.

1. Knowledge Distillation

Core Idea

Train a lightweight “student model” to mimic the output behavior of a complex “teacher model”, thereby transferring the teacher model’s knowledge into the student model.

• Source of knowledge: the teacher model’s output probability distribution (soft labels), intermediate-layer features, or attention mechanisms.
• Goal: the student model reaches performance close to the teacher model while staying small.

Example

Scenario: an image classification task (e.g. the ImageNet dataset)

• Teacher model: a large model (e.g. ResNet-50, 76% accuracy).
• Student model: a lightweight model (e.g. MobileNetV3, only 70% accuracy when trained directly).
• Distillation process:
  1. The teacher model generates “soft labels” for the training data (i.e. a probability distribution over classes, such as [0.7, 0.2, 0.1]).
  2. The student model learns from both the ground-truth labels (hard labels) and the soft labels.
     The loss combines the hard-label loss with a KL-divergence loss on the soft labels: $$L = \alpha L_{CE} + (1 - \alpha) L_{KL}$$ where $L_{CE}$ is the cross-entropy loss and $L_{KL}$ is the KL-divergence loss.

Deeper Thoughts

1.Why can the student model, with fewer parameters, get close to the teacher’s performance?

By analogy with a “derived result”, the student model already has the teacher model’s prior knowledge (the probability distribution) and does not need to learn everything from the ground up. We call this the ability to abstract decision boundaries.

From an information-theoretic perspective: the “useful information” in the teacher model (decision boundaries, feature correlations) is encoded into the student model’s parameters, rather than copying all of the parameters.

2.How is the soft-label probability distribution generated?

Core method: Temperature Scaling
Soft labels are not the teacher model’s raw output used directly; instead, a temperature parameter (Temperature, T) is introduced to smooth the probability distribution so that it conveys the relationships between classes.

Formula:

where is the teacher model’s raw output for class , is the temperature parameter, and is the total number of classes.

The effect of T:

• When , the soft labels are equivalent to the hard labels.
• When , the probability distribution is smoother and carries richer information about the relationships between classes.
• When , the probability distribution approaches a uniform distribution.

3.Why not use soft labels when training the teacher model?

Root cause: the teacher model has a different training objective

• The teacher model’s mission: pursue the highest accuracy, not transfer knowledge
  • The teacher model needs to fit the details in the data as closely as possible; hard labels (definite answers) are a more direct supervision signal.
  • Soft labels would introduce unnecessary “uncertainty” and lower the model’s confidence in the correct class.
• The contradiction in where soft labels come from:
  • In knowledge distillation, the soft labels are generated by a more powerful teacher model (e.g. ResNet-50 teaching MobileNet).
  • If soft labels were used to train the teacher model, another, even stronger model would be needed to generate them, which leads to an infinite regress (who generates the soft labels for that stronger model?)

4.Limitations of knowledge distillation

• Requires a high-quality teacher model: the teacher model must be large enough to provide high-quality soft labels.
• Requires a lot of compute: the teacher model needs substantial computing resources to generate high-quality soft labels.
• Requires a lot of data: the teacher model needs a large amount of data to generate high-quality soft labels.

2. Quantization

Core Idea

Convert the model parameters (weights) and activations from high-precision floating point (e.g. 32-bit) to low-precision values (e.g. 8-bit integers), reducing model size and compute cost.

• Types:
  • Post-training Quantization: quantize an already-trained model directly.
  • Quantization-aware Training: simulate quantization error during training to improve the accuracy of the final quantized model.

Example

Scenario: a speech recognition model on a phone

• Original model: an LSTM-based speech recognition model in FP32 precision, 120MB in size, 50ms latency.
• Quantization steps:
  1. Quantize the weights and activations from FP32 to INT8 (mapping the range to -128~127).
  2. Insert dequantization layers to restore precision at critical compute nodes.
• Result: model size shrinks to 30MB, latency drops to 15ms, accuracy loss under 1%.

Below, using PyTorch as an example, are implementations of post-training quantization and quantization-aware training.

Post-training Quantization

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

import torch
import torch.quantization
from torchvision.models import mobilenet_v2

# Step 1: Load the pretrained model
model = mobilenet_v2(pretrained=True)
model.eval()

# Step 2: Define the quantization config
model.qconfig = torch.quantization.get_default_qconfig('qnnpack')  # mobile-optimized config

# Step 3: Insert observers (Observer) to calibrate the quantization parameters
model_fp32_prepared = torch.quantization.prepare(model)

# Step 4: Run the model on calibration data (random data used here as an example)
input_fp32 = torch.randn(1, 3, 224, 224)  # assume the input size is 224x224
with torch.no_grad():
    model_fp32_prepared(input_fp32)

# Step 5: Convert to a quantized model
model_int8 = torch.quantization.convert(model_fp32_prepared)

# Save the quantized model
torch.save(model_int8.state_dict(), "mobilenet_v2_quantized.pth")

# Check the model size
import os
print("FP32 model size:", os.path.getsize("mobilenet_v2.pth")/1e6, "MB")     # about 14MB
print("INT8 model size:", os.path.getsize("mobilenet_v2_quantized.pth")/1e6, "MB")  # about 3.5MB

Quantization-aware Training

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub

# Step 1: Define a model structure that supports quantization
class QuantizableModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = QuantStub()      # quantization entry point
        self.conv = nn.Conv2d(3, 64, kernel_size=3)
        self.dequant = DeQuantStub()  # dequantization exit point
    
    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.dequant(x)
        return x

# Step 2: Insert fake-quantization nodes
model = QuantizableModel()
model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
model.train()  # switch to training mode
model_prepared = torch.quantization.prepare_qat(model)

# Step 3: Normal training loop (FP32 data required)
optimizer = torch.optim.SGD(model_prepared.parameters(), lr=0.001)
for epoch in range(10):
    for data, target in train_loader:  # assume a data loader already exists
        optimizer.zero_grad()
        output = model_prepared(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()

# Step 4: Convert to the final quantized model
model_int8 = torch.quantization.convert(model_prepared)

Deeper Thoughts

Implementing the quantized computation by hand:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

# Original FP32 computation
W_fp32 = torch.tensor([2.5, -1.3, 0.8], dtype=torch.float32)
x_fp32 = torch.tensor([0.4, 1.2, -0.5], dtype=torch.float32)
y_fp32 = torch.dot(W_fp32, x_fp32)  # output: 2.5*0.4 + (-1.3)*1.2 + 0.8*(-0.5) = -1.56

# Quantize to INT8, assuming the range is [-5, 5]
scale_W = 5 / 127  # symmetric quantization, scale = max(abs(W)) / 127
W_int8 = torch.clamp((W_fp32 / scale_W).round(), min=-128, max=127).to(torch.int8)
scale_x = 5 / 127
x_int8 = torch.clamp((x_fp32 / scale_x).round(), min=-128, max=127).to(torch.int8)

# Integer computation
y_int32 = torch.dot(W_int8.float(), x_int8.float())  # cast to float to avoid overflow
y_dequant = y_int32 * (scale_W * scale_x)  # dequantize

print("FP32 result:", y_fp32.item())        # -1.56
print("Quantized result:", y_dequant.item())     # about -1.55 (a tiny error remains)

The full quantization pipeline:

• Preparation stage
  • Insert observers into the model to collect the distribution of weights and activations in each layer.
  • Code: model_prepared = prepare(model)
• Calibration stage
  • Run the model on representative data; the observers record each layer’s min/max values.
  • Code: model_prepared(input_data)
• Conversion stage
  • Compute the quantization parameters from the calibration results and replace the floating-point operators with quantized operators.
  • Code: model_quantized = convert(model_prepared)

Core formulas:

• Quantization formula: $Q(x) = \text{clamp}\left(\text{round}\left(\frac{x}{\text{scale}} + \text{zero\_point}\right), \text{min}, \text{max}\right)$
  • clamp: restrict the result to the range between min and max
  • round: round to the nearest integer
  • zero_point: the quantization offset, used for calibration, usually 0
• Dequantization formula: $x = \left(\text{Q}(x) - \text{zero\_point}\right) \times \text{scale}$

Practical Applications

• TensorFlow Lite: supports post-training quantization by default and can compress an object detection model (e.g. SSD MobileNet) from 16MB down to 4MB.
• Apple Core ML: runs a quantized StyleGAN model on the iPhone for real-time portrait style transfer.

3. Pruning

Core Idea

Reduce model complexity by removing unimportant parameters (e.g. weights close to zero) or structures (e.g. redundant neurons) from the model.

• Types:
  • Unstructured pruning: remove individual weights (sparsification).
  • Structured pruning: remove entire neurons or channels (better suited to hardware acceleration).

Example

Scenario: compressing a BERT model in natural language processing

• Original model: BERT-base (110 million parameters, 400MB model size).
• Pruning process:
  1. During fine-tuning, prune 30% of the attention heads based on weight magnitude or gradient-based importance scores.
  2. Retrain the remaining parameters to recover accuracy.
• Result: model size reduced to 280MB, inference speed up 1.5x, accuracy on the GLUE benchmark drops by only 0.5%.

Here is an example of implementing pruning in PyTorch:

Unstructured Pruning

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

# Define an example model
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)  # input channels 3, output channels 16
        self.fc = nn.Linear(16*26*26, 10)  # assume the input image size is 28x28

    def forward(self, x):
        x = self.conv1(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()

# --- Pruning steps ---
# Step 1: Choose the pruning target (here, the weights of the conv1 layer)
parameters_to_prune = [(model.conv1, 'weight')]

# Step 2: Apply L1-norm pruning (remove 20% of the weights)
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2  # pruning ratio 20%
)

# Step 3: Inspect the pruning result
print("Weight sparsity after pruning:", 
      torch.sum(model.conv1.weight == 0).item() / model.conv1.weight.nelement())

# Step 4: Permanently remove the pruned weights (optional)
prune.remove(model.conv1, 'weight')

# Step 5: Fine-tune the pruned model
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(5):
    for data, target in train_loader:  # assume a data loader already exists
        optimizer.zero_grad()
        output = model(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()

Why prune with the L1 norm?:
The L1 norm naturally induces sparsity: by minimizing the L1 norm, the model tends to push some weights toward 0 to achieve sparsification, and it is simple to compute.

Structured Pruning

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

from torch.nn.utils.prune import ln_structured, remove_structured

# Step 1: Prune entire channels (based on the L2 norm)
# Prune the output channels of the conv1 layer (remove 20% of the channels)
prune.ln_structured(
    model.conv1,
    name="weight",
    amount=0.2,
    n=2,  # L2 norm
    dim=0  # prune along the output-channel dimension
)

# Step 2: Inspect the weight shape after channel pruning
print("conv1.weight shape after pruning:", model.conv1.weight.shape)  
# original shape [16,3,3,3] → after pruning [13,3,3,3] (assuming 3 channels are removed)

# Step 3: Apply the pruning permanently
remove_structured(model.conv1, 'weight')

# Step 4: Adjust the subsequent layers (important! structured pruning requires adapting the network structure)
# the original fc layer input dimension is 16*26*26; after pruning it becomes 13*26*26 → must be redefined
model.fc = nn.Linear(13*26*26, 10)  # update the input dimension

# Fine-tune the model (same as above)

Why prune with the L2 norm?:

• Avoids extreme values: shrinks uniformly
• Computational efficiency: the L2 norm has relatively low computational complexity

Practical Applications

• NVIDIA’s Nemo framework: applies structured pruning to speech recognition models (e.g. QuartzNet), doubling GPU inference speed.
• Drone obstacle avoidance: a pruned YOLOv5 model detects obstacles in real time on edge devices, with 40% lower power consumption.

Comparison and Combined Use of the Three

A combined-use example:
Google’s MobileNetV4 model combines all three:

1. Use knowledge distillation to transfer knowledge from EfficientNet;
2. Apply mixed-precision quantization to the model (INT8 for some layers, FP16 for critical layers);
3. Prune away 80% of the redundant channels; the final model is 6x smaller and 3x faster, with only a 2% drop in accuracy.

Summary

Knowledge distillation, quantization, and pruning are the three core techniques for compressing AI models on mobile devices:

• Knowledge distillation: transfers knowledge through “teacher-student learning”, suited to transferring model capabilities;
• Quantization: lowers numerical precision, directly shrinking size and accelerating computation;
• Pruning: eliminates redundant parameters, improving hardware execution efficiency.
  In practice, the three are often used together (e.g. a “distillation + quantization + pruning” pipeline) to achieve the most aggressive optimization of mobile AI models while preserving accuracy.

Translated from the Chinese original.