Chưng cất tri thức (Knowledge Distillation)

1. Tổng quan

Chưng cất tri thức (Knowledge Distillation) là kỹ thuật tinh gọn mô hình (Model Compression) huấn luyện một mô hình trò (Student) tương đối nhỏ và nhẹ bắt chước tri thức mà một mô hình thầy (Teacher) lớn và hiệu năng cao đã học được, qua đó giảm mạnh số tham số, khối lượng tính toán và bộ nhớ trong khi vẫn duy trì độ chính xác gần với bản gốc.

Hiệu năng của các mô hình học sâu nhìn chung được cải thiện tỷ lệ với quy mô tham số. Tuy nhiên, các mô hình lớn với hàng tỷ đến hàng trăm tỷ tham số có độ trễ suy luận (latency) lớn, tiêu tốn quá mức bộ nhớ GPU và điện năng, và khó triển khai nguyên vẹn trên các môi trường biên (edge) như điện thoại thông minh, IoT, ECU ô tô. Thực tế, với cùng một bài toán, giảm tham số 10 lần thì chi phí suy luận và thời gian phản hồi được cải thiện vài lần, nhưng nếu chỉ đơn giản huấn luyện một mô hình nhỏ từ đầu (scratch) thì độ chính xác giảm mạnh do thiếu năng lực biểu diễn. Chưng cất tri thức ra đời để giải quyết thế lưỡng nan này. Tức là ý tưởng cốt lõi không phải "huấn luyện tốt một mô hình nhỏ" mà là "để mô hình nhỏ thừa hưởng xu hướng phán đoán của mô hình lớn đã được huấn luyện tốt".

Lý do căn bản khiến chưng cất hiệu quả là đầu ra của mô hình thầy chứa thông tin phong phú hơn nhiều so với việc đúng hay sai (0/1). Ví dụ, khi phân loại chữ số viết tay '7', mô hình thầy biểu diễn độ tương đồng giữa các lớp (dark knowledge, tri thức ngầm) dưới dạng phân phối xác suất như "xác suất là 7 là 0.92, là 1 là 0.05, là 9 là 0.02". So với chỉ cung cấp một nhãn đúng, nếu học cùng những soft target (mục tiêu mềm) như vậy, mô hình trò học được cả cấu trúc quan hệ giữa các lớp nên hiệu năng tổng quát hóa tốt hơn. Kể từ khi Hinton và cộng sự đề xuất năm 2015, chưng cất ngày nay đã trở thành phương tiện tinh gọn tiêu chuẩn cho sLLM (mô hình ngôn ngữ nhỏ), AI trên thiết bị (on-device), dịch vụ suy luận thời gian thực.

2. Cấu trúc cơ bản và nguyên lý của chưng cất tri thức

flowchart LR
  D["Dữ liệu huấn luyện"] --> T["Mô hình thầy (Teacher, lớn)"]
  D --> S["Mô hình trò (Student, nhỏ)"]
  T --> ST["Soft target (phân phối xác suất, làm mềm bằng T)"]
  ST --> LKD["Mất mát chưng cất (KL Divergence)"]
  S --> SP["Phân phối dự đoán của trò"]
  SP --> LKD
  D --> HL["Hard target (nhãn đúng)"]
  HL --> LCE["Mất mát phân loại (Cross-Entropy)"]
  SP --> LCE
  LKD --> LOSS["Tổng mất mát = alpha·mất mát chưng cất + (1-alpha)·mất mát phân loại"]
  LCE --> LOSS
  LOSS --> S

Việc huấn luyện chưng cất được tiến hành bằng cách kết hợp hai hàm mất mát. Thứ nhất là mất mát chưng cất, tối thiểu hóa sự khác biệt giữa soft target của thầy và phân phối đầu ra của trò bằng KL divergence (Kullback-Leibler Divergence). Thứ hai là mất mát phân loại thông thường, tối thiểu hóa cross-entropy giữa đáp án thực (hard target) và đầu ra của trò. Mất mát cuối cùng dung hòa hai số hạng bằng trọng số alpha, thông thường đặt tỷ trọng lớn hơn cho mất mát chưng cất để ưu tiên học xu hướng phán đoán của thầy, đồng thời hiệu chỉnh thiên lệch bằng tín hiệu đáp án.

