Knowledge distillation · model nhỏ học từ tín hiệu của model lớn

1Định nghĩa

Knowledge Distillation (chưng cất tri thức) là phương pháp huấn luyện một mô hình nhỏ và nhẹ hơn, gọi là student, bằng cách cho nó bắt chước đầu ra của một mô hình lớn và mạnh hơn, gọi là teacher, đã bị đóng băng trọng số. Mục tiêu kỹ thuật là giữ phần lớn chất lượng của teacher nhưng giảm mạnh lượng bộ nhớ HBM (High Bandwidth Memory, bộ nhớ băng thông cao gắn trên GPU) và độ trễ khi phục vụ thật, vì lúc đó chỉ student phải chạy.

Có hai cách truyền tín hiệu. Với chưng cất theo phân phối (logit-based), teacher và student cùng đọc một đầu vào, và student tối thiểu hóa một hàm mất mát đo độ lệch KL (Kullback–Leibler, số đo mức hai phân phối xác suất khác nhau) để khớp toàn bộ bảng xác suất của teacher thay vì chỉ một nhãn đúng. Với chưng cất theo câu trả lời, teacher sinh văn bản và student học nó như học SFT; cách này không cần logits của teacher nên dùng được khi hai model khác tokenizer, như bài response distillation mô tả.

Ví dụ đời thường: một đầu bếp trưởng truyền nghề cho phụ bếp trẻ. Thay vì chỉ nói món này nêm 2 thìa muối là đúng, bếp trưởng giải thích vì sao 1,5 thìa thì nhạt còn 3 thìa thì mặn gắt, nên người phụ bếp học được cả mức độ "gần đúng" của từng lựa chọn. Ví dụ dừng lại ở chỗ: người phụ bếp có thể hiểu lý do, còn student chỉ khớp phân phối số, và nếu bếp trưởng nhớ nhầm công thức thì học trò cũng nhầm theo.

2Luồng

Sáu bước đầu tạo ra student, bốn bước sau phục vụ nó. Hai kiểu tín hiệu (logits mềm và văn bản teacher sinh) rẽ nhánh ở bước 2 và gặp lại ở bước 5, nơi chỉ còn student cần gradient. Chọn từng nút để mở bài liên quan.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

4F4 · Scale & Training Methods

teacher soft labels → student

4F · Distillation

Mô hình nhỏ (student) học khớp phân phối xác suất của một mô hình lớn (teacher).

Nhận gì

Một mô hình teacher (lớn, đã huấn luyện) và một mô hình student nhỏ hơn, cùng dữ liệu.

Làm gì, vì sao

Student không chỉ học nhãn cứng mà học khớp phân phối xác suất mềm (soft label) của teacher, đo bằng KL (độ lệch giữa hai phân phối). Nhiệt độ T làm phân phối phẳng hơn để lộ thêm thứ tự ưu tiên giữa các token.

Ra gì

Một student nhỏ, chạy nhanh hơn, bắt chước teacher. Checkpoint student dùng để suy luận.

Ví dụ, cách đọc ô

Nhãn ô: teacher soft labels → student. Ví dụ trong lab: teacher cho token đúng xác suất 0.996 còn nhãn cứng chỉ nói là 1; ở T = 2 con số đó thành 0.802 và phần còn lại lộ ra.

Lưu ý

Chất lượng student bị chặn bởi teacher và dữ liệu. Có biến thể distillation ở mức câu trả lời (response distillation), khi student học từ văn bản do teacher sinh ra.

Đọc bài 4F · Distillation↗Bài liên quan · Response / sequence distillation · fine-tune student trên câu trả lời của teacher↗Giai đoạn 11 trong Bài Lab: Distillation↗

