Why is there a Embedding layer in each stage of the model framework

[WuliSmart]

No description provided.

[Lupin1998]

Hi, @WuliSmart! Thanks for your interesting question, and sorry for the late reply. We apply the embedding layer (Conv2d 2x2 with a norm layer) at the start of each stage to stabilize the training and introduce more nonlinearity for the lightweight and small-size variants. A similar design can be found in the Metaformer baseline (https://arxiv.org/abs/2210.13452).