Trong quá trình này, siêu tham số cốt lõi là nhiệt độ (Temperature, T). Khi chia cho T trong hàm softmax để làm phân phối xác suất trở nên thoải hơn (soften), các chênh lệch xác suất nhỏ của các lớp không phải đáp án vốn bị ép sát về 0 sẽ lộ ra. T=1 thì giống softmax thông thường, còn khi tăng T lên khoảng 2~5 thì thông tin về độ tương đồng tương đối giữa các lớp được khuếch đại, giúp trò hấp thụ tri thức phong phú hơn. Tuy nhiên, nếu T quá lớn thì phân phối trở nên đồng đều và tín hiệu lại yếu đi, nên phải tìm giá trị tối ưu bằng tập kiểm định. Khi huấn luyện, áp dụng cùng T cho cả thầy và trò, còn khi triển khai (suy luận) thì đưa về T=1.

Ví dụ, trên benchmark phân loại ảnh, khi chưng cất tri thức của thầy (ResNet-50, khoảng 25,6 triệu tham số) sang trò (ResNet-18, khoảng 11,7 triệu tham số), có báo cáo cho thấy độ chính xác phục hồi khoảng 1~3%p so với khi huấn luyện trò riêng lẻ, trong khi tốc độ suy luận nhanh hơn gấp 2 lần trở lên. Trong lĩnh vực ngôn ngữ tự nhiên, ví dụ tiêu biểu là DistilBERT chưng cất từ BERT-base (khoảng 110 triệu tham số), giảm khoảng 40% tham số và tăng tốc suy luận khoảng 60% trong khi vẫn duy trì khoảng 97% hiệu năng của bản gốc.

3. Các loại phương thức chưng cất

flowchart TB
  KD["Các loại chưng cất tri thức"] --> RB["Dựa trên phản hồi (Response-based)"]
  KD --> FB["Dựa trên đặc trưng (Feature-based)"]
  KD --> RL["Dựa trên quan hệ (Relation-based)"]
  KD --> SD["Tự chưng cất (Self-Distillation)"]
  RB --> RB1["Bắt chước đầu ra cuối (logit/xác suất)"]
  FB --> FB1["Bắt chước biểu diễn tầng ẩn trung gian"]
  RL --> RL1["Bắt chước quan hệ (cấu trúc) giữa mẫu·tầng"]
  SD --> SD1["Cùng một mô hình tự đóng vai thầy"]

Chưng cất được chia thành các loại tùy theo trò bắt chước 'cái gì' của thầy. Chưng cất dựa trên phản hồi là cách kinh điển·đơn giản nhất, chỉ bắt chước đầu ra cuối (logit·xác suất) của thầy; vì không cần biết cấu trúc bên trong của thầy nên dễ áp dụng cả giữa các kiến trúc khác nhau. Tuy nhiên, do chỉ nhìn đầu ra cuối nên không tận dụng được thông tin chứa trong quá trình suy luận bên trong của thầy.

Chưng cất dựa trên đặc trưng hướng trò mô phỏng cả biểu diễn (feature map) của tầng ẩn (hidden layer) trung gian của thầy. Vì trò hấp thụ từng bước biểu diễn theo tầng của thầy nên có thể chuyển giao tri thức chính xác hơn, nhưng do cấu trúc tầng và số chiều của thầy·trò khác nhau nên cần thêm tầng chiếu (projection) để khớp, và việc huấn luyện trở nên phức tạp. Chưng cất dựa trên quan hệ truyền 'tri thức quan hệ' như quan hệ giữa các mẫu dữ liệu hay cấu trúc tương quan giữa các tầng thay vì từng đầu ra riêng lẻ, bảo toàn cấu trúc hình học của không gian embedding mà thầy đã học.

Mặt khác, tự chưng cất (Self-Distillation) là biến thể không cần thầy lớn riêng biệt, trong đó cùng một mô hình (hoặc tầng sâu) lấy tầng nông·chính nó ở epoch trước làm thầy để học, loại bỏ chi phí huấn luyện thầy trong khi vẫn nâng hiệu năng nhờ hiệu ứng chính quy hóa (regularization). Gần đây, với các mô hình ngôn ngữ lớn, chưng cất mức chuỗi — dùng câu trả lời·quá trình suy luận do mô hình mạnh sinh ra làm dữ liệu để huấn luyện mô hình nhỏ — hay chưng cất Chain-of-Thought truyền cả chuỗi lập luận (rationale) đang được sử dụng rộng rãi.

