Awesome

Roy Miles¹, Ismail Elezi¹, Jiankang Deng¹

¹ Huawei Noah's Ark Lab

ArXiv Preprint (https://arxiv.org/abs/2403.06213)

Abstract

Knowledge distillation is an effective method for training small and efficient deep learning models. However, the efficacy of a single method can degenerate when transferring to other tasks, modalities, or even other architectures. To address this limitation, we propose a novel constrained feature distillation method. This method is derived from a small set of core principles, which results in two emerging components: an orthogonal projection and a task-specific normalisation. Equipped with both of these components, our transformer models can outperform all previous methods on ImageNet and reach up to a 4.4% relative improvement over the previous state-of-the-art methods. To further demonstrate the generality of our method, we apply it to object detection and image generation, whereby we obtain consistent and substantial performance improvements over state-of-the-art.

Structure

We provide the training code and weights for the two main sets of experiments corresponding for DeIT and ViDT. These are based on the repositories co-advise and ViDT respectively.

.
├── scripts/               # Example scripts for training and evaluation.
├── vidt/                  # Object detection with ViDT models.       
├── main_distributed.py    # Model, optimizer, data loader setup.
├── engine.py              # Train and evaluation loops.

Image Classification

Training data-efficient image transformers using orthogonal knowledge distillation. All models are trained for 300 epochs with a CNN based teacher. The dataset should be downloaded and copied to /data/imagenet2012. The pre-trained regnet-y teacher weights (original DeiT) can be found here. Currently the code is setup to use automatically download the weights from Hugging Face for conveniance. However, we did find the original DeiT teacher weights to be much more effective (+0.5%). We have recently re-ran the experiments for DeiT-Ti and improved the performance to 79.2%. This is with a single GPU and a batch size of 512. Multi-GPU training will likely require changing some learning rate/scheduling hyperparameters to be the most effective.

Pretrained Models

We provide the pre-distilled model weights and logs for the DeIT experiments.

Evaluation

Training

We have tested training with 1 and 2 GPUs using effective batches sizes between 256 and 1024. Using larger batch sizes, or more GPUs, will require modifying the distributed training and/or the learning rates.

Object Detection

Training a more efficient and effective transformer-based object detector. All models are trained for 300 epochs with any larger pre-trained ViDT as the teacher. The MSCOCO dataset should be downloaded and copied to /data/coco. See the vidt/ README for any additional requirements.

Pretrained Models

We provide the pre-distilled model weights and full logs for the ViDT experiments. All the teacher weights can be found here. We only provide the mid-training checkpoints, but are currently re-training the models to release the final checkpoints too.

Evaluation

Training

Citation

@InProceedings{miles2024orthogonalKD,
      title      = {VkD : Improving Knowledge Distillation using Orthogonal Projections}, 
      author     = {Roy Miles, Ismail Elezi and Jiankang Deng},
      booktitle  = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
      year       = {2024},
      month      = {March}
}

If you have any questions, feel free to email me!

Please consider citing our paper and staring the repo if you find this repo useful.