Model Distillation: Thu Nhỏ Mô Hình AI Giữ Lại Độ Chính Xác

Model distillation (chưng cất mô hình) là kỹ thuật đào tạo một mô hình nhỏ (student) học cách bắt chước hành vi của một mô hình lớn đã được huấn luyện tốt (teacher). Student học từ hai nguồn: nhãn gốc và soft labels — phân phối xác suất chi tiết của teacher. Kết quả: giảm 10-100 lần kích thước, tăng tốc độ suy luận 5-20 lần, nhưng vẫn giữ được 95-98% độ chính xác của mô hình gốc. Ứng dụng: triển khai AI trên mobile, edge device, hoặc giảm chi phí inference.

Tại sao cần thu nhỏ mô hình AI?

Mô hình AI càng lớn, càng chính xác — nhưng cũng càng tốn kém.

Chi phí triển khai cao. GPT-3.5 với 175 tỷ tham số đòi ~350GB RAM. Chạy trên cloud? $2-5 mỗi triệu token. Startup AI với 100,000 request/ngày sẽ chảy máu $3,000-15,000/tháng chỉ riêng inference.

Không chạy được trên thiết bị thực. Smartphone chỉ có 4-8GB RAM. Camera giám sát IoT? 512MB. Mô hình lớn? Không vừa. Real-time trên thiết bị? Không thể.

Độ trễ cao. Mô hình lớn cần nhiều phép tính. Mỗi request 2-5 giây — quá chậm cho chatbot, nhận diện ảnh, hoặc bất cứ ứng dụng tương tác nào.

Model distillation giải quyết cả ba. Một mô hình student được distill tốt có thể nhỏ hơn teacher 10-100 lần, chạy nhanh hơn 5-20 lần, nhưng vẫn đạt 95-98% độ chính xác. Khác biệt với quantization hay pruning? Bạn thiết kế lại kiến trúc từ đầu — tối ưu cho thiết bị đích, không bị gò bó bởi mô hình gốc.

Knowledge distillation hoạt động như thế nào?

Ý tưởng cốt lõi

Một mô hình lớn đã học được nhiều hơn chỉ “phân loại đúng”. Nó hiểu mối quan hệ giữa các lớp. Khi phân loại một bức ảnh chó, teacher model không chỉ cho điểm 100% “chó” — nó có thể cho:

• Chó: 85%
• Sói: 10%
• Cáo: 3%
• Mèo: 1%
• Khác: 1%

Con số 10% cho “sói” và 3% cho “cáo” chứa tri thức quan trọng: chó gần với sói và cáo hơn là với mèo. Hard label (chỉ “chó”) mất hết thông tin này.

Soft labels (phân phối xác suất đầy đủ) giàu thông tin hơn hard labels nhiều lần. Student học từ soft labels sẽ nắm bắt được cấu trúc tri thức teacher đã học — và đạt độ chính xác cao hơn nhiều so với chỉ học từ data gốc.

Công thức distillation

Loss function của student kết hợp hai phần:

L_total = α × L_hard + (1-α) × L_soft

• L_hard: Cross-entropy với nhãn gốc (hard labels) — đảm bảo student vẫn học được nhiệm vụ chính
• L_soft: KL divergence giữa soft labels của teacher và student — học cấu trúc tri thức
• α: Trọng số cân bằng (thường 0.1-0.3, nghĩa là ưu tiên học từ teacher)

Temperature (T) làm mềm phân phối xác suất:

P_i = exp(z_i / T) / Σ exp(z_j / T)

• T = 1: Phân phối gốc (85% chó, 10% sói…)
• T = 5: Phân phối mềm hơn (70% chó, 15% sói…) — làm nổi bật các xác suất nhỏ
• T càng cao, student càng học được nhiều từ các dự đoán phụ của teacher

Thực tế thường dùng T = 3-5 trong quá trình distillation, rồi set T = 1 khi inference.

Quy trình huấn luyện

1. Huấn luyện teacher model trên toàn bộ dataset (hoặc dùng mô hình pre-trained)
2. Chạy teacher trên training data → lưu soft labels
3. Thiết kế kiến trúc student (nhỏ hơn nhiều lần: ít layer, ít hidden units, hoặc kiến trúc khác hoàn toàn)
4. Huấn luyện student với loss kết hợp (L_hard + L_soft)
5. Fine-tune student nếu cần (trên hard labels)
6. Đánh giá trên test set

Các phương pháp distillation nâng cao

1. Response-based distillation (cơ bản)

Student học từ output cuối cùng của teacher (logits hoặc xác suất). Đơn giản, hiệu quả, dễ triển khai.

