[論文レビュー] Squat: Quant Small Language Models on the Edge

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

Advanced Neural Network Applications被引用数 4

ひとこと要約

EdgeQATは、自己注意機構における情報歪みを最小限に抑えるためにエントロピーおよび分布ガイド付きの手法を用いて、軽量LLMの重みと活性化を共同最適化する画期的な量子化意識学習フレームワークを提案する。Raspberry Pi 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フレームワークを開発すること。
• トークンの重要度に基づく適応的量子化を導入し、精度を損なわずにビット幅をさらに削減すること。
• スマートフォン、Raspberry Pi、FPGAなどのエッジプラットフォーム上で、高速かつリソース効率の良い推論を達成すること。

提案手法

• クエリおよびキーテンソルにおけるエントロピー最大化により、量子化誤差を最小化し、情報含量を回復すること。
• FP16と量子化されたアテンションマップ間のコサイン類似度の差を最小化する分布ガイド付き損失を導入すること。
• 活性化パターンに基づき、動的に4ビットまたは8ビット精度を割り当てるトークンの重要度に配慮したメカニズムを設計すること。
• 重みと活性化の両方に対してQATを適用し、エッジデバイス上で完全に整数演算のみを用いた高速化を実現すること。
• 高重要度のトークンには8ビット精度、低重要度のものには4ビット精度を適用するハイブリッド量子化戦略を採用し、効率性を最適化すること。
• エントロピー、分布、および適応的ビット幅最適化を統合した一貫したトレーニングパイプラインを設計し、エンドツーエンドのQATフレームワークとして統合すること。

実験結果

リサーチクエスチョン

• RQ1 なぜPost-Training Quantization (PTQ) は8ビット未満で重み、活性化、KVキャッシュを同時に量子化する際に失敗するのか？
• RQ2 LLMの低ビット量子化における主な性能劣化要因、特にアテンション機構においては何か？
• RQ3 エントロピー最大化と分布整合性が、アテンションマップにおける量子化誤差を共同で低減できるか？
• RQ4 重要度に配慮した適応的アクティベーション量子化は、平均ビット幅を削減しながら精度を維持するのにどの程度効果的か？
• RQ5 重みと活性化の完全なQATは、Raspberry Pi やFPGAなどの実際のエッジデバイスで顕著な速度向上を達成できるか？

主な発見

• Raspberry Pi 5ではW4A4量子化で最大2.37倍の速度向上を達成し、FP16推論を著しく上回る。
• OnePlus 11スマートフォンでは1.91倍のオンデバイス速度向上を達成し、FPGAプラットフォームではW4A4設定で最大1.9倍の向上を達成した。
• 4ビットの重みと活性化を用いたSuperGLUEベンチマークで最先端の精度を達成し、PACT や LSQ などのQATベースラインを上回った。
• 25％のトークンに8ビット精度を適用する適応的量子化は、同等の非適応的5ビット量子化よりも高い精度を達成した。
• 分布ガイド付き損失により、FP16出力と量子化出力間のアテンションマップの乖離が低減され、精度損失の直接的緩和が実現した。
• エントロピーに基づくクエリおよびキー分布最適化により、特に低ビット環境下で量子化誤差が顕著に低減された。