flowchart TD
  step1["Chọn input và model teacher"]
  step2["Teacher tạo tín hiệu, không nhận gradient"]
  step3["Student forward bằng tham số riêng"]
  step4["Loss khớp teacher"]
  step5["Backward và AdamW chỉ cập nhật student"]
  step6["Lọc output teacher và đánh giá student độc lập"]
  step7["Phục vụ 1 · Nạp chỉ checkpoint của student"]
  step8["Phục vụ 2 · Prefill prompt bằng student"]
  step9["Phục vụ 3 · Decode từng token trên student"]
  step10["Phục vụ 4 · Dừng, giải mã và đo chi phí"]
  step1 --> step2
  step2 --> step3
  step3 --> step4
  step4 --> step5
  step5 --> step6
  step6 --> step7
  step7 --> step8
  step8 --> step9
  step9 --> step10

3Giải pháp

Bài toán: model lớn tốt nhưng quá đắt để phục vụ

Một model lớn cho chất lượng cao nhưng mỗi request tốn nhiều bộ nhớ GPU, nhiều GPU ghép lại và độ trễ cao. Huấn luyện một model nhỏ từ đầu bằng nhãn cứng thường kém hơn nhiều. Distillation chọn con đường giữa: tận dụng teacher làm nguồn tín hiệu giàu thông tin hơn nhãn thật, rồi chỉ phục vụ student. Cần phân biệt với lượng tử hóa: distillation huấn luyện lại trọng số của student từ tín hiệu teacher, còn quantization chỉ đổi cách biểu diễn số; có thể lượng tử hóa student sau đó nhưng phải đo lại chất lượng.

Khái niệm tối thiểu cần nắm trước khi đi vào các bước

Tại mỗi vị trí, model xuất logits (điểm số cho từng token trong từ điển) và hàm softmax đổi logits thành xác suất cộng bằng 1. Softmax có temperature T chia logits cho T trước khi đổi: T = 1 là bình thường, T lớn làm phân phối phẳng hơn và để lộ các phương án phụ mà teacher cân nhắc. Soft target là phân phối đó của teacher, khác nhãn cứng (hard label) chỉ nói "đáp án là phương án số 1".

KL divergence giữa phân phối teacher P và student Q, viết KL(P ‖ Q), bằng tổng của P_i nhân ln(P_i / Q_i) trên mọi phương án. Nó bằng 0 khi hai phân phối trùng nhau và dương khi chúng khác nhau, nên làm loss được. Cross-entropy trên nhãn cứng là trường hợp riêng: trừ ln xác suất student gán cho đáp án đúng.

Ví dụ chạy xuyên suốt cả bài

Để theo dõi từng bước, ta dùng một vị trí duy nhất với ba token ứng viên A, B, C, tự bịa cho mục đích minh họa. Teacher cho logits (4; 2; 1), student ở một thời điểm huấn luyện cho logits (2; 2,5; 0,5). Đáp án thật (nhãn cứng) là A. Ở T = 1, teacher cho xác suất khoảng (0,844; 0,114; 0,042) và student khoảng (0,348; 0,574; 0,078). Ở T = 2, teacher cho (0,629; 0,231; 0,140) và student cho (0,363; 0,466; 0,171). Teacher tin A nhưng cũng cho B nhiều hơn C, còn student đang nghiêng nhầm về B.

Bước huấn luyện 1 · Chọn input và teacher đã học cho transfer set

Đầu vào là tập prompt (transfer set) và một teacher đã train xong. Thao tác: chọn teacher thể hiện tốt đúng hành vi cần chuyển, làm sạch input, đưa về định dạng mà teacher vận hành tốt, và tách riêng ngay tập đánh giá. Lý do: student chỉ học được từ những gì nó thấy; teacher yếu ở đâu thì tín hiệu ở đó kém, và input ngoài vùng teacher giỏi vẫn cho tín hiệu nhưng là tín hiệu của hành vi không đáng tin. Dữ liệu transfer tốt là dữ liệu mà student sẽ gặp lúc phục vụ: đúng phân phối tác vụ, đúng định dạng, đủ đa dạng.

