シグマ：効率的な言語モデルのためのクエリ、キー、値の微分再スケーリング

Sigma: Differential Rescaling of Query, Key and Value for Efficient Language Models

January 23, 2025

著者: Zhenghao Lin, Zihao Tang, Xiao Liu, Yeyun Gong, Yi Cheng, Qi Chen, Hang Li, Ying Xin, Ziyue Yang, Kailai Yang, Yu Yan, Xiao Liang, Shuai Lu, Yiming Huang, Zheheng Luo, Lei Qu, Xuan Feng, Yaoxiang Wang, Yuqing Xia, Feiyang Chen, Yuting Jiang, Yasen Hu, Hao Ni, Binyang Li, Guoshuai Zhao, Jui-Hao Chiang, Zhongxin Guo, Chen Lin, Kun Kuang, Wenjie Li, Yelong Shen, Jian Jiao, Peng Cheng, Mao Yang

cs.AI

要旨

Sigmaは、システム領域に特化した効率的な大規模言語モデルを紹介します。このモデルは、DiffQKVアテンションを含む革新的なアーキテクチャによって強化され、入念に収集されたシステム領域のデータで事前学習されています。DiffQKVアテンションは、アテンションメカニズム内のQuery（Q）、Key（K）、Value（V）の各要素を、モデルのパフォーマンスと効率指標への影響の違いに基づいて異なる方法で最適化することで、Sigmaの推論効率を大幅に向上させます。具体的には、(1) KとVの要素の圧縮に対するモデルの感度の違いを実証する実験を行い、異なる圧縮されたKVを開発し、(2) Qヘッドの次元を拡張する拡張Qを提案し、モデルの表現能力を向上させつつ推論速度への影響を最小限に抑えます。厳密な理論的および実証的分析により、DiffQKVアテンションが効率を大幅に向上させ、長いコンテキストのシナリオにおいて従来のグループ化されたクエリアテンション（GQA）に比べて推論速度を最大33.36%向上させることが示されます。Sigmaは、様々なソースからの6兆トークン、慎重に収集された195億のシステム領域データ、および1兆トークンの合成および書き換えられたデータで事前学習されています。一般領域では、Sigmaは他の最先端モデルと同等のパフォーマンスを達成します。システム領域では、Sigmaは全タスクで優れたパフォーマンスを発揮する初の包括的なベンチマークAIMiciusを導入し、GPT-4を最大52.5%向上させる絶対的な改善を達成します。

English

We introduce Sigma, an efficient large language model specialized for the system domain, empowered by a novel architecture including DiffQKV attention, and pre-trained on our meticulously collected system domain data. DiffQKV attention significantly enhances the inference efficiency of Sigma by optimizing the Query (Q), Key (K), and Value (V) components in the attention mechanism differentially, based on their varying impacts on the model performance and efficiency indicators. Specifically, we (1) conduct extensive experiments that demonstrate the model's varying sensitivity to the compression of K and V components, leading to the development of differentially compressed KV, and (2) propose augmented Q to expand the Q head dimension, which enhances the model's representation capacity with minimal impacts on the inference speed. Rigorous theoretical and empirical analyses reveal that DiffQKV attention significantly enhances efficiency, achieving up to a 33.36% improvement in inference speed over the conventional grouped-query attention (GQA) in long-context scenarios. We pre-train Sigma on 6T tokens from various sources, including 19.5B system domain data that we carefully collect and 1T tokens of synthesized and rewritten data. In general domains, Sigma achieves comparable performance to other state-of-arts models. In the system domain, we introduce the first comprehensive benchmark AIMicius, where Sigma demonstrates remarkable performance across all tasks, significantly outperforming GPT-4 with an absolute improvement up to 52.5%.