Bài 57: Knowledge Distillation — model nhỏ học từ model lớn

Sau bài này bạn sẽ:

• Giải thích được vì sao soft label cung cấp gradient phong phú hơn hard label.
• Viết KD loss function với temperature scaling bằng PyTorch.
• Dùng DistillationTrainer của HuggingFace Transformers để distill BERT-class model.
• Phân biệt response-based, feature-based, relation-based KD.
• Hiểu pattern synthetic data distillation cho LLM và giới hạn pháp lý khi dùng commercial API.
• Quyết định khi nào nên chọn KD thay vì pruning hoặc quantization.

Knowledge Distillation (KD) là kỹ thuật nén model được Hinton, Vinyals, Dean giới thiệu trong "Distilling the Knowledge in a Neural Network" (arXiv:1503.02531, 2015). Ý tưởng: thay vì train model nhỏ từ đầu trên label gốc, hãy train nó để bắt chước output của model lớn đã được train trước.

• Teacher: model lớn, accuracy cao, chi phí inference cao. Ví dụ: Llama-3-70B, BERT-large, ResNet152.
• Student: model nhỏ, rẻ hơn, deploy được trên edge hoặc CPU. Ví dụ: Llama-3-8B, DistilBERT, MobileNetV3.

Teacher không cần online trong lúc production — chỉ cần trong lúc train student. Sau khi student được train xong, deploy student một mình.

KD khác với fine-tuning thông thường ở chỗ: loss function của student không chỉ dùng ground truth label mà còn dùng output distribution của teacher làm supervision signal.

Giả sử bài toán phân loại ảnh 3 class: chó, mèo, thỏ.

Hard label (one-hot): [1, 0, 0] — chỉ biết đây là chó, không biết gì thêm.

Soft label từ teacher: [0.70, 0.25, 0.05] — teacher cho biết ảnh này 70% chó, 25% mèo, 5% thỏ. Thông tin này có nghĩa: ảnh này trông giống mèo hơn là thỏ, có thể do màu lông, hình dáng tai...

Soft label cung cấp hai lợi ích cụ thể:

• Gradient phong phú hơn: cross-entropy trên one-hot chỉ update weight theo class đúng. KL divergence trên soft label update theo toàn bộ distribution → student học nhanh hơn và cần ít data hơn để converge.
• Generalization tốt hơn: student học được inter-class similarity được encode trong teacher. Paper gốc (Hinton 2015) báo cáo student train trên 3% MNIST data với soft label đạt accuracy tương đương train trên 100% data với hard label.

Temperature scaling: teacher logit trước softmax thường rất "nhọn" — class đúng gần như 1.0. Để soft label thực sự "mềm" (có thông tin trong các class phụ), ta chia logit cho temperature T > 1 trước softmax. T cao → distribution phẳng hơn → thông tin inter-class rõ hơn.

Hinton et al. 2015 đề xuất loss kết hợp hai thành phần:

L_KD = α × CE(student_logits, hard_label)
     + (1 - α) × T² × KL(softmax(student_logits/T) || softmax(teacher_logits/T))

Các tham số:

• T (temperature): scale logits trước softmax. Phổ biến: T = 2 đến 10. T = 1 tương đương không dùng soft label. T quá cao (> 20) làm distribution quá phẳng, mất tín hiệu.
• α (alpha): trọng số cho hard label loss. Phổ biến: α = 0.1 đến 0.3 (ưu tiên soft label). α = 1.0 bỏ hoàn toàn soft label.
• T²: nhân thêm để bù lại độ lớn của gradient khi logit bị scale xuống T. Đây là chi tiết quan trọng trong paper gốc — nếu bỏ T², soft loss sẽ bị underweight so với hard loss.

KL divergence đo khoảng cách giữa hai distribution. Khi student distribution gần teacher distribution, KL → 0.

import torch
import torch.nn.functional as F


def kd_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    labels: torch.Tensor,
    T: float = 4.0,
    alpha: float = 0.3,
) -> torch.Tensor:
    """
    KD loss theo Hinton et al. 2015 (arXiv:1503.02531).

    Args:
        student_logits: (batch, num_classes) — raw logit trước softmax.
        teacher_logits: (batch, num_classes) — raw logit của teacher, no_grad.
        labels: (batch,) — ground truth class index.
        T: temperature. Cao hơn → soft label phẳng hơn.
        alpha: trọng số hard label loss. (1-alpha) dành cho soft label.
    """
    # Hard label loss — cross entropy với ground truth
    hard_loss = F.cross_entropy(student_logits, labels)

    # Soft label loss — KL divergence giữa student và teacher ở temperature T
    # log_softmax(student/T) vs softmax(teacher/T)
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",
    ) * (T * T)  # nhân T² để bù gradient scale

    return alpha * hard_loss + (1.0 - alpha) * soft_loss