Ưu điểm: Đơn giản, không cần truy cập vào kiến trúc bên trong của teacher
Nhược điểm: Chỉ học từ output cuối, bỏ qua thông tin từ các lớp trung gian

2. Feature-based distillation

Student học cả output cuối lẫn feature maps của các lớp trung gian. Thêm loss để khớp representation ở nhiều lớp.

Ví dụ: FitNets, Attention Transfer — student học cách chú ý (attention maps) giống teacher.

Ưu điểm: Student học được cách “suy nghĩ” như teacher, không chỉ kết quả cuối
Nhược điểm: Cần truy cập kiến trúc teacher, khó áp dụng khi teacher là black box

3. Self-distillation

Một mô hình làm cả vai trò teacher và student — huấn luyện nhiều lần, mỗi lần học từ chính phiên bản tốt nhất trước đó.

Ứng dụng: Boosting độ chính xác mà không cần mô hình lớn hơn. ResNet self-distillation tăng được 1-2% accuracy trên ImageNet.

4. Online distillation

Teacher và student được huấn luyện đồng thời. Không cần mô hình teacher pre-trained — hai mô hình học từ nhau.

Ưu điểm: Tiết kiệm thời gian (không cần huấn luyện teacher trước)
Nhược điểm: Phức tạp hơn, cần điều chỉnh kỹ thuật

Ví dụ thực tế: DistilBERT

Teacher: BERT-base (110M tham số, 12 layers, 768 hidden units)
Student: DistilBERT (66M tham số, 6 layers, 768 hidden units)

Kết quả:

• Giảm 40% kích thước (110M → 66M)
• Tăng tốc độ suy luận 60% (do ít layer hơn)
• Giữ được 97% độ chính xác của BERT trên GLUE benchmark

Cách triển khai:

• Dùng 6 layer đầu của BERT làm student (kiến trúc giống nhưng nông hơn)
• Huấn luyện với soft labels từ BERT-base
• Temperature T = 2
• Thêm cosine embedding loss để khớp hidden states

Thực tế production: DistilBERT chạy được trên smartphone (4GB RAM), phản hồi <100ms cho các tác vụ NLP như sentiment analysis, Q&A.

TinyBERT: Đẩy xa hơn nữa

Teacher: BERT-base (110M)
Student: TinyBERT (14.5M tham số, 4 layers, 312 hidden units)

Kết quả:

• Giảm 87% kích thước (110M → 14.5M)
• Tăng tốc độ suy luận 9.4x
• Giữ được 96.8% độ chính xác trên GLUE

Kỹ thuật đặc biệt:

• Two-stage distillation: Distill từ pre-trained BERT trước, rồi mới fine-tune trên task cụ thể
• Transformer distillation: Khớp cả embedding, attention matrices, và hidden states
• Data augmentation: Dùng thêm unlabeled data để teacher tạo soft labels

Ứng dụng: Chatbot trên mobile, voice assistant, real-time translation trên edge device.

So sánh distillation với các kỹ thuật nén khác

Kỹ thuật	Giảm kích thước	Tốc độ suy luận	Độ chính xác	Thay đổi kiến trúc	Công cụ
Quantization	4-8x	2-4x	-1-2%	Không	TFLite, PyTorch
Pruning	2-10x	1.5-3x	-2-5%	Có thể	TensorFlow, PyTorch
Distillation	10-100x	5-20x	-2-5%	Hoàn toàn	Hugging Face, Custom
Kết hợp	50-500x	10-50x	-3-8%	Có thể	Tùy chọn

Khi nào dùng distillation?

• Cần giảm kích thước mạnh (>10x)
• Muốn thay đổi kiến trúc (CNN → MobileNet, BERT → DistilBERT)
• Có teacher model tốt sẵn (pre-trained hoặc đã huấn luyện tốn kém)
• Triển khai trên mobile/edge với tài nguyên rất hạn chế

Khi nào KHÔNG dùng?

• Chỉ cần giảm nhẹ (2-4x) → quantization đơn giản hơn
• Không có teacher model tốt → distillation không hiệu quả
• Thời gian huấn luyện giới hạn → pruning hoặc quantization nhanh hơn

Triển khai distillation với PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, temperature=3.0, alpha=0.7):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()
        
    def forward(self, student_logits, teacher_logits, labels):
        # Hard loss (với nhãn gốc)
        hard_loss = self.ce_loss(student_logits, labels)
        
        # Soft loss (với soft labels từ teacher)
        soft_labels = F.softmax(teacher_logits / self.temperature, dim=1)
        soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
        soft_loss = F.kl_div(soft_student, soft_labels, reduction='batchmean')
        soft_loss *= (self.temperature ** 2)  # Scale theo temperature
        
        # Tổng hợp
        total_loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss
        return total_loss

