Sigma：用于高效语言模型的查询、键和数值的微分重新缩放

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注意力通过根据它们对模型性能和效率指标的不同影响优化注意力机制中的查询（Q）、键（K）和值（V）组件，显著提升了Sigma的推理效率。具体来说，我们（1）进行了大量实验，证明了模型对压缩K和V组件的敏感性不同，从而导致了不同压缩的KV的发展，以及（2）提出了增强的Q以扩展Q头维度，从而增强了模型的表示能力，对推理速度的影响最小。严格的理论和实证分析表明，DiffQKV注意力显著提升了效率，在长上下文情况下推理速度比传统的分组查询注意力（GQA）提升了高达33.36%。我们从各种来源预训练了Sigma，包括我们精心收集的195亿系统领域数据和1万亿合成和重写数据。在一般领域中，Sigma的性能与其他最先进的模型相当。在系统领域中，我们引入了第一个全面的基准AIMicius，Sigma在所有任务中表现出色，明显优于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%.