Đầu ra là transfer set cộng tập giữ riêng. Cách kiểm tra: đếm các nhóm tác vụ, độ dài và định dạng; kiểm tra trùng lặp và rò rỉ giữa transfer set và tập giữ riêng; đo chính teacher bằng thước đo tác vụ trước khi chưng cất để biết nó sai ở đâu. Với logit distillation có thể dùng cả văn bản thô chưa gắn nhãn vì nhãn học chính là phân phối của teacher.

Bước huấn luyện 2 · Teacher tạo tín hiệu học, không nhận gradient

Teacher chỉ chạy chiều tiến. Tùy phương pháp, ta thu một trong hai loại tín hiệu. Loại một là logits tại mỗi vị trí, chia cho temperature T để thành soft targets. Loại hai là văn bản: cho teacher sinh trọn câu trả lời hoặc chuỗi suy luận bằng decode (có thể beam search hoặc greedy) rồi lưu thành dataset tĩnh.

Logits có thể tính ngay trong lúc huấn luyện (online) hoặc tính trước rồi lưu (đổi compute lấy dung lượng đĩa). Lưu đủ logits trên từ điển lớn rất tốn; chỉ giữ top-k là một xấp xỉ khác và cần định nghĩa cách xử lý phần xác suất còn lại. Loại văn bản rẻ ở khâu truy cập vì chỉ cần text, nhưng tốn công sinh và lọc; ghi lại prompt, output và cấu hình sinh để biết student đã học từ đâu. Cách đọc: với logits, xem entropy hoặc độ phẳng của soft targets; với văn bản, đếm số chuỗi, độ dài và tỉ lệ bị loại. Chi phí tạo tín hiệu thuộc ngân sách huấn luyện, không đổ cho phục vụ.

Bước huấn luyện 3 · Student forward bằng tham số riêng của mình

Student là mạng có kiến trúc và số tham số riêng, thường ít lớp hơn hoặc vector hẹp hơn, hoặc là model Dense gọn (mọi tham số đều tham gia với mỗi token) trong khi teacher là MoE khổng lồ (Mixture of Experts, gồm nhiều chuyên gia con mà mỗi token chỉ dùng một vài chuyên gia). Kích thước student được chọn theo ngân sách bộ nhớ và độ trễ lúc phục vụ (kích thước kiến trúc), không suy từ số tham số của teacher. Student có thể khởi tạo ngẫu nhiên hoặc lấy từ một checkpoint nhỏ đã pretrain, hoặc từ teacher (DistilBERT lấy một lớp trong hai khi cùng kích thước ẩn); khởi tạo tốt chỉ giúp hội tụ nhanh hơn, không thay thế tín hiệu teacher trong loss.

Student chạy forward trên cùng input để cho logits_student. Student nhỏ không có nghĩa là huấn luyện rẻ: student vẫn cần backward đầy đủ trên đồ thị của mình cộng chi phí teacher tạo tín hiệu; lợi ích nằm ở khâu phục vụ. Nếu student quá nhỏ (capacity mismatch) nó không khớp hết phân phối, và cách chọn loss và dữ liệu phải tính tới giới hạn này.

Bước huấn luyện 4 · Loss khớp teacher: KL trên soft targets hoặc cross-entropy trên văn bản

Với logits, một công thức minh họa là L = alpha × CE(nhãn thật) + (1 − alpha) × T² × KL(P_teacher ‖ P_student), trong đó cả hai phân phối đều lấy softmax ở cùng temperature T. Tính theo ví dụ xuyên suốt: cross-entropy nhãn cứng ở T = 1 là −ln 0,348 ≈ 1,055. KL ở T = 1 là khoảng 0,537; ở T = 2 là khoảng 0,155, nhân T² = 4 được 0,621. Với alpha = 0,5 và T = 2, L ≈ 0,5 × 1,055 + 0,5 × 0,621 ≈ 0,838. Tất cả là số minh họa tính từ các logits giả định, không phải số đo.

