Technical Summary: ScalePredictor

Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance in computer vision but face significant deployment challenges on resource-constrained edge devices due to their computational intensity. Post-Training Quantization (PTQ) offers a solution by compressing models to low-bit formats (e.g., 8-bit or lower) without retraining. However, existing PTQ methods predominantly rely on static quantization, where a single set of quantization parameters (scales) is applied uniformly to all input instances.

This static paradigm is inherently suboptimal for ViTs because natural images exhibit substantial diversity, leading to significant variations in activation distributions across different samples. This heterogeneity is particularly detrimental in ultra-low-bit settings (e.g., W3/A3), where fixed scales fail to capture instance-specific dynamics. Conversely, dynamic quantization strategies that adapt parameters per sample often require runtime full-tensor reductions (min/max calculations) for every layer. These operations introduce severe synchronization barriers and memory-bound overheads on modern hardware, drastically undermining inference speed and negating the efficiency benefits of quantization.

Methodology: ScalePredictor

The authors propose ScalePredictor, a dynamic quantization framework that achieves instance-aware scale learning with minimal computational overhead. The core insight driving the method is a discovered correlation: the optimal quantization scales for deeper layers are strongly dependent on the activation magnitude observed in shallow layers.

The framework consists of three main components:

1. Robust Range Extraction:
   Instead of computing global min/max values (sensitive to outliers) or performing expensive percentile sorting, ScalePredictor extracts a "robust range" ($r(x)$) from the Patch Embedding output (the shallowest layer).

   • The flattened tensor is partitioned into $K$ non-overlapping chunks.
   • Local min/max values are computed for each chunk.
   • These local extrema are averaged to derive a global robust range.
   • This approach acts as a coarse-grained percentile approximation, effectively suppressing outlier influence while maintaining $O(N \cdot C)$ complexity and avoiding sorting.

2. Taylor-Motivated Polynomial Predictor:
   A learnable polynomial module maps the extracted robust range $r(x)$ to the optimal quantization scales for all subsequent layers.

   • Based on the observation that the relationship between shallow statistics and deep optimal scales follows a smooth, non-linear trend, the method models the scale $s_l(x)$ for layer $l$ as a polynomial function of $r(x)$:
     $$s_l(x) = \sum_{k=0}^{d} c_{l,k} \cdot (r(x))^k$$
   • The coefficients $c_{l,k}$ are learned during a calibration phase using a block-wise reconstruction loss.
   • This allows the model to predict instance-specific scales for all layers in a single pass, without accessing deep activation tensors.

3. Calibration and Inference:

   • Calibration: The polynomial coefficients are optimized using a small calibration set. The constant term ($c_{l,0}$) is initialized with the optimal static scale to ensure stability, while higher-order terms are learned to capture instance variations.
   • Inference: For a new input, the robust range is computed once at the Patch Embedding layer. All subsequent layer scales are generated via polynomial evaluation ($O(1)$ FLOPs), eliminating the need for per-layer runtime reductions.

Key Contributions

1. Instance-Aware PTQ for ViTs: ScalePredictor is presented as the first post-training quantization method for ViTs that leverages dynamic quantization to address instance-level distribution heterogeneity without retraining.
2. Efficient Scale Learning Mechanism: The integration of a robust range extraction mechanism (chunk-based min/max averaging) with a Taylor-inspired polynomial projection enables accurate, on-the-fly scale prediction with negligible inference overhead.
3. Discovery of Shallow-Deep Correlation: The work reveals a consistent functional dependency where shallow-layer activation statistics (specifically from Patch Embeddings) can effectively predict optimal quantization scales for deeper transformer blocks across various architectures (ViT, DeiT, Swin).

Experimental Results

Extensive experiments were conducted on the ImageNet dataset across seven ViT architectures (ViT-S/B, DeiT-T/S/B, Swin-S/B) and four PTQ baselines (BRECQ, QDrop, AdaLog, I&S-ViT).

• Accuracy Improvements: ScalePredictor consistently outperforms static PTQ baselines, particularly in low-bit regimes.
  • Notable gains include a 10.58% improvement (from 33.57% to 44.15%) on Swin-Small using I&S-ViT at W2/A3.
  • Significant improvements were also observed at W3/A3, such as a 8.71% gain on ViT-Small using BRECQ.
  • Even in cases where static baselines collapsed (e.g., BRECQ on ViT-Small at W2/A3 with 1.10% accuracy), ScalePredictor recovered performance to 4.15%.
• Comparison with JIT Dynamic Quantization: Unlike Just-In-Time (JIT) dynamic quantization methods (e.g., JIT Min/Max), which severely degrade accuracy by breaking the calibrated weight-activation relationship, ScalePredictor preserves calibration consistency while enabling adaptivity.
• Efficiency:
  • Latency: On an NVIDIA RTX 3090, ScalePredictor introduces only 0.6% end-to-end latency overhead compared to static PTQ.
  • Scale Computation: The method is 40–52× faster than JIT Min/Max and 507–1025× faster than JIT Percentile for computing quantization scales, as it avoids per-layer tensor reductions.

Significance and Claims

The paper claims that ScalePredictor successfully bridges the gap between the accuracy of dynamic quantization and the efficiency of static quantization. By leveraging the hidden correlation between shallow-layer distributions and deep-layer optimal scales, the method avoids the costly "just-in-time" calibration required by traditional dynamic approaches.

The authors emphasize that ScalePredictor offers a favorable accuracy–efficiency trade-off, making it a viable solution for deploying low-bit Vision Transformers on edge devices. The method is designed to be seamlessly integrated into existing PTQ frameworks, requiring only the addition of a lightweight predictor module and a robust range extraction step, without altering the underlying network architecture or requiring full retraining.