Sigma: Differentiële Herschaling van Query, Key en Value voor Efficiënte Taalmodellen

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

January 23, 2025

Auteurs: 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

Samenvatting

Wij introduceren Sigma, een efficiënt groot taalmodel gespecialiseerd voor het systeemdomein, versterkt door een nieuw architectuur met DiffQKV-aandacht, en vooraf getraind op onze zorgvuldig verzamelde systeemdomeingegevens. DiffQKV-aandacht verbetert aanzienlijk de inferentie-efficiëntie van Sigma door de Query (Q), Key (K) en Value (V) componenten in het aandachtsmechanisme differentieel te optimaliseren, gebaseerd op hun variërende impact op de modelprestaties en efficiëntie-indicatoren. Specifiek voeren we (1) uitgebreide experimenten uit die de variërende gevoeligheid van het model voor de compressie van K en V componenten aantonen, wat leidt tot de ontwikkeling van differentieel gecomprimeerde KV, en (2) stellen we augmented Q voor om de Q-hoofddimensie uit te breiden, wat de representatiecapaciteit van het model verbetert met minimale impact op de inferentiesnelheid. Grondige theoretische en empirische analyses tonen aan dat DiffQKV-aandacht de efficiëntie aanzienlijk verbetert, met een verbetering van maximaal 33,36% in inferentiesnelheid ten opzichte van de conventionele gegroepeerde-query-aandacht (GQA) in scenario's met lange context. We trainen Sigma vooraf op 6T tokens uit verschillende bronnen, waaronder 19,5B systeemdomeingegevens die we zorgvuldig verzamelen en 1T tokens van gesynthetiseerde en herschreven gegevens. In algemene domeinen behaalt Sigma vergelijkbare prestaties als andere state-of-the-art modellen. In het systeemdomein introduceren we de eerste uitgebreide benchmark AIMicius, waar Sigma opmerkelijke prestaties levert op alle taken, aanzienlijk beter presterend dan GPT-4 met een absolute verbetering tot 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%.