Sigma: Redimensionamento Diferencial de Consulta, Chave e Valor para Modelos de Linguagem Eficientes

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

January 23, 2025

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

Resumo

Apresentamos Sigma, um modelo de linguagem grande eficiente especializado para o domínio do sistema, potencializado por uma arquitetura inovadora que inclui atenção DiffQKV e pré-treinado em nossos dados meticulosamente coletados no domínio do sistema. A atenção DiffQKV melhora significativamente a eficiência de inferência do Sigma otimizando os componentes de Consulta (Q), Chave (K) e Valor (V) no mecanismo de atenção de forma diferencial, com base em seus diferentes impactos no desempenho do modelo e nos indicadores de eficiência. Especificamente, realizamos experimentos extensivos que demonstram a sensibilidade variável do modelo à compressão dos componentes K e V, resultando no desenvolvimento de KV comprimidos de forma diferencial, e propomos um Q aumentado para expandir a dimensão da cabeça Q, o que melhora a capacidade de representação do modelo com impactos mínimos na velocidade de inferência. Análises teóricas e empíricas rigorosas revelam que a atenção DiffQKV melhora significativamente a eficiência, alcançando uma melhoria de até 33,36% na velocidade de inferência em comparação com a atenção de consulta agrupada convencional (GQA) em cenários de contexto longo. Pré-treinamos o Sigma em 6T tokens de várias fontes, incluindo 19,5B de dados no domínio do sistema que coletamos cuidadosamente e 1T de tokens de dados sintetizados e reescritos. Em domínios gerais, o Sigma alcança desempenho comparável a outros modelos de ponta. No domínio do sistema, introduzimos o primeiro benchmark abrangente AIMicius, onde o Sigma demonstra um desempenho notável em todas as tarefas, superando significativamente o GPT-4 com uma melhoria absoluta de até 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%.