Hệ số T² giữ cho phần loss mềm không bị co nhỏ khi tăng T (gradient của nó co theo 1/T²). Cần nhớ temperature ở đây chỉ dùng khi huấn luyện để làm mềm phân phối; lúc phục vụ student dùng T = 1 và nút temperature của sampling là chuyện khác. T cao không phải lúc nào cũng tốt hơn: thí nghiệm trên MNIST (bộ ảnh chữ số viết tay hay dùng làm bài thử) của bài gốc cho thấy student đủ lớn thì nhiều giá trị T đều tương tự, còn student rất nhỏ thì T mức trung bình tốt hơn đáng kể. Hãy chọn T, alpha và dữ liệu bằng đánh giá student.

Với dạng văn bản, loss đơn giản là cross-entropy có mặt nạ giống SFT, chỉ khác mục tiêu là chuỗi do teacher viết; chất lượng chuỗi teacher quyết định trần học. Nên lọc chuỗi sai hoặc trùng và kiểm tra chuỗi bị cắt cụt vì mẫu cụt làm hỏng nhãn. Biến thể on-policy (GKD) cho student tự sinh chuỗi trước rồi nhờ teacher chấm từng vị trí trên chính chuỗi đó, nhằm tránh lệch giữa dữ liệu lúc học và lúc suy luận. Với chuỗi, chỉ cộng các vị trí hợp lệ và chuẩn hóa theo số vị trí đó.

Bước huấn luyện 5 · Backward và optimizer chỉ cập nhật student

Gradient chảy ngược qua đồ thị của student và AdamW chỉ cập nhật trọng số student. Teacher không nằm trong đồ thị, nên không tốn bộ nhớ cho activation hay trạng thái optimizer của teacher; nếu logits hay văn bản teacher đã xuất sẵn ra đĩa, GPU lúc huấn luyện chỉ cần nạp student. Ngân sách đầy đủ vẫn gồm forward hoặc decode của teacher, lưu trữ tín hiệu và backward của student.

Cách đọc: theo dõi gradient norm và NaN của student, đo thời gian mỗi bước và bộ nhớ đỉnh theo cùng một khối lượng công việc. Chọn on-policy hay không là lựa chọn về dữ liệu, không phải một loss mới bắt buộc; tỉ lệ chuỗi do student tự sinh phải thử cho từng tác vụ, và tài liệu của TRL dùng tên tham số riêng (ví dụ tỉ lệ dữ liệu của student và hệ số nội suy của JSD), nên đọc tài liệu của thư viện đang dùng trước khi đặt tên.

Bước huấn luyện 6 · Lọc output teacher và đánh giá student độc lập

Teacher mạnh đến đâu cũng có thể sinh câu trả lời sai, dài dòng hoặc lỗi định dạng, nên với dạng văn bản cần bước lọc (dùng bộ chấm code, kiểm tra đáp án toán hoặc khử trùng) trước khi dùng. Đếm số mẫu giữ và loại theo từng nguyên nhân, và kiểm tra các nhóm tác vụ còn đại diện sau khi lọc. Bước này không áp dụng máy móc cho logit distillation, vì tín hiệu ở đó là phân phối chứ không phải văn bản.

Đánh giá student phải độc lập: dùng câu hỏi chưa từng vào tập huấn luyện hay tập tạo dữ liệu, chấm bằng đáp án hoặc bài kiểm tra bên ngoài, và đo thêm khả năng chung. Student giống teacher nhiều chưa chứng minh student đúng; nếu teacher sai có hệ thống, student sai theo. Mọi tuyên bố student tốt hơn baseline phải đi kèm thước đo độc lập, các nhóm tác vụ và phân tích lỗi. Lưu checkpoint student cùng chính sách lọc để tái lập; muốn huấn luyện tiếp thì cần thêm trạng thái optimizer.

Bước phục vụ 1 · Nạp chỉ checkpoint của student

