Knowledge Distillation Schedule

The temperature parameter (T) in the softmax function controls the smoothness of the teacher's output probability distribution. A schedule typically starts with a high temperature (e.g., T=10-20) to create a 'soft' target rich in dark knowledge, then anneals it down to 1 over the course of training. This gradually shifts the student's focus from learning the teacher's generalized class relationships to matching the final, hard predictions.

• High-T Phase: Student learns inter-class similarities (e.g., that a 'cat' is more like a 'dog' than a 'car').
• Annealing Phase: Student's task becomes progressively more precise.
• Low-T/Final Phase: Student converges on the teacher's definitive classifications.

The total loss is a weighted sum of the distillation loss (student vs. teacher soft targets) and the student loss (student vs. ground truth hard labels). The weighting parameter α defines their balance.

A standard schedule progresses through phases:

1. Initialization/Warm-up: High α (e.g., 0.9). Student primarily mimics the teacher's softened outputs to bootstrap learning.
2. Balanced Phase: α is reduced (e.g., to 0.5-0.7). Student learns from both teacher and true labels.
3. Final Fine-tuning: Low α (e.g., 0.1-0.3). Student refines its predictions against the ground truth to ensure it matches hard accuracy metrics.

This progression ensures the student first captures the teacher's generalization before specializing.

A distillation schedule is often decomposed into distinct, sequential training phases:

• Teacher Training/Frozen Phase: The large teacher model is pre-trained (or loaded pre-trained) and its weights are frozen. It acts as a static source of knowledge.
• Distillation Phase: The primary phase where the student is trained using the combined loss. The temperature and α schedules are active here.
• Student-Only Fine-tuning Phase: A final phase where the student is trained solely on the ground truth labels (α=0, T=1) to correct any potential bias from the teacher and lock in final task performance.

Some advanced schedules employ co-training, where the teacher's weights are also updated, but this is less common in standard compression scenarios.

Beyond soft label matching, schedules can include intermediate feature distillation. This involves aligning the student's internal feature maps or attention maps with the teacher's at specific layers. The schedule dictates:

• Which layers to match (e.g., every Nth layer, or specific bottleneck layers).
• What representation to match (e.g., normalized feature activations, attention matrices from Transformers).
• The weighting of this auxiliary loss term, which may also be annealed over time.

For example, a schedule might heavily weight feature matching in early epochs to guide the student's representation learning, then phase it out in favor of output distillation later.

This is a multi-generation scheduling strategy. A large teacher distills a medium student. This medium student then becomes the teacher for a smaller student, and so on. The schedule defines the model size reduction per generation and the training regimen for each step.

Key schedule parameters include:

• Number of distillation generations.
• Compression ratio (size reduction) per step.
• Whether to reset or inherit schedules (T, α) for each generation.

This method often yields better results than direct distillation from a very large teacher to a very small student, as each step involves a more manageable knowledge gap.

Advanced schedules incorporate feedback mechanisms rather than fixed timetables.

• Adaptive α/T: The weights α and temperature T are adjusted based on validation performance. If the student loss plateaus, the schedule might increase α to focus more on the teacher's guidance.
• Distillation-Specific Early Stopping: Training is halted not just based on student validation accuracy, but when the divergence between student and teacher predictions on a held-out set stops decreasing. This indicates the student has absorbed all the transferable knowledge it can from that teacher.
• Layer-wise Scheduling: Different layers of the student may have different distillation intensities or phases based on their sensitivity, analogous to layer-wise sensitivity in pruning.

A pruning schedule is a predefined strategy dictating the timing, rate, and criteria for removing parameters from a neural network. Unlike a one-time operation, it plans sparsity introduction over training steps or epochs.

• Key Mechanism: Often uses a sparsity ramp function (e.g., cubic, linear) to increase the percentage of pruned weights from an initial value to a target.
• Objective: Allows the network to adapt gradually, preserving accuracy better than abrupt, one-shot pruning.
• Example: A schedule might prune 20% of weights every 10 epochs until 80% sparsity is reached, with fine-tuning after each pruning step.

A QAT schedule outlines the phased introduction of quantization simulation during training. It defines when to start emulating low-precision arithmetic and how to manage the associated loss.

• Typical Phases: 1) Full-precision warmup to stabilize training. 2) Quantization simulation where fake quantization nodes are inserted. 3) Fine-tuning to recover accuracy under quantized constraints.
• Critical Parameter: The quantization delay (epoch to start QAT) significantly impacts final accuracy. Starting too early can destabilize learning.
• Tools: Frameworks like TensorFlow's tfmot and PyTorch's torch.ao.quantization provide APIs to implement these schedules.

Layer-wise sensitivity measures how much a model's accuracy degrades when a specific layer is compressed (pruned or quantized). It is a foundational metric for creating non-uniform compression schedules.

• Measurement: Often assessed via a sensitivity analysis, compressing each layer in isolation and observing the accuracy drop.
• Scheduling Impact: Informs heterogeneous schedules where sensitive layers are compressed less aggressively or later in the process. For example, early convolutional layers in a vision model are often highly sensitive and may be spared from aggressive pruning.
• Automation: Tools like Neural Magic's SparseML perform automated sensitivity profiling to guide schedule creation.

Adaptive compression is a scheduling paradigm where the rate or type of compression is dynamically adjusted during training based on real-time feedback, rather than following a fixed, pre-defined plan.

• Feedback Signals: Uses metrics like validation loss, gradient norms, or weight distribution statistics to decide when to prune or change quantization bit-width.
• Advantage: Can react to the model's current state, potentially finding a more optimal compression path than a static schedule.
• Example: An algorithm might pause pruning if the validation loss increases beyond a threshold, resuming only after loss recovery.

Multi-stage compression is a scheduling strategy where different compression techniques are applied in separate, sequential phases, typically with recovery fine-tuning between stages.

• Common Sequence: Pruning → Fine-tune → Quantization → Fine-tune. This isolates the effects of each technique and allows accuracy recovery.
• Rationale: Applying pruning and quantization simultaneously can compound distortion and make accuracy recovery difficult. Staging them simplifies the optimization problem.
• Industry Practice: This is a standard methodology for production model deployment to edge devices, as seen in TensorFlow Lite's model optimization toolkit pipeline.

A compression policy is a comprehensive set of rules and algorithms that govern which techniques to apply, their intensity, and their timing across a model's entire lifecycle. It's the overarching plan from which specific schedules are derived.

• Components: Includes the choice of algorithms (e.g., magnitude pruning, QAT), sensitivity thresholds, sparsity targets, and recovery mechanisms.
• Automation: Frameworks for Automated Model Compression (AMC) use reinforcement learning to search for optimal policies per layer.
• Output: The policy generates the concrete pruning schedule, QAT schedule, and knowledge distillation schedule used during training.