How Lightweight Can A Vision Transformer Be

TL;DR

This work tackles the challenge of building ultra-light vision transformers that perform well on small datasets and edge devices. It introduces mLiT, a MoE-based lightweight ViT with SwiGLU experts and shared projections, plus mmLiT with MAE pretraining for self-supervised learning and downstream fine-tuning. Key findings show that mmLiT-S can achieve competitive accuracy with only 0.67M parameters, with mmLiT-XS/XXS remaining viable at larger scales, and pretraining on CIFAR-100 enhances transfer to CIFAR-10, SVHN, and Flowers102. The results suggest that a combination of streamlined MoE, depth-wise scaling, and MAE pretraining can mitigate inductive-bias gaps and enable practical edge-efficient vision transformers.

Abstract

In this paper, we explore a strategy that uses Mixture-of-Experts (MoE) to streamline, rather than augment, vision transformers. Each expert in an MoE layer is a SwiGLU feedforward network, where V and W2 are shared across the layer. No complex attention or convolutional mechanisms are employed. Depth-wise scaling is applied to progressively reduce the size of the hidden layer and the number of experts is increased in stages. Grouped query attention is used. We studied the proposed approach with and without pre-training on small datasets and investigated whether transfer learning works at this scale. We found that the architecture is competitive even at a size of 0.67M parameters.

Figures (6)

• Figure 1: The structure of an MoE based transformer encoder layer. Each expert is a SwiGLU feedforward network. $b$ stands for batch size, $n$ for the number of embeddings and $m$ for embedding size. In this layer there are $t$ number of experts
• Figure 2: The working of MoE assuming the input is a vector. In this example, expert 1 and expert 3 receive the input as directed by the gating network
• Figure 3: The actual implementation of MoE. Unlike the case when the input is just a vector, when $b \cdot n$ embeddings go into the layer, almost all the expert networks will receive some embeddings sent by the dispatcher
• Figure 4: The overall architecture of mLiT. The linear transformation and positional embeddings before the first layer are not included in the illustration. $b$ stands for batch size, $n$ for the number of embeddings and $m$ for embedding size. The input and the output embedding size of each SwiGLU FFN expert are equal to the embedding size $m$. The model in this figure has 9 MoE based transformer encoder layers, with the number of experts increased by 1 at layer 3 and layer 6. $d^{\mathrm{first}}_h$ and $d^{\mathrm{last}}_h$ are 81 and 27, respectively. The hidden size at each layer is calculated using Equation \ref{['eq:hid_size']}.
• Figure 5: The architecture of mmLiT. $b$, $m$ and $p$ stand for batch size, embedding size and patch size, respectively. $n_E$ and $n_D$ are the number of embeddings/patches at the encoder and decoder respectively. $n_E + n_D = n$. $L_E$ is the total number of MoE transformer encoder layers at encoder; $L_D$ is the total number the layers at decoder.