ViT-1.58b: Mobile Vision Transformers in the 1-bit Era

The rapid advancement of Transformer Vaswani et al. (2017) architectures has significantly impacted the field of computer vision, particularly with the introduction of Vision Transformers (ViTs) Dosovitskiy et al. (2020). By treating image patches as tokens and leveraging the attention mechanism to process image patches, ViTs effectively capture complex dependencies across entire images, achieving remarkable performance in various image classification tasks Graham et al. (2021); Liu et al. (2021); Yuan et al. (2021). However, the computational and memory demands of ViTs are substantial, stemming primarily from their attention mechanisms, which scale quadratically with the number of tokens Dosovitskiy et al. (2020). This inherent complexity poses significant challenges for deploying ViTs in resource-constrained environments such as mobile devices and embedded systems, where energy efficiency and low latency are critical

Recently, research in neural network efficiency to mitigate these demands has explored various strategies including model pruning Pan et al. (2023); Yu and Wu (2023); Song et al. (2022), knowledge distillation Yang et al. (2022); Touvron et al. (2021); Hao et al. (2022), and quantization Li and Gu (2023); Sun et al. (2022); Lin et al. (2021). Among these, quantization techniques are particularly effective as they directly reduce the precision of weights and activations, thereby significantly lowering the memory and computational requirements of deep learning models. Traditionally, post-training quantization has been favored for its simplicity as it does not necessitate changes to the training pipeline or retraining of the model. However, this method often results in significant accuracy loss at lower precision levels because the model isn’t optimized for the quantized representation during training Frantar et al. (2022); Chee et al. (2024), limiting its utility for high-performance applications. In contrast, Quantization-Aware Training (QAT) Park et al. (2018) integrates the effects of quantization into the training process itself, simulating quantization effects to typically achieve better accuracy than post-training methods. For example, extreme quantization, like 1-bit models used in Binarized neural networks (BNNs) Courbariaux et al. (2016), which utilize binary weights and activations, significantly decreasing computational and memory demands. Recent adaptations of 1-bit quantization techniques to Transformer models, such as BitNet Wang et al. (2023), and BiVit He et al. (2023) have shown that even severe quantization can maintain performance while substantially cutting resource consumption.

However, these 1-bit models using binary quantization often face challenges in preserving model accuracy due to the extreme reduction in weight precision. To address this limitation, research has ventured into ternary quantization, which offers a more balanced approach. A notable example is "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits" Ma et al. (2024), which explores 1.58-bit quantization where weights can assume values of -1, 0, or 1. This approach refines the balance between performance and computational demands by introducing a zero value, which significantly reduces computational overhead. The inclusion of a non-polar weight (zero) not only allows for sparsity, thereby reducing the number of active computations but also maintains a richer representation than binary weights, potentially leading to less information loss. The success of this quantization strategy in large language models suggests its promising applicability to Vision Transformers, enabling more efficient processing of visual data while preserving model accuracy.

To address the unique demands of large-scale ViT models, we introduce ViT-1.58b, a 1.58-bit quantized ViT model. ViT-1.58b utilizes ternary quantization to optimize both memory and computational efficiency while maintaining competitive performance. This approach leverages the benefits of extreme quantization demonstrated in language models, adapting them to the context of computer vision. Our contributions are summarized as follows:

• •

  We propose ViT-1.58b, the first 1.58-bit quantized ViT model, optimized for efficient scaling in terms of both memory and computation.

• •

  We demonstrate the effectiveness of ViT-1.58b on the CIFAR-10 and ImageNet-1k datasets, showing competitive accuracy with significant reductions in memory and computational costs.

• •

  We provide a comprehensive evaluation, comparing ViT-1.58b with state-of-the-art quantization methods and full-precision ViT models, highlighting its advantages in resource-constrained environments.

As shown in Figure 1, our ViT-1.58b architecture is primarily based on ViT for image classification. The process begins by dividing the input image into patches, followed by applying a linear projection to the flattened patches. These patches then undergo patch and position embedding, as well as class embedding, before being fed into the transformer encoder. The output from the transformer encoder is processed through an MLP to produce the final classification prediction. We employ BitLinear layers from BitNet b1.58 Ma et al. (2024) to replace the conventional nn.Linear layers in the transformer encoder. Different quantization functions are employed to quantize the weights to 1.58-bit precision and the activations to 8-bit precision, ensuring both efficiency and performance. Next, we will introduce how we implement BitLinear layers to achieve 1.58-bit weights and 8-bit activations.

Using this method, we can convert the weight matrix $W$ into a ternary matrix $\tilde{W}$, where each weight value is constrained to {-1, 0, 1}.

