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%。我們在各種來源的6T標記上對Sigma進行了預訓練，包括我們精心收集的195億系統領域數據和1T標記的合成和重寫數據。在一般領域中，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%.