The vanilla ViT is problematic. If you happen to check out the unique ViT paper [1], you’ll discover that though this Deep Learning mannequin proved to work extraordinarily effectively, it requires a whole bunch of hundreds of thousands of labeled coaching photos to attain this. Effectively, that’s lots.
This requirement of an infinite quantity of information is unquestionably an issue, and thus, we’d like an answer for that. Touvron et al. again in December 2020 introduced an thought of their analysis paper titled “Coaching data-efficient picture transformers & distillation via consideration” [2] to make coaching a ViT mannequin to be computationally less expensive. The authors got here up with an thought the place as an alternative of coaching the transformer-based mannequin from scratch, they exploited the data of the present mannequin via distillation. With this method, they managed to unravel the ViT’s data-hungry drawback whereas nonetheless sustaining excessive accuracy. What’s much more fascinating is that this paper got here out solely two months after the unique ViT!
On this article I’m going to debate the mannequin which the authors known as DeiT (Knowledge-efficient picture Transformer) in addition to implement the structure from scratch. Since DeiT is immediately derived from ViT, it’s extremely really helpful to have prior data about ViT earlier than studying this text. You’ll find my earlier article about it in reference [3] on the finish of this publish.
The Thought of DeiT
DeiT leverages the thought of data distillation. In case you’re not but aware of the time period, it’s primarily a way to switch the data of a mannequin (trainer) to a different one (scholar) in the course of the coaching part. On this case, DeiT acts as the coed whereas the trainer is RegNet, a CNN-based mannequin. Later within the inference part, we are going to utterly omit the RegNet trainer and let the DeiT scholar make predictions by itself.
The data distillation method permits the coed mannequin to be taught extra effectively, which is sensible because it not solely learns the patterns within the dataset from scratch but additionally advantages from the data of the trainer throughout coaching. Consider it like somebody studying a brand new topic. They may research purely from books, however it is going to be rather more environment friendly if in addition they had a mentor to offer steerage. On this analogy, the learner acts as the coed, the books are the dataset, whereas the mentor is the trainer. So, with this mechanism, the coed primarily derives data from each the dataset and the trainer concurrently. Because of this, coaching a scholar mannequin requires a lot much less quantity of information. To higher illustrate this, the unique ViT wanted 300 million photos for coaching (JFT-300M dataset), whereas DeiT depends solely on 1 million photos (ImageNet-1K dataset). That’s 300x smaller!
Technically talking, data distillation may be executed with out making any modifications to the coed or trainer fashions. Quite, the modifications are solely made to the loss perform and the coaching process. Nonetheless, authors discovered that they’ll obtain extra by barely modifying the community construction, which on the similar time additionally altering the distillation mechanism. Particularly, as an alternative of sticking with the unique ViT and apply an ordinary distillation course of on it, they modify the structure which they lastly seek advice from as DeiT. It is very important know that this modification additionally causes the data distillation mechanism to be completely different from the traditional one. To be actual, in ViT we solely have the so-called class token, however in DeiT, we are going to make the most of the class token itself and an extra one referred to as distillation token. Take a look at the Determine 1 under to see the place these two tokens are positioned within the community.
DeiT and ViT Variants
There are three DeiT variants proposed within the paper, particularly DeiT-Ti (Tiny), DeiT-S (Small) and DeiT-B (Base). Discover in Determine 2 that the biggest DeiT variant (DeiT-B) is equal to the smallest ViT variant (ViT-B) when it comes to the mannequin dimension. So, this implicitly signifies that DeiT was certainly designed to problem ViT by prioritizing effectivity.

Later within the coding half, I’m going to implement the DeiT-B structure. I’ll make the code as versatile as doable as a way to simply alter the parameters if you wish to implement the opposite variants as an alternative. Taking a better have a look at the DeiT-B row within the above desk, we’re going to configure the mannequin such that it maps every picture patch to a single-dimensional tensor of dimension 768. The weather on this tensor will then be grouped into 12 heads inside the eye layer. By doing so, each single of those consideration heads shall be accountable to course of 64 options. Keep in mind that the eye layer we’re speaking about is basically a part of a Transformer encoder layer. Within the case of DeiT-B, this layer is repeated 12 instances earlier than the tensor is ultimately forwarded to the output layer. If we implement it appropriately in line with these configurations, the mannequin ought to include 86 million trainable parameters.
Experimental Outcomes
There are many experiments reported within the DeiT paper. Beneath is one in every of them that grabbed my consideration probably the most.

