VQ-Logits: Compressing the Output Bottleneck of Large Language Models via Vector Quantized Logits

要約

大規模な言語モデル（LLM）は驚くべき成功を収めていますが、特に広範な出力の語彙により、重要な計算とメモリの課題に直面しています。
非表示状態を語彙サイズのロジットにマッピングする最終的な線形投影層は、多くの場合、モデルのパラメーターのかなりの部分と推論中の計算コストを構成します。
適応型ソフトマックスや階層ソフトマックスなどの既存の方法は、構造的な複雑さを導入します。
この論文では、VQ-Logitsを提案します。これは、ベクター量子化（VQ）を活用してLLM出力層のパラメーター数と計算負荷を大幅に削減する新しいアプローチを提案します。
VQ-Logitsは、大きなv * dmodel出力埋め込みマトリックスを、ベクターの埋め込みベクターの小さな共有コードブックに置き換えます（k << v）。 語彙の各トークンは、これらのKコードブックベクトルの1つにマッピングされます。 LLMは、このコンパクトコードブックのロジットを予測し、学習または事前に署名したマッピングを使用して、完全な語彙スペースに効率的に「散乱」します。 標準言語モデリングベンチマーク（例：Wikitext-103、C4）に関する広範な実験を通じて、VQ-Logitsは、出力層の最大99％のパラメーター削減とロジット計算で6倍のスピードアップを達成できることを実証します。 さらに、コードブックのサイズ、初期化、学習戦略に関する詳細なアブレーション研究を提供し、アプローチの堅牢性と有効性を紹介します。

要約(オリジナル)

Large Language Models (LLMs) have achieved remarkable success but face significant computational and memory challenges, particularly due to their extensive output vocabularies. The final linear projection layer, mapping hidden states to vocabulary-sized logits, often constitutes a substantial portion of the model’s parameters and computational cost during inference. Existing methods like adaptive softmax or hierarchical softmax introduce structural complexities. In this paper, we propose VQ-Logits, a novel approach that leverages Vector Quantization (VQ) to drastically reduce the parameter count and computational load of the LLM output layer. VQ-Logits replaces the large V * dmodel output embedding matrix with a small, shared codebook of K embedding vectors (K << V ). Each token in the vocabulary is mapped to one of these K codebook vectors. The LLM predicts logits over this compact codebook, which are then efficiently 'scattered' to the full vocabulary space using the learned or preassigned mapping. We demonstrate through extensive experiments on standard language modeling benchmarks (e.g., WikiText-103, C4) that VQ-Logits can achieve up to 99% parameter reduction in the output layer and 6x speedup in logit computation, with only a marginal 4% increase in perplexity compared to full softmax baselines. We further provide detailed ablation studies on codebook size, initialization, and learning strategies, showcasing the robustness and effectiveness of our approach.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google