PaperlayerView exact revision on arXiv

Source-linked AI summary

FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, Anurag Ranjan

arXiv:2303.14189v2cs.CV

TL;DR

FastViT addresses the efficiency gap between accurate vision transformers and computationally cheaper convolutional models. It combines RepMixer structural reparameterization with factorized convolutions, train-time overparameterization, and large kernels, achieving strong latency–accuracy results across platforms and tasks. The model also demonstrates robustness to corruption and out-of-distribution samples, although train-time overparameterization modestly increases training time.

  • Problem

    Vision transformers and self-attention can be computationally expensive, motivating architectures that preserve accuracy while improving latency and efficiency.

  • Method

    FastViT uses RepMixer to remove skip connections through structural reparameterization, plus factorized convolutions, train-time overparameterization, and large convolutional kernels.

  • Results

    FastViT achieves improved latency–accuracy trade-offs across image classification, detection, segmentation, and 3D hand mesh estimation on mobile and desktop GPU platforms.

  • Takeaways & Limitations

    FastViT provides a general-purpose hybrid backbone that is efficient across compute fabrics and tasks while remaining robust to corruption and out-of-distribution samples.

  • Takeaways & Limitations

    Train-time overparameterization increases training time by 6.7% for FastViT-SA12 and 4.4% for FastViT-SA36 under the reported settings.

Abstract

from arXiv · show

The recent amalgamation of transformer and convolutional designs has led to steady improvements in accuracy and efficiency of the models. In this work, we introduce FastViT, a hybrid vision transformer architecture that obtains the state-of-the-art latency-accuracy trade-off. To this end, we introduce a novel token mixing operator, RepMixer, a building block of FastViT, that uses structural reparameterization to lower the memory access cost by removing skip-connections in the network. We further apply train-time overparametrization and large kernel convolutions to boost accuracy and empirically show that these choices have minimal effect on latency. We show that - our model is 3.5x faster than CMT, a recent state-of-the-art hybrid transformer architecture, 4.9x faster than EfficientNet, and 1.9x faster than ConvNeXt on a mobile device for the same accuracy on the ImageNet dataset. At similar latency, our model obtains 4.2% better Top-1 accuracy on ImageNet than MobileOne. Our model consistently outperforms competing architectures across several tasks -- image classification, detection, segmentation and 3D mesh regression with significant improvement in latency on both a mobile device and a desktop GPU. Furthermore, our model is highly robust to out-of-distribution samples and corruptions, improving over competing robust models. Code and models are available at https://github.com/apple/ml-fastvit.

1. Introduction

FastViT targets state-of-the-art latency–accuracy trade-offs by combining hybrid vision-transformer design with structural reparameterization, train-time overparameterization, and large convolutional kernels. It achieves strong latency improvements across compute platforms and vision tasks while maintaining accuracy and robustness.

  • Motivation: FastViT combines convolutional and transformer designs to pursue state-of-the-art latency–accuracy trade-offs across vision tasks.The architecture is evaluated on image classification, detection, segmentation, and 3D hand mesh estimation.
  • Architecture: RepMixer removes token-mixer skip connections through structural reparameterization, reducing memory access cost and eliminating inference-time branches.It uses depthwise convolutions for spatial information mixing and can be reparameterized at inference.
  • Architecture: Train-time overparameterization restores capacity after replacing dense convolutions with factorized depthwise-plus-pointwise convolutions.The added branches exist only during training and are reparameterized at inference.
  • Architecture: Large-kernel convolutions replace self-attention in early stages and are incorporated into FFN and patch-embedding layers with minimal latency impact.These changes are reported to improve performance while preserving overall latency.
  • Evaluation: At 83.9% ImageNet Top-1 accuracy, FastViT is 4.9× faster than EfficientNet-B5, 3.5× faster than CMT-S, and 1.9× faster than ConvNeXt-B on iPhone 12 Pro.It is also 1.6× faster than EfficientNetV2-S at the same accuracy.
  • Evaluation: FastViT generalizes across classification, detection, segmentation, and 3D hand mesh estimation while improving latency and robustness to corruption and out-of-distribution samples.The authors report significant speed advantages over competing robust models and across multiple downstream tasks.

2. Related Work

