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

TL;DR

Sigma is a specialized large language model for the system domain that employs DiffQKV attention to improve inference efficiency by differentially compressing components, achieving up to 33.36% speedup and outperforming GPT-4 on domain-specific tasks.

Contribution

The paper introduces DiffQKV attention, a novel architecture that enhances efficiency by differentially compressing Q, K, and V components, and pre-trains Sigma on extensive system domain data.

Findings

DiffQKV attention improves inference speed by up to 33.36%.

Sigma outperforms GPT-4 on system domain benchmarks.

Differential compression of K and V components is effective.

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