# Sử dụng
teacher_model.eval()  # Teacher đã được train sẵn
student_model = StudentModel()  # Mô hình nhỏ hơn
criterion = DistillationLoss(temperature=3.0, alpha=0.7)
optimizer = torch.optim.Adam(student_model.parameters(), lr=1e-4)

for images, labels in dataloader:
    with torch.no_grad():
        teacher_logits = teacher_model(images)
    
    student_logits = student_model(images)
    loss = criterion(student_logits, teacher_logits, labels)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Công cụ và framework hỗ trợ

Hugging Face Transformers: Các mô hình distilled sẵn (DistilBERT, DistilGPT-2, DistilRoBERTa) — chỉ cần load và dùng:

from transformers import DistilBertModel, DistilBertTokenizer

model = DistilBertModel.from_pretrained('distilbert-base-uncased')
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')

TensorFlow Model Optimization: API distillation cho TensorFlow/Keras models

PyTorch: Tự implement (linh hoạt nhất) hoặc dùng thư viện như torch-distill

Neural Network Distiller (Intel): Framework chuyên nghiệp cho distillation, pruning, quantization — hỗ trợ nhiều kỹ thuật kết hợp

Thách thức và giải pháp

1. Capacity gap quá lớn

Vấn đề: Student quá nhỏ so với teacher (ví dụ 1000x) → không học nổi
Giải pháp: Dùng teacher assistant — mô hình trung gian (10x nhỏ hơn teacher) làm cầu nối. Student học từ assistant, assistant học từ teacher.

2. Overfitting trên soft labels

Vấn đề: Student học quá kỹ soft labels → không generalize tốt
Giải pháp: Tăng α (trọng số hard loss), thêm regularization, hoặc early stopping

3. Dataset lớn — không chạy được teacher trên toàn bộ

Giải pháp: Unlabeled data distillation — chạy teacher trên subset hoặc unlabeled data để tạo soft labels, rồi distill student

4. Teacher là black box (API model)

Giải pháp: Response-based distillation — chỉ cần output cuối, không cần truy cập kiến trúc bên trong

Kết hợp distillation với các kỹ thuật khác

Distillation + Quantization: Distill trước để có student nhỏ, rồi quantize student xuống int8 → giảm 100-500x so với teacher gốc.

Ví dụ: BERT-base (110M float32 = 440MB) → DistilBERT (66M float32 = 264MB) → DistilBERT int8 (66M int8 = 66MB). Tổng giảm 6.7x kích thước, tăng tốc 10-15x.

Distillation + Pruning: Distill student, rồi prune student thêm 50-70% → mô hình cực kỳ nhỏ gọn.

Ví dụ: MobileBERT (25M) → distill → prune 60% → 10M tham số, chạy được trên smartwatch.

Khi nào distillation là lựa chọn tốt nhất?

Nên dùng khi:

• Bạn có teacher model tốt (đã pre-trained hoặc đã train tốn kém)
• Cần giảm kích thước mạnh (>10x) mà vẫn giữ độ chính xác cao
• Triển khai trên mobile/edge/IoT với tài nguyên rất hạn chế
• Muốn thay đổi kiến trúc hoàn toàn (CNN → lightweight MobileNet, BERT → TinyBERT)
• Chi phí inference cao — cần giảm mạnh chi phí cloud

Không nên dùng khi:

• Không có teacher model tốt — distillation không thể tạo ra tri thức mới
• Chỉ cần giảm nhẹ (2-4x) — quantization đơn giản hơn và nhanh hơn
• Thời gian huấn luyện giới hạn — distillation cần train lại từ đầu
• Data quá ít — student không học được từ soft labels

Đọc thêm:

• AI Model Compression: Nén Mô Hình AI Hiệu Quả Để Triển Khai Thực Tế — Tổng quan các kỹ thuật nén AI bao gồm quantization, pruning, và distillation, cách kết hợp chúng để đạt tỷ lệ nén tối đa.
• Transfer Learning: Tái Sử Dụng Tri Thức AI Tiết Kiệm 90% Chi Phí — Cách tận dụng mô hình pre-trained làm teacher hoặc làm nền tảng cho distillation, giảm thời gian và chi phí huấn luyện đáng kể.
• MLOps: Vận Hành Mô Hình Machine Learning Trong Production — Quy trình triển khai mô hình đã distill lên production, theo dõi hiệu suất và tối ưu chi phí inference.