10. Prune-Distill-Quantize Pipeline

Chapter 10 of 18 · 20 min

KEY INSIGHT

Combining pruning, knowledge distillation, and quantization in a structured pipeline yields better results than applying these techniques in isolation, but the order and hyperparameters require careful tuning. The three compression techniques—pruning, distillation, and quantization—each target different aspects of model redundancy. A well-designed pipeline applies them sequentially to progressively reduce model size while preserving accuracy. However, the order significantly impacts final model quality. ### Pipeline Architecture A typical three-stage pipeline proceeds as follows: ```python class CompressionPipeline: def __init__(self, model, config): self.model = model self.config = config def compress(self, train_loader, eval_loader): # Stage 1: Pruning print("Stage 1: Structured pruning") pruned_model = self.apply_pruning( self.model, self.config.pruning_ratio ) pruned_model = self.finetune_after_pruning( pruned_model, train_loader ) # Stage 2: Knowledge Distillation print("Stage 2: Knowledge distillation") teacher_outputs = self.collect_teacher_logits( self.model, train_loader ) student_model = self.create_student(pruned_model) distilled_model = self.distill( student_model, teacher_outputs, train_loader ) # Stage 3: Quantization print("Stage 3: Quantization-aware training") quantized_model = self.quantize_aware_training( distilled_model, train_loader, self.config.target_precision ) return quantized_model ``` ### Order Considerations Applying techniques in the wrong order causes problems: - **Pruning after quantization**: Harder to identify which weights matter when they're already clustered by quantization boundaries - **Distillation before pruning**: Wastes capacity on weights that will be removed - **Quantization without calibration data**: Poor activation range estimates The recommended order (prune → distill → quantize) removes structural redundancy first, then transfers remaining knowledge to a leaner architecture, and finally reduces numerical precision. ### Hyperparameter Tuning Each stage's hyperparameters affect subsequent stages. A 70% pruned model may not have enough capacity to benefit from distillation, while an 8-bit quantized model may have insufficient gradient signal for effective distillation. ```python def find_optimal_sequence(config_space): results = [] for pruning_ratio in [0.5, 0.6, 0.7]: for distill_temperature in [2.0, 4.0, 8.0]: for target_bits in [8, 6, 4]: pipeline = CompressionPipeline(model, { 'pruning_ratio': pruning_ratio, 'distill_temperature': distill_temperature, 'target_precision': target_bits }) result = pipeline.compress(train_loader, eval_loader) results.append({ 'config': {...}, 'metrics': evaluate_model(result, test_loader) }) return select_pareto_optimal(results) ``` ### Failure Modes The pipeline fails when stages interact poorly: 1. **Catastrophic forgetting during finetuning**: Pruned model loses knowledge before distillation recovers it 2. **Distribution mismatch**: Teacher logits from original model don't match student's learning dynamics after pruning 3. **Quantization degradation compounds**: Each stage loses small amounts of accuracy, which accumulates Using early stopping with validation loss during each stage prevents excessive degradation.

EXERCISE

Implement a three-stage compression pipeline: apply magnitude pruning at 50% sparsity, then distill from the unpruned teacher, then quantize to int8. Measure the cumulative accuracy loss at each stage.