[Paper Review] Squat: Quant Small Language Models on the Edge

Xuan Shen, Peiyan Dong+10 more|arXiv (Cornell University)|Feb 16, 2024

Advanced Neural Network Applications4 citations

TL;DR

EdgeQAT proposes a novel quantization-aware training framework that jointly optimizes weights and activations for lightweight LLMs using entropy and distribution guidance to minimize information distortion in self-attention mechanisms. It achieves up to 2.37× inference speedup on edge devices like Raspberry Pi 5 with 4-bit weights and activations, outperforming FP16 baselines while preserving model accuracy.

ABSTRACT

A growing trend has emerged in designing high-quality Small Language Models (SLMs) with a few million parameters. This trend is driven by the increasing concerns over cloud costs, privacy, and latency. Considering that full parameter training is feasible for SLMs on mobile devices, Quantization-Aware Training (QAT) is employed to improve efficiency by reducing computational overhead and memory footprint. However, previous QAT works adopt fine-grained quantization methods to compress models with billions of parameters on GPUs, incompatible with current commodity hardware, such as mobile and edge devices, which relies on Single Instruction Multiple Data (SIMD) instructions. Thus, the generalization of these methods to SLMs on mobile devices is limited. In this paper, we propose Squat method, an effective QAT framework with deployable quantization for SLMs on mobile devices. Specifically, we propose entropy-guided and distribution-aligned distillation to mitigate the distortion of attention information from quantization. Besides, we employ sub-8-bit token adaptive quantization, assigning varying bit widths to different tokens based on their importance. Furthermore, we develop a SIMD-based Multi-Kernel Mixed-Precision (MKMP) multiplier to support sub-8-bit mixed-precision MAC on mobile devices. Our extensive experiments verify the substantial improvements of our method compared to other QAT methods across various datasets. Furthermore, we achieve an on-device speedup of up to 2.37x compared with its FP16 counterparts, signaling a great advancement. Code: https://github.com/shawnricecake/squant

Motivation & Objective

• To address the performance degradation in low-bit quantization of both weights and activations in lightweight LLMs on edge devices.
• To identify that information distortion in quantized attention maps—particularly in query and key distributions—is the primary cause of accuracy drop.
• To develop a QAT framework that jointly optimizes weight and activation quantization to enable full utilization of INT8 operations on edge hardware.
• To introduce adaptive quantization based on token importance to further reduce bit width without sacrificing accuracy.
• To achieve high-speed, low-resource inference on edge platforms such as smartphones, Raspberry Pi, and FPGAs.

Proposed method

• Proposes entropy maximization on query and key tensors to minimize quantization error and restore information content.
• Introduces a distribution-guided loss that minimizes cosine similarity difference between FP16 and quantized attention maps.
• Designs a token importance-aware mechanism to dynamically assign 4-bit or 8-bit precision based on activation patterns.
• Applies QAT to both weights and activations, enabling full acceleration via integer-only operations on edge devices.
• Employs a hybrid quantization strategy where high-importance tokens use 8-bit precision and low-importance ones use 4-bit, optimizing for efficiency.
• Uses a unified training pipeline that integrates entropy, distribution, and adaptive bit-width optimization into a single end-to-end QAT framework.

Experimental results

Research questions

• RQ1 Why does Post-Training Quantization (PTQ) fail when quantizing weights, activations, and KV cache together below 8 bits?
• RQ2 What causes the primary performance degradation in low-bit quantization of LLMs, particularly in attention mechanisms?
• RQ3 Can entropy maximization and distribution alignment jointly reduce quantization error in attention maps?
• RQ4 How effective is adaptive, importance-aware activation quantization in reducing average bit width while preserving accuracy?
• RQ5 Can full QAT of weights and activations achieve significant speedup on real edge devices like Raspberry Pi and FPGAs?

Key findings

• EdgeQAT achieves up to 2.37× speedup on Raspberry Pi 5 with W4A4 quantization, significantly outperforming FP16 inference.
• On-device speedup of 1.91× is achieved on the OnePlus 11 smartphone and up to 1.9× on FPGA platforms with W4A4 settings.
• The method achieves state-of-the-art accuracy on SuperGLUE with 4-bit weights and activations, outperforming QAT baselines like PACT and LSQ.
• Adaptive quantization with 25% of tokens using 8-bit precision achieves higher accuracy than equivalent non-adaptive 5-bit quantization.
• The distribution-guided loss reduces attention map discrepancy between FP16 and quantized outputs, directly mitigating accuracy loss.
• The entropy-based optimization of query and key distributions effectively reduces quantization error, especially in low-bit regimes.