Knowledge Distillation là gì?
Knowledge Distillation (KD) là kỹ thuật huấn luyện mô hình nhỏ (student) để bắt chước hành vi của mô hình lớn (teacher) đã được pre-trained. Được giới thiệu bởi Geoffrey Hinton năm 2015, phương pháp này cho phép ta nén trí tuệ của các mô hình khổng lồ như GPT-4, Claude hay Llama 3.1 vào các mô hình nhỏ hơn đáng kể — từ vài trăm triệu đến vài tỷ tham số — mà vẫn giữ được hầu hết khả năng.
Trong thực tế, việc chạy mô hình 70B parameters đòi hỏi GPU A100/H100 với VRAM từ 80GB trở lên. Knowledge distillation cho phép bạn đạt ~90-95% hiệu suất của teacher model chỉ với ~5-10% kích thước, chạy được trên GPU consumer RTX 4090, Mac Studio M2 Ultra, thậm chí smartphone và Raspberry Pi.
Cơ chế hoạt động của Knowledge Distillation
Thay vì học trực tiếp từ dữ liệu gốc (hard labels), student model học từ “soft labels” của teacher — tức là phân phối xác suất trên toàn bộ vocabulary mà teacher model dự đoán. Soft labels mang thông tin quan trọng hơn: ví dụ khi teacher dự đoán “con mèo” với xác suất 0.6, nó vẫn gán xác suất 0.25 cho “con hổ” và 0.12 cho “con chó” — thông tin này cho student biết “đây là vật có lông, có 4 chân, thường nuôi làm thú cưng”.
| Loại label | Ví dụ phân phối xác suất | Thông tin chứa |
|---|---|---|
| Hard label (gốc) | Mèo: 1.0, Hổ: 0.0, Chó: 0.0 | Chỉ đúng/sai |
| Soft label (teacher) | Mèo: 0.6, Hổ: 0.25, Chó: 0.12 | Quan hệ ngữ nghĩa giữa các lớp |
Các kỹ thuật Knowledge Distillation phổ biến
- Response-based KD: Student học trực tiếp từ logits (đầu ra cuối của teacher). Đơn giản nhất, dùng cho các tác vụ classification cơ bản. Ví dụ: DistilBERT, TinyBERT.
- Feature-based KD: Match activation maps giữa các layer trung gian của teacher và student. Phức tạp hơn nhưng hiệu quả cao với các kiến trúc sâu (ResNet, ViT, Transformer).
- Relation-based KD: Student học mối quan hệ giữa các sample hoặc layer, không chỉ giá trị tuyệt đối. Ví dụ: RKD (Relational Knowledge Distillation) dùng distance và angle giữa feature vectors.
Áp dụng Knowledge Distillation với LLM
Với large language model, distillation đi kèm một số kỹ thuật nâng cao:
- Logit Distillation: Match softmax output giữa teacher (ví dụ Llama 3.1 70B) và student (Llama 3.2 3B). Có thể thêm MSE loss giữa logits hoặc KL divergence.
- Chain-of-Thought Distillation: Teacher sinh reasoning traces cho từng câu hỏi, student học cả “cách nghĩ” chứ không chỉ đáp án cuối. Tích cực nhất với các mô hình reasoning phức tạp.
- MiniLLM / GKD (Generative Knowledge Distillation): Thay vì dùng ground-truth dataset cố định, teacher sinh sample mới dựa trên student’s mistakes. Student học từ những gì nó sai — tương tự curriculum learning.
- Speculative Decoding: Dùng student model làm “draft” để đề xuất candidates, teacher model “verify” — tăng tốc inference real-time mà không giảm chất lượng.
Case study thực tế
Một số sản phẩm distillation thành công trên thị trường:
- DistilBERT (110M params): Đạt 97% BERT-base trên GLUE benchmark, inference nhanh 60%, nhẹ 40%.
- TinyLlama 1.1B: Distilled từ Llama 2 7B, chạy trên CPU/MacBook Air, phù hợp cho edge device và chatbot nội bộ.
- Phi-3.5 Mini (3.8B): Microsoft distill từ mô hình lớn hơn, đạt MMLU 82% — ngang GPT-3.5 nhưng chỉ cần 4GB VRAM.
- Apple OpenELM (1B – 3B): Fine-tune từ distillation, chạy trực tiếp trên iPhone 15 Pro với Neural Engine.
Trade-offs và khi nào nên dùng
Knowledge distillation không phải lúc nào cũng là lựa chọn tốt. Trước khi distill, cân nhắc:
- Latency yêu cầu: Nếu cần inference dưới 50ms trên edge device → distillation bắt buộc.
- Chi phí inference: Chạy teacher model 70B tốn ~$0.01/request trên A100; student 3B chỉ ~$0.0005 trên A10.
- Độ chính xác chấp nhận được: Nếu task đòi hỏi accuracy 99.9% (medical, legal), distillation có thể làm mất edge cases quan trọng.
- Khai báo bản quyền: Nhiều model lớn (GPT-4, Claude) cấm distill output của mình. Ưu tiên mô hình open-source (Llama, Mistral, Qwen).
Công cụ hỗ trợ Knowledge Distillation
- Hugging Face TRL (Transformer Reinforcement Learning): Có sẵn `SFTTrainer` + custom distillation loss cho LLM.
- Unsloth: Hỗ trợ fine-tune + distillation cho Llama, Mistral trên GPU consumer với ít VRAM.
- Microsoft Olive: Auto-optimize distillation pipeline cho ONNX Runtime.
- NVIDIA NeMo: Sklearn-style API cho distillation với báo cáo metrics so với teacher.
Nguồn tham khảo: Hinton et al. – Distilling the Knowledge in a Neural Network, MiniLLM – Smaller Language Models Knowledge Distillation, GKD – Generative Knowledge Distillation
