Sigma: Differential Rescaling of Query, Key and Value for Efficient
  Language Models
Authors: name: Zhenghao Lin; name: Zihao Tang; name: Xiao Liu; name: Yeyun Gong; name: Yi Cheng; name: Qi Chen; name: Hang Li; name: Ying Xin; name: Ziyue Yang; name: Kailai Yang; name: Yu Yan; name: Xiao Liang; name: Shuai Lu; name: Yiming Huang; name: Zheheng Luo; name: Lei Qu; name: Xuan Feng; name: Yaoxiang Wang; name: Yuqing Xia; name: Feiyang Chen; name: Yuting Jiang; name: Yasen Hu; name: Hao Ni; name: Binyang Li; name: Guoshuai Zhao; name: Jui-Hao Chiang; name: Zhongxin Guo; name: Chen Lin; name: Kun Kuang; name: Wenjie Li; name: Yelong Shen; name: Jian Jiao; name: Peng Cheng; name: Mao Yang
Published: 2025-01-23 12:58:14

Abstract: 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%.