How does the embeddings work in vision transformer from paper?

Question

I get the part from the paper where the image is split into P say 16x16 (smaller images) patches and then you have to Flatten the 3-D (16,16,3) patch to pass it into a Linear layer to get what they call "Liner Projection". After passing from the Linear layer, the patches will be vectors but with some "meaning" to them.

Can someone please explain how the two types of embeddings are working?

I visited this implementation on github, looked at the code too and looked like a maze to me.

If someone could just explain how these embeddings are working in laymen's terms, I'll look at the code again and understand.

Answer 1

In Machine Learning "embedding" means taking some set of raw inputs (like natural language tokens in NLP or image patches in your example) and converting them to vectors somehow. The embeddings usually have some interesting dot-product structure between vectors (like in word2vec for example). The Transformer machinery then uses this embedding in the dot-product attention pipeline. The dimension of the embedding $D$ should stay constant throughout the transformer blocks due to ResNet skip connections.

The simplest idea in the case of the image patches would be to just take all the channels of all the pixels and treat them as a single vector. For example, if you've got (16, 16, 3) patches then you'll have 768-dimensional "embeddings". The problem with such a naive "embedding" is that dot-products between them won't make much sense. So we also multiply these vectors by a trainable matrix $W$ and add a bias vector.

For example, if you've got (16, 16, 3) patches and the transformer downstream uses $D=128$ dimensional embeddings, then you first flatten the patches into 16 * 16 * 3 = 768-dimensional vectors and then multiply by a 768$\times$128 matrix and add a 128-dimensional bias vector.

Looking at the code, it seems that the authors kept improving on that idea by adding one or several early convolutions with nonlinearities (the conv_stem branches in the code). The simplest execution branch before these embellishments seems to be here. And the matrix multiplication I've been talking about is here