Training loop:

teacher.eval()  # teacher không update weight
student.train()

optimizer = torch.optim.AdamW(student.parameters(), lr=2e-4)

for batch in dataloader:
    x, y = batch

    # Teacher forward — không cần gradient
    with torch.no_grad():
        teacher_logits = teacher(x)

    # Student forward
    student_logits = student(x)

    loss = kd_loss(
        student_logits=student_logits,
        teacher_logits=teacher_logits,
        labels=y,
        T=4.0,
        alpha=0.3,
    )

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Lưu ý: teacher.eval() bắt buộc — nếu teacher có BatchNorm hoặc Dropout, để train mode sẽ làm teacher output không ổn định.

Response-based (output distillation)

Student học trực tiếp output (logit / probability) của teacher. Đây là phương pháp trong section 4-5. Đơn giản nhất, không cần biết kiến trúc bên trong teacher. Phù hợp khi teacher và student khác architecture hoàn toàn.

Feature-based distillation

Ngoài output, còn match intermediate representation (hidden state, feature map) giữa student và teacher. Ví dụ: loss thêm MSE giữa hidden layer thứ N của student và hidden layer thứ M của teacher.

def feature_distill_loss(student_feat, teacher_feat):
    # student_feat: (batch, d_student) — chiều có thể khác teacher
    # cần projection layer nếu d_student != d_teacher
    return F.mse_loss(student_feat, teacher_feat.detach())

Phù hợp khi student và teacher có kiến trúc tương đồng (cùng họ BERT, cùng họ ResNet). TinyBERT dùng cách này để match attention matrix.

Relation-based distillation

Match mối quan hệ / distance giữa các sample trong batch, thay vì match representation trực tiếp. Ví dụ: gram matrix của feature, pairwise cosine similarity. Ít phổ biến hơn hai cách trên, thường dùng trong task metric learning.

Thực tế: phần lớn use case production dùng response-based. Feature-based chỉ cần khi response-based không đủ (student quá nhỏ, task phức tạp).

Các model đã được distill sẵn, dùng trực tiếp qua HuggingFace:

Load và dùng DistilBERT:

from transformers import pipeline

# DistilBERT đã được distill — dùng trực tiếp
clf = pipeline(
    "text-classification",
    model="distilbert-base-uncased-finetuned-sst-2-english",
)
result = clf("This movie was surprisingly good.")
# [{'label': 'POSITIVE', 'score': 0.9998}]

Nếu muốn fine-tune DistilBERT thêm cho task riêng, coi nó như base model bình thường và fine-tune bằng Trainer chuẩn — không cần distill lại.

Khi cần tự distill từ teacher custom sang student custom, extend Trainer và override compute_loss:

import torch
import torch.nn.functional as F
from transformers import Trainer, TrainingArguments


def kd_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.5):
    hard_loss = F.cross_entropy(student_logits, labels)
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",
    ) * (T * T)
    return alpha * hard_loss + (1.0 - alpha) * soft_loss


