{Image Classification; Transformer; Vision Transformer; Teacher Student; Pseudo Label; Distillation}
Paper: https://arxiv.org/abs/2012.12877
Code: https://github.com/facebookresearch/deit/blob/main/README_deit.md
1) Motivation, Objectives and Related Works :
Motivation:
Prior Vision Transformer, ViT, needs to be pre-trained with hundreds of millions of images using external data.
ViT does not generalize well when trained on insufficient amounts of data.
Objectives: Data-efficient image Transformer (DEiT)
Knowledge distillation:
Aims to transfer knowledge from a bigger model, i.e., a teacher network, to a smaller, target model, i.e., a student network.
Modify the architecture of ViT by appending another extra token called a Distillation Token.
DeiT, generates two outputs at:
Classification token: compared with a ground truth label.
Distillation token: compared with the logit output from the teacher network.
Distillation process in DeiT
2) Methodology:
Knowledge Distillation
Geoffery Hinton, in 2015: The basic idea is that we have two networks, a stronger, larger, pre-trained teacher network and a weak, small, randomly initialized student network.
==> The student network learns and tries to become more like the teacher, by minimizing a loss function in which the target label is the teacher’s probability distribution (aka. output of softmax layer of teacher).
a. Soft Distillation:
When the student tries to recreate the distribution predicted by the teacher.
Soft distilllation minimizes the Kullback-Leibler (KL) divergence between the softmax of the teacher and the softmax of the student model.
Soft Distillation function:
ψ =Softmax
Lce = cross-entropy loss
Zs = Logits of the student
Zt = Logits of Teacher
y= ground truth
yt= teacher’s predicted labels
λ = Weight of loss terms
The first term (red box) computes the cross-entropy loss of students’ prediction and ground truth labels
The second term(green box) is the distillation component of loss. This term computes KLL divergence loss between the softmax of students’ and teachers’ logits.
The contribution of each part is controlled by λ.
τ is softmax temperature:
Higher temperature(τ) means more melted / softer / distributed softmax distribution.
Lower temperature, gives more spiky distribution.
b. Hard Distillation:
When the student tries to recreate the labels predicted by the teacher.
In this, we reduce cross-entropy loss between labels (argmax of softmax) of teacher and softmax of student.
Hard Distillation function:
ψ =Softmax
Lce = cross-entropy loss
Zs = Logits of the student
Zt = Logits of Teacher
y= ground truth
yt= teacher’s predicted labels
λ = Weight of loss terms
In hard distillation, the first term(red box)is similar to soft distillation. The second term reflects(orange box) the change in soft and hard distillation, this time the distillation component is the cross-entropy loss between softmax of students and labels of a teacher. Both components contribute equally to the total loss.
* [NerdFact-🤓: KLL Divergence loss is similar to Cross-Entropy, it tries to compute difference between two probability distributions, and can be used interchangeably with cross-entropy, read here!]
Distillation Token:
The second modification to convert ViT to DeiT, is the addition of a distillation token to architecture.
In a vanilla ViT, just one learnable CLS token was added to the network, which captured the self-attention among all the input image patches, and was passed through simple MLP and softmax, to perform classification at the end. In DeiT architecture, authors added an additional class-type token called distillation token. This token is exactly like the CLS token, it’s initialized randomly, is learnable, and has a fixed last position, unlike the fixed first position of the CLS token.
DeiT Architecture:
So let’s pass an image through DeiT. We divide the image into 16x16(patch size from ViT) patches, pass it through an embedding layer, and get fixed size patch embeddings of size d. We add one randomly initialized d -sized vector at the start, called a class token, and one at the end called a distillation token. We add position embeddings to all tokens and pass them through a stack of encoders (3 in the above picture). Each encoder comprises a self-attention and FFN layer. The size of the output remains, the same as the input. We discard all tokens and just take the output of the CLS token and Distillation token, pass them through two separate linear layers to project them to the number of classes, and compute loss (training) or predict a class(inference).
Source: https://www.crcv.ucf.edu/wp-content/uploads/2018/11/deit_presentation.pdf
Loss Function:
Both CLS token and distillation token have their own objective functions to learn from. These objective functions are:
class token’s objective function = first term in the loss function (red box), True labels
distillation token’s objective function = the second term in loss function (orange box), Tacher’s predicted labels (Teacher here is RegNet Y-16GF)
[NerdFact-🤓: At inference time, authors averaged the softmax of distillation and CLS token, then passed through softmax +argmax to get better results]
Loss Function
The loss function is computed from both outputs, which allows the model to leverage the knowledge from the teacher while also learning from the ground truths.
CE is cross-entropy
psi is the softmax function
Z_s denotes student predictions
y denotes true labels
y_t denotes teacher predictions
They also incorporate some bag of tricks including data augmentation and regularization to further improve the performance.
With this technique, they reported that DeiT could be trained a single 8-GPU node in three days (53 hours for pre-training and 23 hours for optional fine-tuning) while the original ViT required 30 days to train with an 8-core TPUv3 machine.
