Sigma : Rééchelonnement différentiel de la requête, de la clé et de la valeur pour des modèles de langage efficaces

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

Résumé

Nous présentons Sigma, un modèle de langage large efficace spécialisé pour le domaine système, renforcé par une architecture novatrice comprenant une attention DiffQKV, et pré-entraîné sur nos données du domaine système collectées méticuleusement. L'attention DiffQKV améliore significativement l'efficacité d'inférence de Sigma en optimisant différemment les composants de Requête (Q), Clé (K) et Valeur (V) dans le mécanisme d'attention, en fonction de leurs impacts variables sur les performances du modèle et les indicateurs d'efficacité. Plus précisément, nous (1) menons des expériences approfondies qui démontrent la sensibilité variable du modèle à la compression des composants K et V, conduisant au développement de KV compressés de manière différentielle, et (2) proposons un Q augmenté pour étendre la dimension de la tête Q, ce qui améliore la capacité de représentation du modèle avec des impacts minimes sur la vitesse d'inférence. Des analyses théoriques et empiriques rigoureuses révèlent que l'attention DiffQKV améliore significativement l'efficacité, atteignant jusqu'à une amélioration de 33,36% en vitesse d'inférence par rapport à l'attention de requête groupée conventionnelle (GQA) dans des scénarios à long contexte. Nous pré-entraînons Sigma sur 6T tokens provenant de diverses sources, y compris 19,5B de données du domaine système que nous collectons soigneusement et 1T de tokens de données synthétisées et réécrites. Dans des domaines généraux, Sigma atteint des performances comparables à d'autres modèles de pointe. Dans le domaine système, nous introduisons le premier banc d'essai complet AIMicius, où Sigma démontre des performances remarquables sur toutes les tâches, surpassant significativement GPT-4 avec une amélioration absolue allant jusqu'à 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%.