[Paper Review] LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference

Ben Graham, Alaaeldin El-Nouby+5 more|arXiv (Cornell University)|Apr 2, 2021

Advanced Neural Network Applications66 references91 citations

TL;DR

LeViT proposes a pyramid-structured Vision Transformer that remains competitive in speed with narrower DeiT blocks while achieving faster inference due to design choices such as wider blocks and reduced MLP expansion. Supplementary material provides detailed block timings and attention bias visualizations.

ABSTRACT

We design a family of image classification architectures that optimize the trade-off between accuracy and efficiency in a high-speed regime. Our work exploits recent findings in attention-based architectures, which are competitive on highly parallel processing hardware. We revisit principles from the extensive literature on convolutional neural networks to apply them to transformers, in particular activation maps with decreasing resolutions. We also introduce the attention bias, a new way to integrate positional information in vision transformers. As a result, we propose LeVIT: a hybrid neural network for fast inference image classification. We consider different measures of efficiency on different hardware platforms, so as to best reflect a wide range of application scenarios. Our extensive experiments empirically validate our technical choices and show they are suitable to most architectures. Overall, LeViT significantly outperforms existing convnets and vision transformers with respect to the speed/accuracy tradeoff. For example, at 80% ImageNet top-1 accuracy, LeViT is 5 times faster than EfficientNet on CPU. We release the code at https://github.com/facebookresearch/LeViT

Motivation & Objective

• Motivate faster inference for Vision Transformers by rethinking block design and pyramid structure.
• Characterize runtime of LeViT blocks versus DeiT blocks on comparable resolutions and compute budgets.
• Investigate the impact of pyramid structure and block width on overall efficiency.
• Provide ablations and visualizations to explain attention behavior across LeViT blocks.

Proposed method

• Compare DeiT-tiny and LeViT-256 block designs at 14x14 resolution and compute side-by-side runtimes.
• Analyze the contributions of LayerNorm, Q/K, V, QK^T, AV, attention projection, and MLP to total runtime.
• Show ablations for removing pyramid structure and for widening/block adjustments to understand efficiency gains.
• Visualize attention bias maps to interpret how different heads attend to relative pixel positions.

Experimental results

Research questions

• RQ1 Does LeViT achieve comparable or faster inference than DeiT with a pyramid/convnet-inspired design?
• RQ2 How do pyramid structure and block width affect runtime components and overall efficiency?
• RQ3 What is the effect of making MLP expansion smaller and attention computations on speed?
• RQ4 What do attention bias visualizations reveal about head specialization and information flow across LeViT blocks?

Key findings

• LeViT-256 has a total runtime close to DeiT-tiny, with LeViT’s total runtime around 2365 μs versus 2474 μs for DeiT-tiny on the same benchmark setup.
• LeViT spends less time on QK^T and more time on the subsequent AV product, despite having a wider block width (C=256 vs C=192).
• LeViT reduces MLP runtime by halving the expansion factor from four to two, offsetting some width-related costs.
• Attention bias visualizations show some heads focusing on nearby pixels, with others displaying uniform or directional patterns across stages, illustrating diverse attention strategies.
• Ablations demonstrate the impact of removing pyramid structure or widening blocks on overall performance and FLOP counts.