[논문 리뷰] Squat: Quant Small Language Models on the Edge

Xuan Shen, Peiyan Dong외 10명|arXiv (Cornell University)|2024. 02. 16.

Advanced Neural Network Applications인용 수 4

한 줄 요약

EdgeQAT는 자가주의 메커니즘 내의 정보 왜곡을 최소화하기 위해 엔트로피 및 분포 가이던스를 사용하여 경량 LLM의 가중치와 활성화를 공동 최적화하는 새로운 양자화 인식 훈련 프레임워크를 제안한다. 이는 라즈베리 파이 5와 같은 엣지 디바이스에서 4비트 가중치와 활성화를 사용할 때 최대 2.37배의 추론 속도 향상을 달성하며, 정확도를 유지하면서도 FP16 기준선을 능가한다.

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

연구 동기 및 목표

• 엣지 디바이스에서 경량 LLM의 저비트 양자화 시 가중치와 활성화의 성능 저하 문제를 해결하기 위해.
• 정확도 저하의 주요 원인이 양자화된 어텐션 맵 내의 정보 왜곡, 특히 쿼리 및 키 분포의 왜곡임을 규명하기 위해.
• 엣지 하드웨어에서 INT8 연산을 전면적으로 활용할 수 있도록 가중치 및 활성화 양자화를 공동 최적화하는 QAT 프레임워크를 개발하기 위해.
• 토큰 중요도 기반의 적응형 양자화를 도입하여 정확도를 훼손하지 않은 채 비트 폭을 추가로 감소시키기 위해.
• 스마트폰, 라즈베리 파이, FPGA와 같은 엣지 플랫폼에서 고속·저자원 추론을 달성하기 위해.

제안 방법

• 쿼리 및 키 텐서에 대한 엔트로피 최대화를 제안하여 양자화 오차를 최소화하고 정보 내용을 복원하기 위해.
• FP16 및 양자화된 어텐션 맵 간의 코사인 유사도 차이를 최소화하는 분포 가이드 손실을 도입하기 위해.
• 활성화 패턴에 기반해 동적으로 4비트 또는 8비트 정밀도를 할당하는 토큰 중요도 인식 메커니즘을 설계하기 위해.
• 가중치 및 활성화 모두에 QAT를 적용하여 엣지 디바이스에서 정수 연산만을 사용한 전면 가속을 가능하게 하기 위해.
• 고중요도 토큰은 8비트 정밀도, 저중요도 토큰은 4비트 정밀도를 사용하는 하이브리드 양자화 전략을 적용하여 효율성 최적화하기 위해.
• 엔트로피, 분포, 적응형 비트 폭 최적화를 하나의 종단 간 QAT 프레임워크에 통합한 통합 훈련 파이프라인을 사용하기 위해.

실험 결과

연구 질문

• RQ1 8비트 이하로 동시에 가중치, 활성화, KV 캐시를 양자화할 때 Post-Training Quantization (PTQ)가 실패하는 이유는 무엇인가?
• RQ2 LLM의 저비트 양자화에서 특히 어텐션 메커니즘에서 성능 저하의 주요 원인은 무엇인가?
• RQ3 엔트로피 최대화와 분포 정렬이 어텐션 맵의 양자화 오차를 공동으로 감소시킬 수 있는가?
• RQ4 중요도 인식 기반의 적응형 활성화 양자화가 평균 비트 폭을 줄이는 데 얼마나 효과적인가?
• RQ5 가중치 및 활성화의 전면 QAT가 실제로 라즈베리 파이 및 FPGA와 같은 엣지 디바이스에서 상당한 속도 향상을 이끌 수 있는가?

주요 결과

• Raspberry Pi 5에서 W4A4 양자화를 사용할 경우 EdgeQAT는 최대 2.37배의 속도 향상을 달성하며, FP16 추론을 크게 능가한다.
• OnePlus 11 스마트폰에서의 현장 속도 향상은 1.91배이며, FPGA 플랫폼에서는 W4A4 설정에서 최대 1.9배의 속도 향상을 기록한다.
• 4비트 가중치와 활성화를 사용할 때 SuperGLUE에서 최신 기준 정확도를 확보하며, PACT 및 LSQ와 같은 QAT 기준선을 능가한다.
• 25%의 토큰이 8비트 정밀도를 사용하는 적응형 양자화는 동등한 비적응형 5비트 양자화보다 높은 정확도를 달성한다.
• 분포 가이드 손실은 FP16 출력과 양자화된 출력 간의 어텐션 맵 불일치를 감소시켜 정확도 손실을 직접 완화한다.
• 엔트로피 기반의 쿼리 및 키 분포 최적화는 특히 저비트 환경에서 양자화 오차를 효과적으로 감소시킨다.