Knowledge Distillation

Chưng cất tri thức là kỹ thuật nén model, trong đó một mạng neural "student" nhỏ được huấn luyện để tái tạo hành vi của một model "teacher" lớn hơn và chính xác hơn. Student học từ xác suất đầu ra hoặc đặc trưng nội bộ của teacher, bên cạnh hoặc thay cho nhãn thực tế; các mục tiêu giàu thông tin hơn này giúp student đạt độ chính xác cao hơn so với cùng kiến trúc được huấn luyện chỉ bằng nhãn. Kết quả là model đủ nhanh và nhỏ để chạy trên điện thoại, bo mạch nhúng và phần cứng AI biên khác mà vẫn giữ được phần lớn chất lượng của teacher. Năm 2006, Buciluǎ, Caruana và Niculescu-Mizil chứng minh rằng có thể nén tri thức của một ensemble vào một model; năm 2015, Hinton, Vinyals và Dean đặt tên phương pháp này là chưng cất tri thức và giới thiệu các mục tiêu được làm mềm bằng temperature mà hầu hết phương pháp hiện nay vẫn sử dụng.

Cách chưng cất tri thức hoạt động#

Nhãn chỉ cho biết ảnh thể hiện một con sói. Đầu ra softmax của teacher đã huấn luyện cung cấp thêm thông tin: ảnh trông hơi giống husky, hơi giống mèo và chẳng giống ô tô. Các xác suất tương đối này của những lớp sai, thường gọi là tri thức ẩn, được truyền đến student dưới dạng mục tiêu mềm. Chia logits cho một temperature T lớn hơn 1 trước khi áp dụng softmax sẽ làm phẳng phân phối, giúp các xác suất nhỏ có đủ trọng số để mô hình học. Loss chưng cất là độ phân kỳ Kullback-Leibler giữa các phân phối đã làm mềm của teacher và student, nhân với T² để gradient có độ lớn tương đương với gradient của loss nhãn thông thường:

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

Ở T=1, teacher gán xác suất 0.875 cho lớp sói và chỉ 0.006 cho mèo, 0.0 cho ô tô. Ở T=4, cùng các logits đó trở thành 0.493, 0.299, 0.141 và 0.067, cho student biết rằng sói giống husky hơn mèo, còn mèo giống hơn ô tô. Student tối thiểu hóa loss này, ở đây là 1.244, thường dưới dạng tổng có trọng số với loss tác vụ thông thường trên nhãn thực tế. Teacher được giữ nguyên trọng số trong suốt quá trình và chỉ thực hiện forward pass:

Các loại chưng cất#

• Dựa trên đầu ra: Student khớp với đầu ra cuối cùng của teacher như trong phương pháp mục tiêu mềm ở trên, và không cần thông tin nào từ teacher ngoài logits. DistilBERT kết hợp phương pháp này với mô hình hóa ngôn ngữ và loss trạng thái ẩn trong quá trình pretraining, giúp BERT nhỏ hơn 40% và nhanh hơn 60% trong khi vẫn giữ được 97% năng lực hiểu ngôn ngữ.
• Dựa trên đặc trưng: Student khớp với các feature map trung gian của teacher, dùng một projector nhỏ đã học để nối các khác biệt về số kênh. FitNets giới thiệu các "gợi ý" này và huấn luyện một student có số tham số ít hơn teacher khoảng 10.4 lần nhưng đạt kết quả tốt hơn trên CIFAR-10.
• Ở cấp độ chuỗi: Teacher là model ngôn ngữ tạo văn bản, sau đó student được fine-tune trên các đầu ra đó. DeepSeek fine-tune các model R1 đã chưng cất, dựa trên Qwen2.5 và Llama 3, với 800,000 mẫu được tuyển chọn bằng DeepSeek-R1.

Chưng cất dựa trên đặc trưng phù hợp với detector, vốn có đầu ra gồm hàng nghìn box và điểm số thay vì một phân phối lớp. Ultralytics YOLO dùng phương pháp này khi quá trình huấn luyện nhận một teacher distill_model: một model YOLO26 lớn hơn, được đóng băng, cung cấp ba feature map của neck đưa vào đầu phát hiện, còn student khớp các feature map đó qua projector tích chập 1×1 với loss L2 được điều chỉnh theo độ tin cậy của teacher và nhân với dis (mặc định là 6.0). Trên COCO val2017, model nano YOLO26 đã chưng cất đạt 41.5 mAP_50-95, so với 40.9 của cùng model được huấn luyện không có teacher, với chi phí suy luận như nhau.

from ultralytics import YOLO

# Student: model YOLO26 nano nhỏ gọn sẽ được triển khai
model = YOLO("yolo26n.pt")

# Teacher: checkpoint YOLO26 lớn hơn dùng để định hướng huấn luyện và bị loại bỏ sau đó
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

Chưng cất và các thuật ngữ liên quan#

• Transfer learning: Tái sử dụng trọng số của chính model làm điểm khởi đầu cho tác vụ mới. Chưng cất huấn luyện một kiến trúc khác, thường nhỏ hơn, để bắt chước model trên cùng tác vụ.
• Cắt tỉa model: Loại bỏ trọng số, kênh hoặc layer khỏi mạng đã huấn luyện. Model đã cắt tỉa thường được fine-tune bằng cách chưng cất từ model gốc để khôi phục độ chính xác.
• Lượng tử hóa model: Lưu trọng số và activation với độ chính xác thấp hơn, chẳng hạn INT8. Phương pháp này thu nhỏ cùng một kiến trúc và có thể kết hợp với chưng cất, vốn chọn một kiến trúc nhỏ hơn.
• Ensemble model: Lấy trung bình đầu ra của nhiều model tại thời điểm suy luận. Động lực ban đầu của kỹ thuật này là chưng cất một ensemble vào một student duy nhất.
• Chưng cất dataset: Dù tên gọi tương tự, phương pháp này nén dataset thay vì model. Wang và cộng sự nén 60.000 ảnh huấn luyện MNIST thành 10 ảnh tổng hợp, mỗi lớp một ảnh; với một mạng có khởi tạo cố định, các ảnh này giúp đạt gần độ chính xác ban đầu chỉ sau vài bước gradient.