The core component of ViT-1.58b is the BitLinear layer, which replaces the traditional nn.Linear layer in the Transformer. Initially, the activations undergo a normalization layer Ba et al. (2016) to ensure that the activation values have a variance of 1 and a mean of 0, mathematically represented as $\text{LN}(x)=\frac{x-\mathbb{E}(x)}{\sqrt{\text{Var}(x)+\epsilon}}$. Following this, normalized activations are quantized using the absmax quantization function, resulting in $\text{Quant}(\text{LN}(x))$. With weights quantized to 1.58-bit precision, matrix multiplication results in the output $y=\tilde{W}\cdot\text{Quant}(\text{LN}(x))$, and the output is subsequently dequantized to rescaled with $\beta$ and $\gamma$ to the original precision, expressed as $y_{dequant}=y\times\frac{\beta\gamma}{Q_{b}}$, where $\beta=\frac{1}{nm}\|W\|_{1}$

These steps collectively enable the BitLinear layer to maintain the performance of the Transformer model while significantly reducing computational cost and storage requirements.

Additionally, we use mixed precision training, where weights and activations are quantized to low precision, but gradients and optimizer states are kept in high precision to ensure training stability and accuracy.

The computational framework for this study comprised four NVIDIA TESLA A100 GPUs, each with 80 GB of VRAM. The system’s processing core utilized an AMD EPYC 7552 48-core Processor with 80 GB of system RAM to manage extensive datasets efficiently. We employed PyTorch version 2.0.0 integrated with CUDA 11.8, optimizing tensor operations across GPUs.

For CIFAR-10, the full-precision ViT-L model achieved a test accuracy of 76.28%, which serves as a baseline for comparison. When the bit precision was reduced to 16-bit, there was a moderate decline in accuracy to 74.61%. Further reduction to 8-bit and 4-bit resulted in more pronounced declines to 72.20% and 70.69%, respectively. This trend suggests that lower bit precision generally degrades the model’s performance, likely due to the increased quantization error and reduced capacity to capture the variability in the data. The proposed 1.58-ViT-1.58b-L model, which operates at an even lower bit precision than the other quantized variants, recorded a test accuracy of 72.27%. Interestingly, the performance of this model is closer to the 8-bit version than to the 4-bit.

On the more complex and diverse ImageNet-1k dataset, a similar pattern is observed. The full-precision ViT-L model achieved a Top-1 accuracy of 76.54% and a Top-3 accuracy of 90.23%. Reduction to 16-bit precision caused declines to 75.29% and 87.44% in Top-1 and Top-3 accuracies, respectively. These declines became more substantial at 8-bit (Top-1: 74.11%, Top-3: 85.26%) and 4-bit (Top-1: 70.88%, Top-3: 82.50%). The 1.58-bit model managed a Top-1 accuracy of 74.25% and a Top-3 accuracy of 85.78%, showcasing a performance that surpasses the 8-bit version in Top-1 accuracy and nearly matches it in Top-3 accuracy.

As shown in Figure 2, ViT-1.58b-L shows promising results in both training loss and memory consumption compared to the full-precision ViT-L and its quantized versions. The left panel depicts training loss curves, where ViT-1.58b-L’s loss closely follows that of the full-precision ViT-L. This suggests that the 1.58-bit quantization retains the model’s learning capacity effectively. The right panel highlights the memory consumption. The full-precision ViT-L requires 1.14 GB of memory, while ViT-1.58b-L drastically reduces this to just 57 MB. This significant reduction in memory usage makes our model ideal for deployment in resource-constrained environments. Overall, our ViT-1.58b-L model balances competitive performance with substantial memory savings, demonstrating its efficiency and practicality for real-world applications.

In Figure 3, as bit precision decreases, computational performance (measured in TFLOPs) significantly increases. Full-precision achieves 0.692 TFLOPs, while ours reaches 11.069 TFLOPs. This substantial increase in performance with lower precision highlights the efficiency gains achievable through extreme quantization. Our model offers a dramatic boost in computational throughput, making it highly suitable for high-performance computing environments where both speed and resource efficiency are critical.

In this paper, we introduced ViT-1.58b, a 1.58-bit quantized Vision Transformer that efficiently addresses computational and memory challenges in vision model deployment through ternary quantization. Our results show that ViT-1.58b achieves competitive accuracy on benchmarks like CIFAR-10 and ImageNet-1k with significantly lower resource requirements. This model demonstrates the potential of advanced quantization techniques in complex Transformer architectures, highlighting their role in developing sustainable AI solutions. Future work will explore ViT-1.58b’s scalability and integration into larger systems, enhancing its practical utility and environmental benefits.

While the ViT-1.58b model exhibits promising performance and demonstrates efficient computational usage, it also presents certain limitations that need to be addressed. This section outlines the primary challenges associated with our approach.