Why class embedding token is added to the Visual Transformer?

Question

In the famous work on the Visual Transformers, the image is split into patches of a certain size (say 16x16), and these patches are treated as tokens in the NLP tasks. In order to perform classification, a CLS token is added at the beginning of the resulting sequence: $$ [\textbf{x}_{class}, \textbf{x}_{p}^{1}, \ldots, \textbf{x}_{p}^{N}] ,$$ where $ \textbf{x}_{p}^{i}$ are image patches. There multiple layers in the architecture and the state of the CLS token on the output layer is used for classification.

I think this architectural solution is done in the spirit of NLP problems (BERT in particular). However, for me, it would be more natural not to create a special token, but perform 1d Global Pooling in the end, and attach an nn.Linear(embedding_dim, num_classes) as more conventional CV approach.

Why it is not done in this way? Or is there some intuition or evidence that this would perform worse than the approach used in the paper?

Answer 1

In their official repository, the author confirmed in the issue that the cls is not really important for the ViT, but they wanted to keep the ViT to be exactly the same with the NLP-Transformer, so they added it.

Issue link: https://github.com/google-research/vision_transformer/issues/61#issuecomment-802233921

Answer 2

However, for me, it would be more natural not to create a special token, but perform 1d Global Pooling in the end, and attach an nn.Linear(embedding_dim, num_classes) as more conventional CV approach.

Surprisingly, this approach was experimented in this new paper Better plain ViT baselines for ImageNet-1k and implemented as default in the torch-vit repository.

Along with other trivial improvements, the authors actually outperform other methods such as DieT or ViT with strong data augmentation. The modification you described provides 1.8% accuracy gain in the ablation study.

The paper is essentially 1.5 pages so I definitely recommend having a look!

link: https://arxiv.org/abs/2205.01580