4. So sánh với các kỹ thuật tinh gọn tương tự

Chưng cất tri thức là một trong nhiều trục của tinh gọn mô hình, và trong thực tế được dùng bổ trợ lẫn nhau cùng với cắt tỉa (Pruning)·lượng tử hóa (Quantization). Ba kỹ thuật có cách tiếp cận khác nhau về căn bản. Cắt tỉa loại bỏ các trọng số·nơ-ron có mức đóng góp thấp để làm mô hình 'mỏng' đi, còn lượng tử hóa biểu diễn số thực dấu phẩy động 32 bit bằng số nguyên 8 bit, v.v. để hạ độ chính xác lưu trữ·tính toán. Ngược lại, chưng cất khác biệt ở chỗ không thay đổi bản thân cấu trúc mà nâng hiệu năng bằng cách 'dạy tốt hơn' cho mô hình nhỏ.

Vì ba kỹ thuật nhắm vào những điểm khác nhau trên đánh đổi hiệu năng-hiệu quả, trong pipeline triển khai thực tế việc kết hợp tuần tự là phổ biến. Ví dụ, thu được trò nhỏ bằng chưng cất từ thầy lớn, cắt tỉa trò đó để loại các kết nối không cần thiết, rồi cuối cùng lượng tử hóa để đưa lên bộ tăng tốc tính toán số nguyên (NPU). Làm như vậy có thể đạt tỷ lệ nén và tốc độ suy luận ở mức khó đạt được bằng một kỹ thuật đơn lẻ.

5. Lưu ý và hàm ý

Từ góc nhìn Kỹ sư chuyên nghiệp (Professional Engineer), việc đưa vào chưng cất tri thức không chỉ là chọn thuật toán mà phải tiếp cận như thiết kế đánh đổi tổng hợp giữa chi phí dịch vụ, khả năng phản hồi, độ chính xác và khả năng bảo trì.

• Chiến lược áp dụng (quản lý khoảng cách thầy·trò): Nếu khoảng cách năng lực giữa thầy và trò quá lớn sẽ phát sinh vấn đề 'khoảng cách năng lực (capacity gap)' khiến trò không theo kịp thầy. Khi đó, chiến lược chưng cất nhiều bước đặt từng bước mô hình trợ giảng (Teacher Assistant) có kích thước trung gian để thu hẹp khoảng cách là hiệu quả.

• Đánh đổi chất lượng và hệ thống kiểm chứng: Mô hình trò đã chưng cất có thể bảo toàn độ chính xác trung bình nhưng đặc tính ở các trường hợp đuôi dài (hiếm), thiên lệch·ảo giác (hallucination) có thể khác thầy. Do đó, trước khi triển khai cần trang bị hệ thống kiểm chứng hồi quy riêng không chỉ cho các chỉ số đại diện mà còn cho kịch bản yếu, tính công bằng, tính an toàn.

• Dữ liệu·giấy phép·quản trị: Đặc biệt với LLM, việc dùng đầu ra của mô hình đóng thương mại làm dữ liệu chưng cất có thể bị ràng buộc bởi điều khoản dịch vụ, nên bắt buộc phải xem xét trước nguồn gốc và giấy phép dữ liệu. Thiên lệch của mô hình thầy có thể được chuyển giao·khuếch đại nguyên vẹn sang trò, nên cần quản lý từ góc độ quản trị dữ liệu.

• Triển vọng (lan rộng on-device·sLLM): Khi nhu cầu AI on-device và sLLM tăng lên, chưng cất đang trở thành kênh cốt lõi đưa năng lực của foundation model xuống thiết bị biên. Trong tương lai, dự kiến pipeline tinh gọn tích hợp kết hợp cắt tỉa, lượng tử hóa, thích ứng hạng thấp (LoRA) và chưng cất truyền cả năng lực suy luận (quá trình tư duy) sẽ được chuẩn hóa.

Tóm tắt một câu: Chưng cất tri thức là kỹ thuật tinh gọn mô hình tiêu biểu, chuyển giao soft target (tri thức ngầm) của mô hình thầy lớn sang mô hình trò nhỏ bằng điều chỉnh nhiệt độ và mất mát KL, nhằm giảm kích thước·tính toán·độ trễ trong khi bảo toàn tối đa độ chính xác, và khi kết hợp với cắt tỉa·lượng tử hóa trở thành chiến lược triển khai cốt lõi của thời đại on-device·sLLM.