Vision models increasingly combine convolutional and transformer components to balance local information, global context, accuracy, and efficiency. Related work reduces transformer costs through hybrid structures, while structural reparameterization and factorized convolutions target memory access, parameters, and FLOPs.

  • Transformer efficiency: Vision transformers model long-range dependencies through self-attention but often incur high computational costs.Efficient alternatives seek lower latency while retaining global-context modeling.
  • Hybrid Vision Transformers: Hybrid vision transformers combine convolutional and transformer designs to capture local and global information efficiently.Prior designs use convolutional stems, early convolutional stages, or windowed attention, while others build explicit hybrid structures.
  • Structural Reparameterization: Structural reparameterization lowers memory access cost by converting skip-connected training structures into simpler inference-time components.FastViT applies this idea through RepMixer, a reparameterizable token mixer.
  • Structural Reparameterization: Factorized convolutions use depthwise or grouped convolutions followed by 1×1 pointwise convolutions to reduce efficiency costs, though lower parameter counts can reduce capacity.Train-time overparameterization is introduced to address the capacity reduction associated with factorization.

3. Architecture

FastViT uses a four-stage hybrid architecture that combines RepMixer-based convolutional token mixing with self-attention in the final stage. Its design reparameterizes skip connections and factorized convolutions for efficient inference, while train-time overparameterization and large kernels improve capacity and receptive field.

  • Large Kernel Convolutions: Large-kernel depthwise convolutions replace self-attention in early stages to enlarge the receptive field while improving model robustness.They are incorporated into FFN and patch-embedding layers.
  • Reparameterizing Skip Connections: RepMixer rearranges convolutional mixing so its skip connection can be reparameterized into a single depthwise convolution at inference.This removes skip connections and reduces memory access cost.
  • Overview: FastViT has four stages operating at different scales, with RepMixer in early stages and self-attention used for token mixing in stage 4.
  • Reparameterizing Skip Connections: 25.1% lower latency at 384×384 and 43.9% lower latency at 1024×1024 are obtained with RepMixer instead of Pooling at approximately 1.8G FLOPs.The comparison was measured on an iPhone 12 Pro using MetaFormer S12 variants.
  • Linear Train-time Overparameterization: Factorizing dense k×k convolutions into depthwise and pointwise convolutions improves efficiency, while linear train-time overparameterization restores capacity and is removed at inference.The approach is applied to convolutional stem, patch embedding, and projection layers.
  • Linear Train-time Overparameterization: Train-time overparameterization increases training time by 6.7% for FastViT-SA12 and 4.4% for FastViT-SA36 under the reported settings.Only selected factorized-convolution layers are overparameterized, limiting the training-time overhead.

4. Experiments

FastViT is evaluated across classification, robustness, detection, segmentation, and 3D hand mesh regression on mobile and desktop hardware. It achieves strong accuracy-latency trade-offs and remains competitive or superior across these settings.

  • Image Classification: ImageNet-1k experiments compare FastViT with state-of-the-art models using standard training and deployment-aware evaluation settings.Models are trained for 300 epochs with AdamW and evaluated across mobile and desktop compute fabrics.
  • Image Classification: At 83.9% ImageNet Top-1 accuracy, FastViT-MA36 is 1.9× faster than optimized ConvNeXt-B on iPhone 12 Pro and 2.0× faster on GPU.At 84.9% accuracy, it matches NFNet-F1 GPU speed while being 66.7% smaller and using 50.1% fewer FLOPs.
  • 3D Hand Mesh Estimation: FastViT outperforms other real-time methods on joint- and vertex-error metrics for FreiHAND while being 1.9× faster than MobileHand and 2.8× faster than MobRecon.The approach uses ImageNet-1k pretraining followed by training exclusively on FreiHAND.
  • Semantic Segmentation: FastViT backbones are also evaluated for ADE20K semantic segmentation using Semantic FPN, with FLOPs and latency measured on 512×512 crops.Models are initialized from corresponding image-classification weights.
  • Object Detection: On MS-COCO, FastViT achieves state-of-the-art performance across multiple latency regimes; FastViT-MA36 matches CMT-S while being 2.4× faster on GPU and 4.3× faster on mobile.Evaluation uses Mask-RCNN with a 1x training schedule and pretrained classification backbones.

5. Conclusion

FastViT is presented as a general-purpose hybrid vision transformer designed for efficient execution on mobile devices and desktop GPUs. Structural reparameterization reduces memory access cost while the architecture improves performance across downstream tasks and robustness evaluations.

  • Efficiency: FastViT is efficient across mobile devices and desktop-grade GPUs, with reduced memory access cost from structural reparameterization.The paper reports improvements in runtime, especially at higher resolutions.
  • Generalization: FastViT improves performance on ImageNet classification, object detection, semantic segmentation, and 3D hand mesh estimation.The same backbone is evaluated across multiple computer vision tasks.
  • Robustness: FastViT is reported to be robust to out-of-distribution samples while being significantly faster than competing robust models.The conclusion links robustness and deployment efficiency within the evaluated scope.

A. Ablations

