Home

Awesome

Jumping Through Local Minima: Quantization in the Jagged Loss Landscape of Vision Transformers (Evol-Q) <br><sub>Official PyTorch implementation of the ICCV 2023 paper</sub>

Jumping Through Local Minima: Quantization in the Jagged Loss Landscape of Vision Transformers<br> Natalia Frumkin, Dibakar Gope, and Diana Marculescu <br>https://arxiv.org/pdf/2308.10814<br>

Abstract: Quantization scale and bit-width are the most important parameters when considering how to quantize a neural network. Prior work focuses on optimizing quantization scales in a global manner through gradient methods (gradient descent & Hessian analysis). Yet, when applying perturbations to quantization scales, we observe a very jagged, highly non-smooth test loss landscape. In fact, small perturbations in quantization scale can greatly affect accuracy, yielding a 0.5−0.8% accuracy boost in 4-bit quantized vision transformers (ViTs). In this regime, gradient methods break down, since they cannot reliably reach local minima. In our work, dubbed Evol-Q, we use evolutionary search to effectively traverse the non-smooth landscape. Additionally, we propose using an infoNCE loss, which not only helps combat overfitting on the small calibration dataset (1,000 images) but also makes traversing such a highly non-smooth surface easier. Evol-Q improves the top-1 accuracy of a fully quantized ViT-Base by 10.30%, 0.78%, and 0.15% for 3-bit, 4-bit, and 8-bit weight quantization levels. Extensive experiments on a variety of CNN and ViT architectures further demonstrate its robustness in extreme quantization scenarios.

Requirements

Getting Started

To run an initial experiment with 8-bit quantized DeiT-Tiny:

sh scripts/run.sh

This script contains details for how to quantize 3,4,8-bit models as well as model architectures.

Our repository is adapted from FQ-ViT and Meta's LeVit implementation. We have created our novel block-wise evolutionary algorithm in joint_evol_opt.py which can quantize the following model flavors out of the box:

Below, we show the accuracy for each 4-bit quantized model, averaged across seeds {0 ,1, 2}:

ModelTop-1 Accuracy
ViT-Base
DeiT-Tiny
DeiT-Small
DeiT-Base
LeViT-128S
LeViT-128
LeViT-192
LeViT-246
LeViT-384