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: 加载预训练模型
model = mobilenet_v2(pretrained=True)
model.eval()

# Step 2: 定义量化配置
model.qconfig = torch.quantization.get_default_qconfig('qnnpack')  # 移动端优化配置

# Step 3: 插入观察器（Observer）校准量化参数
model_fp32_prepared = torch.quantization.prepare(model)

# Step 4: 用校准数据运行模型（此处用随机数据示例）
input_fp32 = torch.randn(1, 3, 224, 224)  # 假设输入尺寸为224x224
with torch.no_grad():
    model_fp32_prepared(input_fp32)

# Step 5: 转换为量化模型
model_int8 = torch.quantization.convert(model_fp32_prepared)

# 保存量化模型
torch.save(model_int8.state_dict(), "mobilenet_v2_quantized.pth")

# 检查模型大小
import os
print("FP32模型大小:", os.path.getsize("mobilenet_v2.pth")/1e6, "MB")     # 约14MB
print("INT8模型大小:", os.path.getsize("mobilenet_v2_quantized.pth")/1e6, "MB")  # 约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: 定义支持量化的模型结构
class QuantizableModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = QuantStub()      # 量化入口
        self.conv = nn.Conv2d(3, 64, kernel_size=3)
        self.dequant = DeQuantStub()  # 反量化出口
    
    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.dequant(x)
        return x

# Step 2: 插入伪量化节点
model = QuantizableModel()
model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
model.train()  # 切换到训练模式
model_prepared = torch.quantization.prepare_qat(model)

# Step 3: 正常训练流程（需使用FP32数据）
optimizer = torch.optim.SGD(model_prepared.parameters(), lr=0.001)
for epoch in range(10):
    for data, target in train_loader:  # 假设已有数据加载器
        optimizer.zero_grad()
        output = model_prepared(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()

# Step 4: 转换为最终量化模型
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

# 原始FP32计算
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)  # 输出：2.5*0.4 + (-1.3)*1.2 + 0.8*(-0.5) = -1.56

# 量化到INT8（范围假设为[-5, 5]）
scale_W = 5 / 127  # 对称量化，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)

# 整数计算
y_int32 = torch.dot(W_int8.float(), x_int8.float())  # 转为float避免溢出
y_dequant = y_int32 * (scale_W * scale_x)  # 反量化

print("FP32结果:", y_fp32.item())        # -1.56
print("量化结果:", y_dequant.item())     # 约-1.55（存在微小误差）

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

# 定义示例模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)  # 输入通道3，输出通道16
        self.fc = nn.Linear(16*26*26, 10)  # 假设输入图像尺寸为28x28

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

model = SimpleCNN()

# --- 剪枝步骤 ---
# Step 1: 选择剪枝目标（这里剪枝conv1层的权重）
parameters_to_prune = [(model.conv1, 'weight')]

# Step 2: 应用L1范数剪枝（剪去20%的权重）
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2  # 剪枝比例20%
)

# Step 3: 查看剪枝效果
print("剪枝后的权重稀疏度：", 
      torch.sum(model.conv1.weight == 0).item() / model.conv1.weight.nelement())

# Step 4: 永久移除剪枝的权重（可选）
prune.remove(model.conv1, 'weight')

# Step 5: 微调剪枝后的模型
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(5):
    for data, target in train_loader:  # 假设已有数据加载器
        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: 剪枝整个通道（基于L2范数）
# 对conv1层的输出通道进行剪枝（移除20%的通道）
prune.ln_structured(
    model.conv1,
    name="weight",
    amount=0.2,
    n=2,  # L2范数
    dim=0  # 沿输出通道维度剪枝
)

# Step 2: 查看通道剪枝后的权重形状
print("剪枝后的conv1.weight形状:", model.conv1.weight.shape)  
# 原始形状[16,3,3,3] → 剪枝后[13,3,3,3]（假设移除3个通道）

# Step 3: 永久应用剪枝
remove_structured(model.conv1, 'weight')

# Step 4: 调整后续层（重要！结构化剪枝需适配网络结构）
# 原fc层输入维度为16*26*26，剪枝后变为13*26*26 → 需要重新定义
model.fc = nn.Linear(13*26*26, 10)  # 修改输入维度

# 微调模型（同上）

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.