Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AMD: Automatic Multi-step Distillation of Large-scale Vision Models

Cheng Han, Qifan Wang, Sohail A. Dianat, Majid Rabbani, Raghuveer M. Rao, Yi Fang, Qiang Guan, Lifu Huang, Dongfang Liu

TL;DR

This paper tackles the challenge of compressing large transformer-based vision models for resource-constrained deployment by proposing Automatic Multi-step Distillation (AMD). AMD automatically discovers an optimal teacher-assistant (TA) through a three-stage cascade: Structural Pruning to generate TA candidates, Joint Optimization to evaluate them with shared parameters, and Optimal Selection using the Negative Performance-Scale Derivative (NPSD) to pick the best TA. The method jointly distills from the full teacher to the selected TA and then to the student, using a composite loss that includes cross-entropy, a distillation logit term, and a feature-mimicking term. Across CIFAR-10, CIFAR-100, and ImageNet, AMD consistently surpasses single-step and multi-step baselines for ViT and Swin architectures, while also delivering significantly faster training, validating its effectiveness for large-scale vision model compression and deployment on resource-limited devices.

Abstract

Transformer-based architectures have become the de-facto standard models for diverse vision tasks owing to their superior performance. As the size of the models continues to scale up, model distillation becomes extremely important in various real applications, particularly on devices limited by computational resources. However, prevailing knowledge distillation methods exhibit diminished efficacy when confronted with a large capacity gap between the teacher and the student, e.g, 10x compression rate. In this paper, we present a novel approach named Automatic Multi-step Distillation (AMD) for large-scale vision model compression. In particular, our distillation process unfolds across multiple steps. Initially, the teacher undergoes distillation to form an intermediate teacher-assistant model, which is subsequently distilled further to the student. An efficient and effective optimization framework is introduced to automatically identify the optimal teacher-assistant that leads to the maximal student performance. We conduct extensive experiments on multiple image classification datasets, including CIFAR-10, CIFAR-100, and ImageNet. The findings consistently reveal that our approach outperforms several established baselines, paving a path for future knowledge distillation methods on large-scale vision models.

AMD: Automatic Multi-step Distillation of Large-scale Vision Models

TL;DR

This paper tackles the challenge of compressing large transformer-based vision models for resource-constrained deployment by proposing Automatic Multi-step Distillation (AMD). AMD automatically discovers an optimal teacher-assistant (TA) through a three-stage cascade: Structural Pruning to generate TA candidates, Joint Optimization to evaluate them with shared parameters, and Optimal Selection using the Negative Performance-Scale Derivative (NPSD) to pick the best TA. The method jointly distills from the full teacher to the selected TA and then to the student, using a composite loss that includes cross-entropy, a distillation logit term, and a feature-mimicking term. Across CIFAR-10, CIFAR-100, and ImageNet, AMD consistently surpasses single-step and multi-step baselines for ViT and Swin architectures, while also delivering significantly faster training, validating its effectiveness for large-scale vision model compression and deployment on resource-limited devices.

Abstract

Transformer-based architectures have become the de-facto standard models for diverse vision tasks owing to their superior performance. As the size of the models continues to scale up, model distillation becomes extremely important in various real applications, particularly on devices limited by computational resources. However, prevailing knowledge distillation methods exhibit diminished efficacy when confronted with a large capacity gap between the teacher and the student, e.g, 10x compression rate. In this paper, we present a novel approach named Automatic Multi-step Distillation (AMD) for large-scale vision model compression. In particular, our distillation process unfolds across multiple steps. Initially, the teacher undergoes distillation to form an intermediate teacher-assistant model, which is subsequently distilled further to the student. An efficient and effective optimization framework is introduced to automatically identify the optimal teacher-assistant that leads to the maximal student performance. We conduct extensive experiments on multiple image classification datasets, including CIFAR-10, CIFAR-100, and ImageNet. The findings consistently reveal that our approach outperforms several established baselines, paving a path for future knowledge distillation methods on large-scale vision models.
Paper Structure (14 sections, 2 equations, 5 figures, 6 tables)

This paper contains 14 sections, 2 equations, 5 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Large-scale Vision Models
  4. Methodology
  5. Problem Formulation
  6. Negative Performance-Scale Derivative
  7. Definition of NPSD.
  8. Preliminary Results.
  9. Automatic Multi-step Distillation
  10. Experiment
  11. Experimental Setup
  12. Main Results
  13. Diagnostic Experiments
  14. Conclusion and Discussion

Figures (5)

  • Figure 1: A preliminary study on the impact of teacher-assistants with different scales and performance w.r.t. the performance of the student. In (a) and (b), ViT-Tiny and ViT-Base are used as the teacher models, and distilled to a 10% student via teacher-assistants mirzadeh2020improved at different scales on CIFAR-10 and CIFAR-100 krizhevsky2009learning respectively. There are several key observations: ❶ The performance of the teacher-assistant degrades when its scale decreases (see green curve); ❷ The performance of the student varies with different teacher-assistants in scales (see yellow curve); ❸ We ascertain that the Negative Performance-Scale Derivative (NPSD) metric (see §\ref{['subsec:s_p_tradeoff']}) exhibits a positive correlation with the performance of student models (see red curve).
  • Figure 2: Overview of different knowledge distillation approaches. (a) Traditional distillation methods directly distill the teacher to the student. (b) Multi-step distillation methods first distill the teacher to a teacher-assistant (requires a large search), which is then further distilled to the student. (c) Manual Multi-step Distillation (MMD) effectively identifies a set of teacher-assistants with different scales and performs multi-step distillation. (d) Automatic Multi-step Distillation (AMD) efficiently and effectively selects the optimal teacher-assistant through one single optimization, including three stages: Structural Pruning, Joint Optimization and Optimal Selection.
  • Figure 3: The sufficiency of using one teacher-assistant.colors represent the performance of $\rm{AMD_{Min}}$ using zero, one, two, and three teacher-assistants, respectively. The results from the training on the ViT-Base CIFAR-100 is posited at (a), CIFAR-10 at (b).
  • Figure 4: Efficient training schedule. Our proposed AMD enjoys superior performance among the overall training scene.
  • Figure 5: Weight values from different loss objectives. We present the performance on different values of $\alpha$ and $\beta$ from Eq. \ref{['eq:overall_objective']}. The highest performance is marked in red (i.e., $\alpha=0.2, \beta=100$). colors represent performance with respect to different $\beta \in \left\{1, 10, 50, 100\right\}$.