class DistillationTrainer(Trainer):
    def __init__(self, *args, teacher_model=None, alpha=0.5, T=2.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher = teacher_model
        self.teacher.eval()  # teacher không train
        self.alpha = alpha
        self.T = T

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        outputs_student = model(**inputs)

        with torch.no_grad():
            outputs_teacher = self.teacher(**inputs)

        loss = kd_loss(
            student_logits=outputs_student.logits,
            teacher_logits=outputs_teacher.logits,
            labels=inputs["labels"],
            T=self.T,
            alpha=self.alpha,
        )
        return (loss, outputs_student) if return_outputs else loss


# Sử dụng
from transformers import AutoModelForSequenceClassification, AutoTokenizer

teacher = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
student = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

training_args = TrainingArguments(
    output_dir="./distilled-student",
    num_train_epochs=3,
    per_device_train_batch_size=32,
    learning_rate=5e-5,
    fp16=True,
)

trainer = DistillationTrainer(
    model=student,
    teacher_model=teacher,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    alpha=0.5,
    T=4.0,
)
trainer.train()

Lưu ý về memory: teacher và student cùng nằm trên GPU trong lúc train. Nếu teacher lớn (BERT-large, T5-3B) và GPU ít VRAM, có thể:

• Load teacher lên CPU, student lên GPU — chậm hơn nhưng fit memory.
• Hoặc pre-compute teacher logit cho toàn bộ train set, lưu ra file, rồi train student offline (không cần teacher forward trong lúc train).

Với LLM, distillation kiểu "bắt chước logit" khó thực hiện vì:

• Teacher thường là API-only (GPT-4, Claude) — không có access logit.
• Vocab size hàng trăm nghìn token → KL divergence trên full distribution rất tốn memory.

Thay vào đó, pattern phổ biến là synthetic data generation: dùng teacher generate text, fine-tune student (SFT) trên text đó.

# ===== Bước 1: Generate data từ teacher =====
import openai

client = openai.OpenAI()

source_prompts = load_prompts("domain_prompts.jsonl")  # prompt domain-specific
teacher_dataset = []

for prompt in source_prompts:
    response = client.chat.completions.create(
        model="gpt-4o",  # teacher
        messages=[{"role": "user", "content": prompt}],
        max_tokens=512,
    )
    teacher_dataset.append({
        "prompt": prompt,
        "response": response.choices[0].message.content,
    })

save_jsonl(teacher_dataset, "teacher_data.jsonl")


# ===== Bước 2: Fine-tune student trên teacher data =====
from trl import SFTTrainer, SFTConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

config = SFTConfig(
    output_dir="./llama3-8b-distilled",
    num_train_epochs=2,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    bf16=True,
)

trainer = SFTTrainer(
    model=model,
    args=config,
    train_dataset=teacher_dataset,
)
trainer.train()

Ví dụ thực tế của pattern này:

• Vicuna (LMSYS, 2023): LLaMA fine-tuned trên 70k conversation từ ChatGPT output.
• Orca (Microsoft, 2023): model nhỏ học từ explanation trace của GPT-4.
• Phi-1 / Phi-2 / Phi-3 (Microsoft, 2023-2024): model nhỏ train trên "textbook-quality" synthetic data generate từ LLM lớn.

Phần 14 đề cập giới hạn pháp lý khi dùng output của commercial API.

Cặp teacher/student phổ biến cho image classification:

• Teacher: ResNet-152, EfficientNet-B7, ViT-Large.
• Student: ResNet-18, MobileNetV3-Small, EfficientNet-B0.

Use case chính: deploy lên edge device (Raspberry Pi, NVIDIA Jetson, mobile). Ví dụ MobileNetV3-Small chỉ ~2.5M param, inference < 10ms trên CPU smartphone.

import torchvision.models as models

teacher = models.efficientnet_b7(weights="IMAGENET1K_V1")
student = models.mobilenet_v3_small(weights=None)  # train từ đầu với KD

# Đảm bảo cả hai cùng output 1000 class (ImageNet)
# nếu num_classes khác, adjust classifier head trước
teacher.eval()

# Dùng hàm kd_loss ở section 5, train loop tương tự

Một số paper distillation CV nổi bật: FitNets (Romero et al., 2015) — feature-based cho CNN; DeiT (Touvron et al., 2020) — distill từ CNN teacher sang ViT student với distillation token.

Sau khi distill xong, cần đo ba chiều:

Accuracy / quality

• Classification: accuracy, F1 macro trên test set hold-out.
• NLP: GLUE benchmark, task-specific metric (Rouge, BLEU cho generation).
• LLM: perplexity trên test corpus, task benchmark (MMLU, HumanEval, MT-Bench).
• So sánh student vs teacher: KD thường giữ được 90-99% performance.
• So sánh student (KD) vs student (hard label only): KD thường tốt hơn 1-3% trên task cùng size.

Size và speed

import time
import torch

def benchmark_latency(model, input_tensor, n_runs=100):
    model.eval()
    # Warmup
    with torch.no_grad():
        for _ in range(10):
            model(input_tensor)
    # Benchmark
    start = time.perf_counter()
    with torch.no_grad():
        for _ in range(n_runs):
            model(input_tensor)
    elapsed = time.perf_counter() - start
    return elapsed / n_runs * 1000  # ms per inference

teacher_ms = benchmark_latency(teacher, sample_input)
student_ms = benchmark_latency(student, sample_input)
print(f"Teacher: {teacher_ms:.1f} ms | Student: {student_ms:.1f} ms | Speedup: {teacher_ms/student_ms:.1f}x")

Validate trên test set riêng

Student chỉ được đánh giá trên test set không dùng trong train. Nếu teacher đã thấy test set (vd teacher là GPT-4 đã được train trên internet), cần dùng test data mới hoàn toàn để tránh data contamination.

Ba kỹ thuật có thể kết hợp: distill trước (model nhỏ) → quantize sau (INT8) → deploy. Trong thực tế, nhiều pipeline production làm đúng thứ tự này.

Nên dùng KD khi

• Teacher có accuracy tốt nhưng serving cost quá cao (latency, memory, API cost).
• Cần deploy model lên edge device (mobile, IoT) — model lớn không fit.
• Có dữ liệu unlabeled lớn — dùng teacher generate label thay vì label tay.
• Cần architecture student khác hoàn toàn teacher (vd teacher là Transformer, student là MobileNet cho real-time video).
• Student cần generalize tốt hơn trên ít data — soft label giúp regularization.

Không nên dùng KD khi

• Model đã đủ nhỏ (< 100M param) — quantization INT8 đủ để giảm thêm mà không cần train lại.
• Teacher không đủ tốt (accuracy thấp, hallucination nhiều) — student sẽ học sai theo teacher.
• Chi phí compute để train student lớn hơn tổng chi phí serving teacher trong thời gian dự kiến — KD không có lợi về kinh tế.
• Không có data để train student — KD không tạo ra dữ liệu, chỉ cải thiện label quality.

Khi dùng output của commercial LLM để train model riêng, cần đọc kỹ Terms of Service:

• OpenAI ToS (mục 3): cấm dùng output để "develop models that compete with OpenAI's products and services".
• Anthropic ToS: có điều khoản tương tự về việc dùng output cho training model competitor.

Thực tế an toàn hơn:

• Distill từ open-source LLM: Llama-3 (Meta), Mistral, Mixtral, Qwen — không có restriction này.
• Dùng output commercial API cho nội bộ công ty (không release model ra ngoài) — rủi ro thấp hơn nhưng vẫn cần đọc ToS cụ thể.
• Nếu không chắc, hỏi legal team trước khi scale up.

Ngoài ra, cần filter teacher output trước khi dùng làm training data. LLM lớn vẫn có thể hallucinate hoặc output sai — student sẽ học sai theo. Thiết lập pipeline kiểm tra chất lượng (ví dụ: dùng model khác validate, hoặc human spot-check subset) trước khi dùng toàn bộ synthetic data.

• T = 1: temperature bằng 1 không làm mềm distribution — soft label giống hard label, mất lợi ích của KD. Bắt đầu thử T = 4, tìm bằng hyperparameter search.
• Quên nhân T²: bỏ T * T trong soft loss làm gradient của soft term bị underscale so với hard term — alpha có hiệu lực khác mong muốn.
• Student quá nhỏ: student capacity không đủ để học distribution của teacher → underfit bất kể T hay alpha. Tăng capacity student hoặc chọn architecture phù hợp hơn.
• Teacher output có noise: teacher accuracy 70% → soft label cũng sai 30% → student học theo sai. Filter hoặc chọn teacher tốt hơn trước khi distill.
• Bias từ teacher truyền xuống student: synthetic data distillation copy cả bias của teacher (về gender, ngôn ngữ, culture). Cần bias evaluation riêng cho student.
• Teacher ở train mode: nếu quên teacher.eval(), BatchNorm và Dropout của teacher chạy theo train mode → teacher output ngẫu nhiên theo batch → student nhận supervision signal không ổn định.
• Không đo latency thực tế: student accuracy tốt nhưng latency trên hardware deploy (CPU edge, ARM) có thể vẫn chưa đạt target. Benchmark trên hardware thực sớm, không chỉ benchmark trên GPU train.
• Data contamination: test set của student trùng với data teacher đã thấy trong training → accuracy bị inflate. Dùng held-out test set mới để đánh giá.

Đây là bài cuối của Module 9 và của Series 5 — AI Systems & Deployment. 9 module đã cover:

• M1 FastAPI: build API inference — endpoint, async, lifespan, streaming.
• M2 Demo UI: Gradio, Streamlit, deploy lên HuggingFace Spaces.
• M3 Vector DB: ChromaDB, Pinecone, Qdrant — index và tìm kiếm embedding.
• M4 LangChain: LCEL, document loader, retriever, memory, tools.
• M5 LangGraph: stateful agent, conditional edge, human-in-the-loop.
• M6 Containerization: Docker, multi-stage build, Compose, deploy cloud.
• M7 MLOps: MLflow, W&B, model registry, DVC, CI/CD.
• M8 Monitoring: logging, metrics, data drift, concept drift, Prometheus+Grafana.
• M9 Production Optimization: Redis cache, batching, rate limiting, API key management, prompt injection defense, quantization, ONNX/TensorRT, knowledge distillation.

Series tiếp theo — Series 6: Projects & Job Readiness — tập trung vào capstone project, xây portfolio, chuẩn bị phỏng vấn AI Engineer.