FeaTrim-ViT: Vision Transformer Trimming with One Shot Neural Architecture Search in Continuous Optimisation Space and Efficient Feature Selection

This paper introduces FeaTrim-ViT, a novel framework for trimming Vision Transformer (ViT) models, which are known for their exceptional performance but also for their large number of parameters, making them challenging to deploy on memory-limited devices. Our approach innovatively combines structured pruning techniques with nature-inspired algorithms to efficiently identify and eliminate redundant parameters within ViT models. Existing studies fail to maintain accuracy at higher compression rates during efficient finetuning. The framework is designed to comprehensively search for an optimal submodel, targeting multiple dimensions including input tokens, Multi-Head Self-Attention (MHSA), and Multilayer Perceptron (MLP) modules. A learnable sparsity constraint with preset factors that reflect global importance across these dimensions. The search process is based on a single-shot training scheme making it exceptionally efficient. The resulting substructures boast diverse dimensionalities, tailored to specific module requirements. A budget threshold, based on the desired accuracy-FLOPs trade-off, guides the retraining process to produce the optimised model which finally goes through efficient feature selection using Grey Wolf Optimiser (GWO). The extensive experiments demonstrate that the framework can reduce up to 35% of parameters and FLOPs in the ViT model, while maintaining accuracy obtained by the original model of 80%. This architecture offers a step forward in the practical deployment of high-capability ViT models in resource-constrained environments.