Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Following the Teacher's Footsteps: Scheduled Checkpoint Distillation for Domain-Specific LLMs

Cheng Feng, Chaoliang Zhong, Jun Sun, Yusuke Oishi

TL;DR

This work analyzes when a smaller student model can outperform its larger teacher in domain-specific LLM distillation by partitioning the task into a Student-Favored Subdomain ($SFS$) and a Teacher-Favored Subdomain ($TFS$). It introduces Scheduled Checkpoint Distillation (SCD) to systematically reduce the $TFS$ deficit by mimicking the teacher's convergence trajectory and a sample-wise Adaptive Weighting (AW) to preserve the student's strengths on $SFS$. Theoretical bounds and a proximal teacher selection criterion guide the checkpoint schedule, while AW assigns per-sample distillation weights based on relative training losses, enhancing transfer where the teacher excels. Empirical results on QA, NER, and text classification across English and Japanese domains show that SCD and especially SCD with AW consistently outperform baselines and, in several tasks, allow the student to match or exceed the fine-tuned teacher, enabling more efficient deployment of domain-specific LLMs.

Abstract

Large language models (LLMs) are challenging to deploy for domain-specific tasks due to their massive scale. While distilling a fine-tuned LLM into a smaller student model is a promising alternative, the capacity gap between teacher and student often leads to suboptimal performance. This raises a key question: when and how can a student model match or even surpass its teacher on domain-specific tasks? In this work, we propose a novel theoretical insight: a student can outperform its teacher if its advantage on a Student-Favored Subdomain (SFS) outweighs its deficit on the Teacher-Favored Subdomain (TFS). Guided by this insight, we propose Scheduled Checkpoint Distillation (SCD), which reduces the TFS deficit by emulating the teacher's convergence process during supervised fine-tuning (SFT) on the domain task, and a sample-wise Adaptive Weighting (AW) mechanism to preserve student strengths on SFS. Experiments across diverse domain tasks--including QA, NER, and text classification in multiple languages--show that our method consistently outperforms existing distillation approaches, allowing the student model to match or even exceed the performance of its fine-tuned teacher.

Following the Teacher's Footsteps: Scheduled Checkpoint Distillation for Domain-Specific LLMs

TL;DR

This work analyzes when a smaller student model can outperform its larger teacher in domain-specific LLM distillation by partitioning the task into a Student-Favored Subdomain () and a Teacher-Favored Subdomain (). It introduces Scheduled Checkpoint Distillation (SCD) to systematically reduce the deficit by mimicking the teacher's convergence trajectory and a sample-wise Adaptive Weighting (AW) to preserve the student's strengths on . Theoretical bounds and a proximal teacher selection criterion guide the checkpoint schedule, while AW assigns per-sample distillation weights based on relative training losses, enhancing transfer where the teacher excels. Empirical results on QA, NER, and text classification across English and Japanese domains show that SCD and especially SCD with AW consistently outperform baselines and, in several tasks, allow the student to match or exceed the fine-tuned teacher, enabling more efficient deployment of domain-specific LLMs.

Abstract

Large language models (LLMs) are challenging to deploy for domain-specific tasks due to their massive scale. While distilling a fine-tuned LLM into a smaller student model is a promising alternative, the capacity gap between teacher and student often leads to suboptimal performance. This raises a key question: when and how can a student model match or even surpass its teacher on domain-specific tasks? In this work, we propose a novel theoretical insight: a student can outperform its teacher if its advantage on a Student-Favored Subdomain (SFS) outweighs its deficit on the Teacher-Favored Subdomain (TFS). Guided by this insight, we propose Scheduled Checkpoint Distillation (SCD), which reduces the TFS deficit by emulating the teacher's convergence process during supervised fine-tuning (SFT) on the domain task, and a sample-wise Adaptive Weighting (AW) mechanism to preserve student strengths on SFS. Experiments across diverse domain tasks--including QA, NER, and text classification in multiple languages--show that our method consistently outperforms existing distillation approaches, allowing the student model to match or even exceed the performance of its fine-tuned teacher.
Paper Structure (19 sections, 5 theorems, 24 equations, 5 figures, 1 table)

This paper contains 19 sections, 5 theorems, 24 equations, 5 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Theoretical Analysis
  5. Generalized Single Distillation
  6. When and How Can Students Surpass Their Teachers
  7. Scheduled Checkpoint Distillation
  8. Principles for Schedule
  9. Schedule Metric
  10. Adaptive Weight
  11. Metric for Adaptive Weight
  12. Apply AW to SCD
  13. Experiment
  14. Experimental Setup
  15. Overall results
  16. ...and 4 more sections

Key Result

theorem thmcountertheorem

Let $f_s \in \mathcal{F}_s$ be the student model and $f \in \mathcal{F}$ be the true target function. When learning from $n$ i.i.d. data samples, the student's excess risk is bounded by: where $|\mathcal{F}_s|_C$ measures the complexity of the student's hypothesis class, $O(\cdot)$ is the statistical estimation error decaying at rate $1/\sqrt{n}$, and $\epsilon_s := \inf_{g \in \mathcal{F}_s} \ma

Figures (5)

  • Figure 1: On the Visualization of the proposed methods: Tracking Sample-wise Output Distribution Evolution. Scheduled Checkpoint Distillation (SCD, left) and sample-wise Adaptive Weighting (AW, right). SCD mimics the teacher's training trajectory to ease student learning. AW calibrates distillation weights per sample based on post-SFT output distributions, enhancing transfer where the teacher excels and suppressing it where the student is proficient.
  • Figure 2: Toy Experiment for AW Visualization. Data generated from sine function boundary $y=\sin(0.5x)+\mathcal{N}(0,1)$. Teacher model: 2-layer MLP (128 hidden units, $\sim$16K params). Student models: 2-layer MLP (8 hidden units, $\sim$50 params). (1) SFT Baselines. Distillation variants: (2) Distill Only; (3) TD: $0.5\mathcal{L}_{CE} + 0.5\mathcal{L}_{KL}$; (4) Hard AW: teacher-better: only $\mathcal{L}_{KL}$, student-better: only $\mathcal{L}_{CE}$, others: TD; (5) Soft AW: weighting based on loss ratio $\sigma(\log(\mathcal{L}_S/\mathcal{L}_T))$. (6) AW visualization.
  • Figure 3: Schedule visualization based on the SCD w/ AW method. The blue and orange bars represent Metric 1 and Metric 2 (introduced in Section 4.2), respectively. The horizontal axis indicates the teacher checkpoint ID, while the vertical axis shows the sum of Metric 1 and Metric 2. Red arrows mark the optimal teacher selected.
  • Figure 4: AW Visualization. This demonstrates the distribution of AW on the training set, based on teacher SFT and student SFT.
  • Figure 5: TFS & SFS Visualization on NRNER. In all subfigures, the vertical axis represents the sample-level F1 score difference (teacher – student), while the horizontal axis corresponds to the sample IDs. All plots are sorted in descending order along the vertical axis. SCD outperformed other methods in mitigating TFS, and the addition of AW successfully preserved SFS.

Theorems & Definitions (5)

  • theorem thmcountertheorem: Student Generalization Bound
  • theorem thmcountertheorem: Teacher Generalization Bound
  • theorem thmcountertheorem: Student Distillation Generalization Bound
  • theorem thmcountertheorem: Sufficient Condition for Student Surpassing Teacher
  • theorem thmcountertheorem: Individual Step Checkpoint Distillation Error Bound