The ablations examine replacing self-attention in early stages with efficient alternatives and selecting convolutional kernel sizes for accuracy-latency efficiency. RepMixer becomes increasingly advantageous at higher resolutions, while 7×7 kernels provide the chosen trade-off.

  • Self-Attention Substitution: Early-stage self-attention is costly because token counts are high, motivating efficient alternatives with small accuracy degradation and lower latency.The analysis compares these choices across FastViT-S12 stages.
  • Kernel-Size Ablation: Performance stagnates beyond a 7×7 kernel while FLOPs, latency, and parameter count continue increasing, motivating the 7×7 choice.The ablation measures ImageNet-1k Top-1 accuracy and iPhone 12 Pro latency.

A.3. Training Time

Training-time overparameterization is applied selectively rather than throughout the architecture. The ablation reports no benefit from overparameterizing ConvFFN layers and avoids significant training-time degradation through partial overparameterization.

  • Training-Time Overparameterization: FastViT overparameterizes only the convolutional stem and patch embedding layers during training, rather than every architectural element.The additional branches are intended for training-time use and are not applied to dense ConvFFN blocks.
  • Training-Time Overparameterization: Overparameterizing ConvFFN layers provides no empirical improvement, while partial overparameterization avoids significant training-time degradation.The claim is supported by the reported ablation and train-time measurements.

B.1. Benchmarking

The benchmarking protocol evaluates FastViT and competing architectures on an iPhone 12 Pro and an NVIDIA RTX-2080Ti, while also examining implementation and training configurations. ConvNeXt is modified to reduce mobile latency from reshape operations.

  • Benchmarking: Median latency is measured over 100 estimates after 100 TensorRT warmup iterations on the RTX-2080Ti.Image classification uses batchsize 8; other batch-size details are truncated in the passage.
  • Benchmarking: ConvNeXt latency increases because reshape operations are inefficient on mobile devices.A channel-first LayerNorm implementation and 1×1 convolutions replace the original operations to improve runtime.
  • Benchmarking: Table 14 benchmarks ConvNeXt before and after the described modifications.The supplied passage identifies the table but does not provide its numerical results.
  • Benchmarking: ImageNet-1k models are trained at 256×256 resolution and fine-tuned at 384×384 resolution.FastViT variants use the listed training hyperparameters, with RegNetY-16GF as the distillation teacher.
  • Benchmarking: Figure 4 plots accuracy against mobile latency for recent mobile architectures and FastViT variants on an iPhone 12 Pro.The benchmark uses the image sizes specified in Table 16.
  • Benchmarking: Different training seeds produce Top-1 accuracy results within ±0.2%.This reports the observed variation across seeds.

B.3. Comparison with Mobile Architectures

FastViT is compared with efficient mobile architectures on ImageNet-1k, including the convolutional MobileOne and lightweight transformer EdgeViT. The models are designed within a compact parameter regime using staged width and compute configurations.

  • B.3. Comparison with Mobile Architectures: FastViT outperforms the purely convolutional MobileOne architecture and the lightweight ViT architecture EdgeViT.The supplied passages do not include the corresponding numerical accuracy or latency values.
  • B.3. Comparison with Mobile Architectures: The mobile-architecture comparison is reported in Table 16 for ImageNet-1k classification.HardSwish is replaced with GELU when needed because Core ML support is limited.
  • B.3. Comparison with Mobile Architectures: FastViT architectures in this work contain fewer than 50M parameters for efficient deployment.Most variants use a 1:1:3:1 stage compute ratio, while the smallest uses 1:1:2:1.
  • B.3. Comparison with Mobile Architectures: FastViT widths double at each new stage, with configurations including [48, 96, 192, 384], [64, 128, 256, 512] and [76, 152, 304, 608].The supplied passage lists these width configurations for the paper’s variants.

B.5. 3D Hand mesh estimation

FastViT is used for 3D hand mesh estimation by regressing camera, pose, and MANO shape parameters from images. On FreiHand, qualitative results show complicated gestures and reliable poses despite occlusion from hand-held objects.

  • B.5. 3D Hand mesh estimation: The framework uses single fully connected regression layers for weak-perspective camera, MANO pose, and shape parameters.It regresses 6D rotations for all MANO joints and minimizes 3D vertex and joint losses among its objectives.
  • B.5. 3D Hand mesh estimation: FreiHand contains 130,240 training images and 3,960 test images for the mesh-estimation experiments.Training uses 224×224 images, Adam, 200 epochs, and ImageNet-1k-pretrained backbone weights.
  • B.5. 3D Hand mesh estimation: The framework models complicated gestures on the FreiHand test set.Figure 6 presents qualitative predictions projected onto images using a weak-perspective camera model.
  • B.5. 3D Hand mesh estimation: The model predicts reliable hand poses despite occlusion from hand-held objects.The camera-model parameters used for projection are also predicted by the model.
Loading 2303.14189v2…