EdgeQAT: Entropy and Distribution Guided Quantization-Aware Training for the Acceleration of Lightweight LLMs on the Edge

Anonymous

Abstract: Despite the remarkable strides of Large Language Models (LLMs) in various fields, the wide applications of LLMs on edge devices are limited due to their massive parameters and computations. To address this, quantization is commonly adopted to generate lightweight LLMs with efficient computations and fast inference. However, Post-Training Quantization (PTQ) methods dramatically degrade in quality when quantizing weights, activations, and KV cache together to below 8 bits. Besides, many Quantization-Aware Training (QAT) works quantize model weights, leaving the activations untouched, which do not fully exploit the potential of quantization for inference acceleration on the edge. In this paper, we propose EdgeQAT, the Entropy and Distribution Guided QAT for the optimization of lightweight LLMs to achieve inference acceleration on Edge devices. We first identify that the performance drop of quantization primarily stems from the information distortion in quantized attention maps, demonstrated by the different distributions in quantized query and key of the self-attention mechanism. Then, the entropy and distribution guided QAT is proposed to mitigate the information distortion. Moreover, we design a token importance-aware adaptive method to dynamically quantize the tokens with different bit widths for further optimization and acceleration. Our extensive experiments verify the substantial improvements with our framework across various datasets. Furthermore, we achieve an on-device speedup of up to 2.37x compared with its FP16 counterparts across multiple edge devices, signaling a groundbreaking advancement.

Paper Type: long

Research Area: Efficient/Low-Resource Methods for NLP

Contribution Types: Approaches to low-resource settings, Approaches low compute settings-efficiency

Languages Studied: English

Preprint Status: We are considering releasing a non-anonymous preprint in the next two months (i.e., during the reviewing process).

A1: yes

A1 Elaboration For Yes Or No: 6

A2: n/a

A2 Elaboration For Yes Or No: Our work is based on the public dataset. We compress models which is not related to any ethical issues or negative societal aspects.

A3: yes

A3 Elaboration For Yes Or No: 1

B: no

B1: n/a

B2: n/a

B3: n/a

B4: no

B4 Elaboration For Yes Or No: We use the public datasets, and the same setting as previous works.

B5: n/a

B6: yes

B6 Elaboration For Yes Or No: 5

C: yes

C1: yes

C1 Elaboration For Yes Or No: 5.1

C2: yes

C2 Elaboration For Yes Or No: 5.1

C3: yes

C3 Elaboration For Yes Or No: 5.3

C4: yes

C4 Elaboration For Yes Or No: 5.1

D: no

D1: n/a

D2: n/a

D3: n/a

D4: n/a

D5: n/a

E: no

E1: n/a