Trên cụm phục vụ, worker chỉ nạp checkpoint, config và tokenizer của student; teacher ở lại phía huấn luyện và không bao giờ được gọi. Nhờ student nhỏ hơn nên lượng HBM giữ trọng số và số GPU phải ghép song song giảm, còn quy trình nạp và kiểm tra tương thích giống hệt model thường. Tokenizer và template dùng lúc phục vụ là của student; kiểm tra phiên bản tương thích.

Student nhỏ hơn teacher là điều hay gặp, không phải luật: nếu mục tiêu distill là chuyển hành vi chứ không phải nén thì checkpoint phục vụ vẫn có thể lớn. Hãy đo bộ nhớ và độ trễ của chính student dưới khối lượng thật thay vì suy từ tỉ lệ tham số. Nếu bạn chạy cả teacher cạnh student để bỏ phiếu thời gian thực thì đó là ensemble, không còn là phục vụ một student đã chưng cất.

Bước phục vụ 2 · Prefill prompt bằng student

Worker chạy prefill trên các lớp của student để điền KV cache và tính logits chọn token đầu. Không có cơ chế prefill riêng cho distillation: chunked prefill hay prefix caching chạy như thường. Lợi ích về thời gian chờ token đầu (TTFT) đến từ việc student ít lớp hoặc ít chiều hơn, nên mỗi phép nhân ma trận xong nhanh hơn.

Đừng tìm điểm đặc biệt của distillation ở đây: không có loss, teacher hay temperature distillation. Temperature nếu xuất hiện lúc này là temperature sampling của request. Đổi checkpoint student là đổi khóa cache; không chia sẻ KV giữa teacher và student hay giữa hai student.

Bước phục vụ 3 · Decode từng token trên student

Vòng decode tự hồi quy chạy hoàn toàn trên trọng số và KV cache của student. Vì decode bị giới hạn bởi tốc độ đọc dữ liệu từ HBM (bộ nhớ băng thông cao của GPU), student có trọng số và KV nhỏ hơn rút ngắn thời gian mỗi token và cho phép scheduler gộp batch lớn hơn trên cùng GPU. Với student học từ văn bản teacher (Seq-KD), phân phối thường tập trung quanh đáp án của teacher nên greedy đủ tốt, nhưng vẫn cần kiểm tra trên tập giữ riêng trước khi bỏ beam search.

Các kỹ thuật như speculative decoding hay batching nâng cao là tối ưu của engine, không phải thuộc tính của distillation; student nhỏ chỉ làm mỗi forward rẻ hơn, tốc độ thật vẫn phải đo. Decode không cập nhật trọng số và không gọi teacher dưới bất kỳ hình thức nào; mọi tín hiệu của teacher đã nằm sẵn trong checkpoint. Xem thêm decode.

Bước phục vụ 4 · Dừng, giải mã và đo chi phí thật

Đường kết thúc giống hệt bước hoàn tất suy luận chuẩn: dừng khi gặp EOS, chuỗi dừng hoặc giới hạn độ dài, giải mã token bằng tokenizer của student, rồi giải phóng KV cache. Chất lượng và chi phí mỗi triệu token lúc này đo trực tiếp trên student. Đừng đọc output để suy ngược teacher: response của student phản ánh những gì student đã học, kể cả lỗi chép từ teacher sai, nên mọi đánh giá vận hành vẫn dùng bài kiểm tra độc lập.

Kết luận được rút ra và không được rút ra

Có thể kết luận: trên tập giữ riêng và thước đo độc lập đã nêu, student đạt chất lượng đo được với độ trễ và bộ nhớ phục vụ nhỏ hơn teacher. Không thể kết luận: student giữ toàn bộ khả năng của teacher, T cao luôn tốt hơn, hoặc độ giống teacher bằng độ đúng. Các checkpoint công khai như DeepSeek-R1-Distill (student Dense học từ dữ liệu suy luận của teacher MoE) minh họa hướng đi, nhưng cần đọc kết quả của từng checkpoint thay vì giả định.

Thực hành trên Bài Lab tổng hợp

