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

TL;DR

Sigma introduces DiffQKV attention to differentiate how Q, KV are represented and cached in self-attention, enabling aggressive KV compression and augmented Q to boost efficiency. Theoretical analysis yields a KV-cache reduction of $37.5\%$ in asymptotic regimes, while empirical results show up to $33.36\%$ speedups in long-context inference and strong general-domain performance. Pre-trained on $6\,\text{T}$ tokens with $19.5\text{B}$ system-domain data and $1\text{T}$ synthetic data, Sigma achieves competitive results across standard benchmarks and delivers substantial improvements on the AIMicius system-domain benchmark (up to $52.5\%$ absolute). AIMicius covers 4 tasks—CMDGen, Infrawise, Optiflow, NL2KQL—demonstrating Sigma-System's capability to diagnose AI infrastructure, optimize topologies, and generate actionable commands, indicating strong practical impact for scalable, domain-specific LLMs.

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

Figures (13)

• Figure 1: KET comparison of FlexHeadFA between Standard model(Std) and Sigma.
• Figure 2: Comparison of total CEET cost between Standard model(Std) and Sigma. The gray dashed line indicates where the inference costs of both models are equal. As the output length increases, this intersection point moves progressively earlier.
• Figure 3: Overview of our proposed method for differential rescaling of QKV, compared alongside Multi-Head Attention (MHA), Multi-Query Attention (MQA), and Grouped Query Attention (GQA). Specifically, our method involves: (1) differentially compressed KV: applying more aggressive compression on the number of K heads and their dimensions than on the V components, which more significantly reduces the size of K cache. We can also optionally adopt selective V cache fetching for V compression; (2) augmented Q: adopting a higher dimension for the Q head compared to the KV heads.
• Figure 4: CEET comparison of augmented Q between Standard model(Std) and Sigma. From (a) to (f), the output length increases progressively from 2k to 64k tokens.
• Figure 5: CEET comparison of KV cache between Standard model(Std) and Sigma. From (a) to (f), the output length increases progressively from 2k to 64k tokens.