The above determine was obtained by coaching a number of fashions on ImageNet-1K dataset, together with EfficientNet, ViT, and the DeiT itself. In reality, there are two DeiT variations displayed within the determine: DeiT and DeiT⚗ — sure with that unusual image for the latter (referred to as “alembic”), which mainly refers back to the DeiT mannequin educated utilizing their proposed distillation mechanism.
It’s seen within the determine that the accuracy of ViT is already far behind DeiT with typical distillation whereas nonetheless having the same processing velocity. The accuracy improved even additional when the novel distillation mechanism was utilized and the mannequin was fine-tuned utilizing the identical photos upscaled to 384×384 — therefore the identify DeiT-B⚗↑384. In concept, ViT ought to have carried out higher than its present outcome, but on this experiment it couldn’t unleash its full potential because it wasn’t allowed to be educated on the large JFT-300M dataset. And that’s only one outcome that proves the prevalence of DeiT over ViT in a data-limited scenario.
I believe that was in all probability all of the issues you should perceive to implement the DeiT structure from scratch. Don’t fear for those who haven’t absolutely grasped all the thought of this mannequin but since we are going to get into the small print in a minute.
DeiT Implementation
As I discussed earlier, the mannequin we’re about to implement is the DeiT-B variant. However since I additionally wish to present you the novel data distillation mechanism, I’ll particularly concentrate on the one known as DeiT-B⚗↑384. Now let’s begin by importing the required modules.
# Codeblock 1
import torch
import torch.nn as nn
from timm.fashions.layers import trunc_normal_
from torchinfo import abstract
Because the modules have been imported, what we have to do subsequent is to initialize some configurable parameters within the Codeblock 2 under, that are all adjusted in line with the DeiT-B specs. On the line #(1)
, the IMAGE_SIZE
variable is about to 384 since we’re about to simulate the DeiT model that accepts the upscaled photos. Regardless of this increased decision enter, we nonetheless hold the patch dimension the identical as when working with 224×224 photos, i.e., 16×16, as written at line #(2)
. Subsequent, we set EMBED_DIM
to 768 (#(3)
), whereas the NUM_HEADS
and NUM_LAYERS
variables are each set to 12 (#(4–5)
). Authors determined to make use of the identical FFN construction because the one utilized in ViT, wherein the scale of its hidden layer is 4 instances bigger than the embedding dimension (#(6)
). The variety of patches itself may be calculated utilizing a easy components proven at line #(7)
. On this case, since our picture dimension is 384 and the patch dimension is 16, the worth of NUM_PATCHES
goes to be 576. Lastly, right here I set NUM_CLASSES
to 1000, simulating a classification job on ImageNet-1K dataset (#(8)
).
# Codeblock 2
BATCH_SIZE = 1
IMAGE_SIZE = 384 #(1)
IN_CHANNELS = 3
PATCH_SIZE = 16 #(2)
EMBED_DIM = 768 #(3)
NUM_HEADS = 12 #(4)
NUM_LAYERS = 12 #(5)
FFN_SIZE = EMBED_DIM * 4 #(6)
NUM_PATCHES = (IMAGE_SIZE//PATCH_SIZE) ** 2 #(7)
NUM_CLASSES = 1000 #(8)
Treating an Picture as a Sequence of Patches
In the case of processing photos utilizing transformers, what we have to do is to deal with them as a sequence of patches. Such a patching mechanism is carried out within the Patcher
class under.
# Codeblock 3
class Patcher(nn.Module):
def __init__(self):
tremendous().__init__()
self.conv = nn.Conv2d(in_channels=IN_CHANNELS, #(1)
out_channels=EMBED_DIM,
kernel_size=PATCH_SIZE, #(2)
stride=PATCH_SIZE) #(3)
self.flatten = nn.Flatten(start_dim=2) #(4)
def ahead(self, x):
print(f'originalt: {x.dimension()}')
x = self.conv(x) #(5)
print(f'after convt: {x.dimension()}')
x = self.flatten(x) #(6)
print(f'after flattent: {x.dimension()}')
x = x.permute(0, 2, 1) #(7)
print(f'after permutet: {x.dimension()}')
return x
You possibly can see in Codeblock 3 that we use an nn.Conv2d
layer to take action (#(1)
). Needless to say the operation executed by this layer will not be meant to truly carry out convolution like in CNN-based fashions. As an alternative, we use it as a trick to extract the knowledge of every patch in a non-overlapping method, which is the rationale that we set each kernel_size
(#(2)
) and stride
(#(3)
) to PATCH_SIZE
(16). The operation executed by this convolution layer entails the patching mechanism solely — we haven’t truly put these patches into sequence simply but. So as to take action, we are able to merely make the most of an nn.Flatten
layer which I initialize at line #(4)
within the above codeblock. What we have to do contained in the ahead()
technique is to move the enter tensor via the conv
(#(5)
) and flatten
(#(6)
) layers. Additionally it is essential to carry out the permute operation afterwards as a result of we would like the patch sequence to be positioned alongside axis 1 and the embedding dimension alongside axis 2 (#(7)
).
Now let’s take a look at the Patcher()
class above utilizing the next codeblock. Right here I take a look at it with a dummy tensor which the dimension is about to 1×3×384×384, simulating a single RGB picture of dimension 384×384.
# Codeblock 4
patcher = Patcher()
x = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
x = patcher(x)
And under is what the output seems like. Right here I print out the tensor dimension after every step as a way to clearly see the circulate contained in the community.
# Codeblock 4 Output
authentic : torch.Measurement([1, 3, 384, 384])
after conv : torch.Measurement([1, 768, 24, 24]) #(1)
after flatten : torch.Measurement([1, 768, 576]) #(2)
after permute : torch.Measurement([1, 576, 768]) #(3)
Discover at line #(1)
that the spatial dimension of the tensor modified from 384×384 to 24×24. This means that our convolution layer efficiently executed the patching course of. By doing so, each single pixel within the 24×24 picture now represents every 16×16 patch of the enter picture. Moreover, discover in the identical line that the variety of channels elevated from 3 to EMBED_DIM
(768). In a while, we are going to understand this because the variety of options that shops the knowledge of a single patch. Subsequent, we are able to see at line #(2)
that our flatten
layer efficiently flattened the 24×24 tensor right into a single-dimensional tensor of size 576, which signifies that we already acquired our picture represented as a sequence of patch tokens. The permute operation I discussed earlier was primarily executed as a result of within the case of time-series information PyTorch treats the axis 1 of a tensor as a sequence (#(3)
).
Transformer Encoder
Now let’s put our Patcher
class apart for some time since on this part we’re going to implement the transformer encoder layer. This layer is immediately derived from the unique ViT paper which the structure may be seen within the Determine 4 under. Check out Codeblock 5 to see how I implement it.

# Codeblock 5
class Encoder(nn.Module):
def __init__(self):
tremendous().__init__()
self.norm_0 = nn.LayerNorm(EMBED_DIM) #(1)
self.multihead_attention = nn.MultiheadAttention(EMBED_DIM, #(2)
num_heads=NUM_HEADS,
batch_first=True)
self.norm_1 = nn.LayerNorm(EMBED_DIM) #(3)
self.ffn = nn.Sequential( #(4)
nn.Linear(in_features=EMBED_DIM, out_features=FFN_SIZE),
nn.GELU(),
nn.Linear(in_features=FFN_SIZE, out_features=EMBED_DIM),
)
def ahead(self, x):
residual = x
print(f'residual dimt: {residual.dimension()}')
x = self.norm_0(x)
print(f'after normt: {x.dimension()}')
x = self.multihead_attention(x, x, x)[0]
print(f'after attentiont: {x.dimension()}')
x = x + residual
print(f'after additiont: {x.dimension()}')
residual = x
print(f'residual dimt: {residual.dimension()}')
x = self.norm_1(x)
print(f'after normt: {x.dimension()}')
x = self.ffn(x)
print(f'after ffnt: {x.dimension()}')
x = x + residual
print(f'after additiont: {x.dimension()}')
return x
Based on the above determine, there are 4 layers have to be initialized within the __init__()
technique, particularly a multihead consideration layer (#(2)
), an MLP layer — which is equal to FFN in Determine 1 (#(4)
), and two layer normalization layers (#(1,3)
). I’m not going to get deeper into the above code since it’s precisely the identical as what I defined in my earlier article about ViT [4]. So, I do suggest you verify that article to raised perceive how the Encoder
class works. And moreover, for those who want an in-depth clarification particularly concerning the consideration mechanism, you may also learn my earlier transformer article [5] the place I carried out all the transformer structure from scratch.
We are able to now simply go forward to the testing code to see how the tensor flows via the community. Within the following codeblock, I assume that the enter tensor x
is a picture that has already been processed by the Patcher
block we created earlier, which is the rationale why I set it to have the scale of 1×576×768.
# Codeblock 6
encoder = Encoder()
x = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
x = encoder(x)
# Codeblock 6 Output
residual dim : torch.Measurement([1, 576, 768])
after norm : torch.Measurement([1, 576, 768])
after consideration : torch.Measurement([1, 576, 768])
after addition : torch.Measurement([1, 576, 768])
residual dim : torch.Measurement([1, 576, 768])
after norm : torch.Measurement([1, 576, 768])
after ffn : torch.Measurement([1, 576, 768])
after addition : torch.Measurement([1, 576, 768])
Based on the above outcome, we are able to see that the ultimate output tensor dimension is strictly the identical as that of the enter. This property permits us to stack a number of encoder blocks with out disrupting all the community construction. Moreover, though the form of the tensor seems to be fixed alongside its approach to the final layer, there are literally plenty of dimensionality modifications taking place particularly inside the eye and the FFN layers. Nonetheless, these modifications are usually not printed because the processes are executed internally by nn.MultiheadAttention
and nn.Sequential
, respectively.
The Total DeiT Structure
All of the codes I defined within the earlier sections are literally similar to these used for setting up the ViT structure. On this part, you’ll lastly discover those that clearly differentiate DeiT from ViT. Let’s now concentrate on the layers we have to initialize within the __init__()
technique of the DeiT
class under.
# Codeblock 7a
class DeiT(nn.Module):
def __init__(self):
tremendous().__init__()
self.patcher = Patcher() #(1)
self.class_token = nn.Parameter(torch.zeros(BATCH_SIZE, 1, EMBED_DIM)) #(2)
self.dist_token = nn.Parameter(torch.zeros(BATCH_SIZE, 1, EMBED_DIM)) #(3)
trunc_normal_(self.class_token, std=.02) #(4)
trunc_normal_(self.dist_token, std=.02) #(5)
self.pos_embedding = nn.Parameter(torch.zeros(BATCH_SIZE, NUM_PATCHES+2, EMBED_DIM)) #(6)
trunc_normal_(self.pos_embedding, std=.02) #(7)
self.encoders = nn.ModuleList([Encoder() for _ in range(NUM_LAYERS)]) #(8)
self.norm_out = nn.LayerNorm(EMBED_DIM) #(9)
self.class_head = nn.Linear(in_features=EMBED_DIM, out_features=NUM_CLASSES) #(10)
self.dist_head = nn.Linear(in_features=EMBED_DIM, out_features=NUM_CLASSES) #(11)
The primary part I initialized right here is Patcher
we created earlier (#(1)
). Subsequent, as an alternative of solely utilizing class token, DeiT makes use of one other one named distillation token. These two tokens, which within the above code are known as class_token
(#(2)
) and dist_token
(#(3)
), will later be appended to the patch token sequence. We set these two extra tokens to be trainable, permitting them to work together with and be taught from the patch tokens later in the course of the processing within the consideration layer. Discover that I initialized these two trainable tensors utilizing trunc_normal_()
with an ordinary deviation of 0.02 (#(4–5)
). In case you’re not but aware of the perform, it primarily generates a truncated regular distribution, which ensures that no worth lies past two customary deviations from the imply, avoiding the presence of maximum values for weight initialization. This method is definitely higher than immediately utilizing torch.randn()
since this perform doesn’t have such a worth truncation mechanism.
Afterwards, we create a learnable positional embedding tensor utilizing the identical method which I do at traces #(6)
and #(7)
. It is very important remember that this tensor will then be element-wise summed with the sequence of patch tokens that has been appended with the category and distillation tokens. Resulting from this motive, we have to set the size of axis 1 of this embedding tensor to NUM_PATCHES+2
. In the meantime, the transformer encoder layer is initialized inside nn.ModuleList
which permits us to repeat the layer NUM_LAYERS
(12) instances (#(8)
). The output produced by the final encoder layer within the stack shall be processed with a layer norm (#(9)
) earlier than ultimately being forwarded to the classification (#(10)
) and distillation heads (#(11)
).
Now let’s transfer on to the ahead()
technique which you’ll be able to see within the Codeblock 7b under.
# Codeblock 7b
def ahead(self, x):
print(f'originaltt: {x.dimension()}')
x = self.patcher(x) #(1)
print(f'after patchertt: {x.dimension()}')
x = torch.cat([self.class_token, self.dist_token, x], dim=1) #(2)
print(f'after concattt: {x.dimension()}')
x = x + self.pos_embedding #(3)
print(f'after pos embedtt: {x.dimension()}')
for i, encoder in enumerate(self.encoders):
x = encoder(x) #(4)
print(f"after encoder #{i}t: {x.dimension()}")
x = self.norm_out(x) #(5)
print(f'after normtt: {x.dimension()}')
class_out = x[:, 0] #(6)
print(f'class_outtt: {class_out.dimension()}')
dist_out = x[:, 1] #(7)
print(f'dist_outtt: {dist_out.dimension()}')
class_out = self.class_head(class_out) #(8)
print(f'after class_headt: {class_out.dimension()}')
dist_out = self.dist_head(dist_out) #(9)
print(f'after dist_headtt: {class_out.dimension()}')
return class_out, dist_out
After taking uncooked picture because the enter, this ahead()
technique will course of the picture utilizing the patcher
layer (#(1)
). As now we have beforehand mentioned, this layer is accountable to transform the picture right into a sequence of patches. Subsequently, we are going to concatenate the category and distillation tokens to it utilizing torch.cat()
(#(2)
). It may be value noting that though the illustration in Determine 1 locations the category token to start with of the sequence and the distillation token on the finish, however the code within the official GitHub repository [6] says that the distillation token is positioned proper after the category token. Thus, I made a decision to observe this method in our implementation. Determine 5 under illustrates what the ensuing tensor seems like.

Nonetheless with Codeblock 7b, what we have to do subsequent is to inject the positional embedding tensor to the token sequence which the method is completed at line (#(3)
). We then move the tensor via the stack of encoders utilizing a easy loop (#(4)
) and normalize the output produced by the final encoder layer (#(5)
). At traces #(6)
and #(7)
we extract the knowledge from the category and distillation tokens we appended earlier utilizing an ordinary array slicing technique. These two tokens ought to now include significant data for classification job since they already realized the context of the picture via the self-attention layers. The ensuing class_out
and dist_out
tensors are then forwarded to 2 similar output layers and can endure processing independently (#(8–9)
). Since this mannequin is meant for classification, these two output layers will produce tensors containing logits, wherein each single factor represents the uncooked prediction rating of a category.
We are able to see the circulate of the DeiT mannequin with the next testing code, the place we initially begin with the uncooked enter picture (#(1)
), turning it into sequence of patches (#(2)
), concatenating class and distillation tokens (#(3)
), and so forth till ultimately getting the output from each classification and distillation heads (#(4–5)
).
# Codeblock 8
deit = DeiT()
x = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
class_out, dist_out = deit(x)
# Codeblock 8 Output
authentic : torch.Measurement([1, 3, 384, 384]) #(1)
after patcher : torch.Measurement([1, 576, 768]) #(2)
after concat : torch.Measurement([1, 578, 768]) #(3)
after pos embed : torch.Measurement([1, 578, 768])
after encoder #0 : torch.Measurement([1, 578, 768])
after encoder #1 : torch.Measurement([1, 578, 768])
after encoder #2 : torch.Measurement([1, 578, 768])
after encoder #3 : torch.Measurement([1, 578, 768])
after encoder #4 : torch.Measurement([1, 578, 768])
after encoder #5 : torch.Measurement([1, 578, 768])
after encoder #6 : torch.Measurement([1, 578, 768])
after encoder #7 : torch.Measurement([1, 578, 768])
after encoder #8 : torch.Measurement([1, 578, 768])
after encoder #9 : torch.Measurement([1, 578, 768])
after encoder #10 : torch.Measurement([1, 578, 768])
after encoder #11 : torch.Measurement([1, 578, 768])
after norm : torch.Measurement([1, 578, 768])
class_out : torch.Measurement([1, 768])
dist_out : torch.Measurement([1, 768])
after class_head : torch.Measurement([1, 1000]) #(4)
after dist_head : torch.Measurement([1, 1000]) #(5)
You may as well run the next code if you wish to see much more particulars of the structure. It’s seen within the ensuing output that this community comprises 87 million variety of parameters, which is barely increased than reported within the paper (86 million). I do acknowledge that the code I wrote above is certainly a lot less complicated than the one within the documentation, so I would in all probability miss one thing that results in such a distinction within the variety of params — please let me know for those who spot any errors in my code!
# Codeblock 9
abstract(deit, input_size=(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE))
# Codeblock 9 Output
==========================================================================================
Layer (sort:depth-idx) Output Form Param #
==========================================================================================
DeiT [1, 1000] 445,440
├─Patcher: 1-1 [1, 576, 768] --
│ └─Conv2d: 2-1 [1, 768, 24, 24] 590,592
│ └─Flatten: 2-2 [1, 768, 576] --
├─ModuleList: 1-2 -- --
│ └─Encoder: 2-3 [1, 578, 768] --
│ │ └─LayerNorm: 3-1 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-2 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-3 [1, 578, 768] 1,536
│ │ └─Sequential: 3-4 [1, 578, 768] 4,722,432
│ └─Encoder: 2-4 [1, 578, 768] --
│ │ └─LayerNorm: 3-5 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-6 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-7 [1, 578, 768] 1,536
│ │ └─Sequential: 3-8 [1, 578, 768] 4,722,432
│ └─Encoder: 2-5 [1, 578, 768] --
│ │ └─LayerNorm: 3-9 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-10 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-11 [1, 578, 768] 1,536
│ │ └─Sequential: 3-12 [1, 578, 768] 4,722,432
│ └─Encoder: 2-6 [1, 578, 768] --
│ │ └─LayerNorm: 3-13 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-14 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-15 [1, 578, 768] 1,536
│ │ └─Sequential: 3-16 [1, 578, 768] 4,722,432
│ └─Encoder: 2-7 [1, 578, 768] --
│ │ └─LayerNorm: 3-17 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-18 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-19 [1, 578, 768] 1,536
│ │ └─Sequential: 3-20 [1, 578, 768] 4,722,432
│ └─Encoder: 2-8 [1, 578, 768] --
│ │ └─LayerNorm: 3-21 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-22 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-23 [1, 578, 768] 1,536
│ │ └─Sequential: 3-24 [1, 578, 768] 4,722,432
│ └─Encoder: 2-9 [1, 578, 768] --
│ │ └─LayerNorm: 3-25 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-26 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-27 [1, 578, 768] 1,536
│ │ └─Sequential: 3-28 [1, 578, 768] 4,722,432
│ └─Encoder: 2-10 [1, 578, 768] --
│ │ └─LayerNorm: 3-29 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-30 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-31 [1, 578, 768] 1,536
│ │ └─Sequential: 3-32 [1, 578, 768] 4,722,432
│ └─Encoder: 2-11 [1, 578, 768] --
│ │ └─LayerNorm: 3-33 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-34 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-35 [1, 578, 768] 1,536
│ │ └─Sequential: 3-36 [1, 578, 768] 4,722,432
│ └─Encoder: 2-12 [1, 578, 768] --
│ │ └─LayerNorm: 3-37 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-38 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-39 [1, 578, 768] 1,536
│ │ └─Sequential: 3-40 [1, 578, 768] 4,722,432
│ └─Encoder: 2-13 [1, 578, 768] --
│ │ └─LayerNorm: 3-41 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-42 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-43 [1, 578, 768] 1,536
│ │ └─Sequential: 3-44 [1, 578, 768] 4,722,432
│ └─Encoder: 2-14 [1, 578, 768] --
│ │ └─LayerNorm: 3-45 [1, 578, 768] 1,536
│ │ └─MultiheadAttention: 3-46 [1, 578, 768] 2,362,368
│ │ └─LayerNorm: 3-47 [1, 578, 768] 1,536
│ │ └─Sequential: 3-48 [1, 578, 768] 4,722,432
├─LayerNorm: 1-3 [1, 578, 768] 1,536
├─Linear: 1-4 [1, 1000] 769,000
├─Linear: 1-5 [1, 1000] 769,000
==========================================================================================
Whole params: 87,630,032
Trainable params: 87,630,032
Non-trainable params: 0
Whole mult-adds (Items.MEGABYTES): 398.43
==========================================================================================
Enter dimension (MB): 1.77
Ahead/backward move dimension (MB): 305.41
Params dimension (MB): 235.34
Estimated Whole Measurement (MB): 542.52
==========================================================================================
How Classification and Distillation Heads Work
I want to speak a bit bit concerning the tensors produced by the 2 output heads. Throughout the coaching part, the output from the classification head is in contrast with the unique floor fact (one-hot label) which the classification efficiency is evaluated utilizing cross entropy loss. In the meantime, the output from the distillation head is in contrast with the output produced by the trainer mannequin, i.e., RegNet. We all the time understand the output of the trainer as a fact no matter whether or not its prediction is right. And that’s primarily how data is distilled from RegNet to DeiT.
There are literally two strategies doable for use to carry out data distillation: gentle distillation and laborious distillation. The previous is a way the place we use the logits produced by the trainer mannequin as is (reasonably than the argmaxed logits) for the label. This type of extra floor fact is known as gentle label. If we determined to make use of this system, we should always use the so-called Kullback-Leibler (KL) loss, which is appropriate for evaluating two logits: one from the distillation head and one other one from the trainer output. However, laborious distillation is a way the place the prediction made by the trainer is argmaxed previous to being in contrast with the output from the distillation head. On this case, the trainer output is known as laborious label, which is analogous to a typical one-hot-encoded label. Due to this motive, if we have been to make use of laborious label as an alternative, we are able to merely use the usual cross-entropy loss for this head. Though the authors discovered that arduous distillation carried out higher than gentle distillation, I nonetheless suppose that it’s value experimenting with the 2 approaches for those who plan to make use of DeiT on your upcoming venture to see if this notion additionally applies to your case.
Throughout the inference part, we are going to not use the trainer mannequin. Consider it like the coed has graduated and is able to work by itself. Regardless of the absence of the trainer, the output from the distillation head continues to be utilized. Based on their GitHub documentation [6], the logits produced by each the classification and distillation heads are mixed utilizing an ordinary averaging mechanism earlier than being argmaxed to acquire the ultimate prediction.
Ending
I believe that’s every thing about the primary thought and implementation of DeiT. It is very important word that there are nonetheless plenty of issues I haven’t lined on this article. So, I do suggest you learn the paper [2] if you wish to get even deeper into the small print of this deep studying mannequin.
Thanks for studying, I hope you be taught one thing new immediately!
By the best way you’ll be able to entry the code used on this article within the hyperlink at reference quantity [7].
References
[1] Alexey Dosovitskiy et al. An Picture is Price 16×16 Phrases: Transformers for Picture Recognition at Scale. Arxiv. https://arxiv.org/abs/2010.11929 [Accessed February 17, 2025].
[2] Hugo Touvron et al. Coaching Knowledge-Environment friendly Picture Transformers & Distillation Via Consideration. Arxiv. https://arxiv.org/abs/2012.12877 [Accessed February 17, 2025].
[3] Picture initially created by writer.
[4] Muhammad Ardi. Paper Walkthrough: Vision Transformer (ViT). In the direction of Knowledge Science. https://towardsdatascience.com/paper-walkthrough-vision-transformer-vit-c5dcf76f1a7a/ [Accessed February 17, 2025].
[5] Muhammad Ardi. Paper Walkthrough: Consideration Is All You Want. In the direction of Knowledge Science. https://towardsdatascience.com/paper-walkthrough-attention-is-all-you-need-80399cdc59e1/ [Accessed February 17, 2025].
[6] facebookresearch. GitHub. https://github.com/facebookresearch/deit/blob/main/models.py [Accessed February 17, 2025].
[7] MuhammadArdiPutra. Imaginative and prescient Transformer on a Finances. GitHub. https://github.com/MuhammadArdiPutra/medium_articles/blob/main/Vision%20Transformer%20on%20a%20Budget.ipynb [Accessed February 17, 2025].