Giai đoạn 11 (Distillation) của Bài Lab tổng hợp cho bạn so student học từ tín hiệu teacher với student học chỉ bằng nhãn cứng. Trước khi chạy, hãy tự tính lại KL của ví dụ xuyên suốt với T = 1 và T = 2 để chắc bạn đọc được con số loss mà lab in ra; sau đó đối chiếu sự khác nhau giữa độ giống teacher và độ đúng trên tập giữ riêng.

4Profiling

• Độ gần teacher và chất lượng thật

  Đo KL hoặc tỉ lệ đồng ý với teacher trên input giữ riêng, đồng thời đánh giá độ đúng bằng nhãn hoặc bài kiểm tra độc lập. Hai con số này tách nhau: student giống teacher sai vẫn là sai.

• Chi phí teacher và hiệu quả student

  Tính chi phí chạy teacher và lưu tín hiệu vào ngân sách huấn luyện. Đo kích thước checkpoint, độ trễ, throughput và bộ nhớ của student dưới cùng một khối lượng công việc.

• Chất lượng dữ liệu teacher

  Với dạng văn bản, đếm mẫu bị loại theo từng lý do, mức trùng lặp, độ dài và nhóm tác vụ; kiểm tra không lẫn câu hỏi giữ riêng vào dữ liệu huấn luyện.

• Các thành phần của loss

  Theo dõi riêng phần loss mềm và phần nhãn cứng, kiểm tra tỉ trọng tương đối vẫn ổn khi đổi T, và đếm số vị trí hợp lệ để chắc padding hay vùng bị che không góp loss.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Teacher / student

Teacher là model lớn đã học xong và bị đóng băng, cung cấp tín hiệu dạy; student là model nhỏ được cập nhật từ tín hiệu đó và là thứ duy nhất được phục vụ.

Transfer set

Tập input dùng để teacher tạo tín hiệu cho student; có thể là dữ liệu huấn luyện gốc hoặc văn bản chưa gắn nhãn, nhưng phải giống những gì student sẽ gặp lúc phục vụ.

Soft targets

Phân phối xác suất teacher gán cho mọi phương án, làm mềm bằng temperature; chứa cả thông tin về những phương án gần đúng, nhiều hơn một nhãn cứng.

Temperature T (lúc huấn luyện)

Hệ số chia logits trước softmax để làm mềm phân phối; lúc phục vụ student dùng T bằng 1, và T này khác nút temperature của sampling.

Kullback–Leibler divergence (KL)

Số đo mức khác biệt giữa hai phân phối xác suất; bằng 0 khi hai phân phối giống hệt nhau, càng lớn khi student càng lệch teacher.

Trọng số mix nhãn cứng (alpha)

Tỉ trọng giữa loss theo nhãn thật và loss theo teacher; alpha là siêu tham số phải chọn bằng thí nghiệm.

Hệ số T²

Nhân vào phần loss mềm để độ lớn gradient không bị co lại khi tăng temperature.

Response / sequence distillation

Chưng cất qua văn bản teacher sinh ra; student học các chuỗi đó bằng cross-entropy như SFT, không cần logits của teacher.

On-policy distillation

Student tự sinh chuỗi, teacher chấm xác suất từng token trên chính chuỗi đó, để giảm lệch giữa lúc học và lúc suy luận.

Capacity mismatch

Student quá nhỏ so với teacher thì không thể khớp hết phân phối; cần chọn loss và dữ liệu có tính đến giới hạn này.

Curation

Lọc và kiểm tra output của teacher (đúng sai, trùng lặp, định dạng, độ dài) trước khi dùng làm dữ liệu cho student.

Đánh giá độc lập

Đo student bằng đáp án đúng hoặc bài kiểm tra bên ngoài, không chỉ đo mức giống teacher; student giống teacher sai vẫn là sai.

Student checkpoint

Artifact phục vụ duy nhất, gồm trọng số student đã distill xong cùng config và tokenizer; teacher không có mặt lúc phục vụ.

7Bài liên quan

8Nguồn