【AI论文】Sigma：对查询、键和值进行差分缩放，以实现高效语言模型

摘要：我们推出了Sigma，这是一个专为系统领域设计的高效大型语言模型，其独特之处在于采用了包括DiffQKV注意力机制在内的新型架构，并在我们精心收集的系统领域数据上进行了预训练。DiffQKV注意力机制通过根据查询（Q）、键（K）和值（V）组件对模型性能和效率指标的不同影响，对它们进行差异化优化，从而显著提升了Sigma的推理效率。具体来说，我们（1）进行了大量实验，展示了模型对K和V组件压缩的不同敏感性，进而开发出了差异化压缩的KV方法；（2）提出了增强型Q，通过扩展Q头的维度，在几乎不影响推理速度的情况下，增强了模型的表示能力。严格的理论和实证分析表明，DiffQKV注意力机制显著提高了效率，在长上下文场景中，与传统分组查询注意力（GQA）相比，推理速度提升了高达33.36%。我们在6T个token上对Sigma进行了预训练，这些数据来自多个来源，包括我们精心收集的195亿个系统领域数据，以及1T个合成和重写数据。在通用领域，Sigma的表现与其他最先进的模型相当。在系统领域，我们引入了首个全面的基准测试AIMicius，Sigma在所有任务中均表现出色，显著优于GPT-4，绝对提升幅度高达52.5%。Huggingface链接：Paper page，论文链接：2501.13629

1. 引言

近年来，大型语言模型（LLMs）在各个领域取得了显著进展，展现了强大的性能。然而，随着模型规模的扩大，其计算复杂度和内存需求也随之增加，给实际应用带来了挑战。特别是在系统领域，即利用AI模型自动验证、评估、诊断和优化AI基础设施（如硬件、配置、云服务、数据库和工作负载）的领域，尽管具有巨大潜力，但尚未得到足够的重视。本文介绍了Sigma，一个专为系统领域设计的高效大型语言模型，通过采用包括DiffQKV注意力机制在内的新型架构，显著提高了推理效率。

2. Sigma模型概述

2.1 DiffQKV注意力机制

DiffQKV注意力机制是Sigma模型的核心创新点。在标准的多头注意力机制（MHA）中，查询（Q）、键（K）和值（V）通常使用相同数量的头和相同维度的头。然而，DiffQKV注意力机制打破了这一传统做法，允许Q、K、V组件具有不同数量的头和不同维度的头。此外，在推理过程中，K和V缓存的检索策略也各不相同。

2.1.1 差分压缩的KV

实验发现，模型性能对V向量的压缩比K向量更为敏感。因此，DiffQKV注意力机制对K向量采用更激进的压缩算法，而对V向量则采用较轻的压缩形式。尽管V向量的压缩程度相对较低，但在推理过程中可以通过仅加载与最高注意力分数对应的V向量来进一步优化，从而在保持模型性能的同时大幅减少内存使用。

2.1.2 增强的Q

增强的Q涉及采用比KV头更高的维度来扩展Q头维度，从而增强模型的表示能力，同时对推理速度的影响最小。实验表明，向Q头组件引入额外参数可以有效提升模型性能。

2.2 Sigma模型架构

Sigma模型基于DiffQKV注意力机制构建，并采用了两种模型规模：15亿参数和100亿参数，分别称为Sigma-1.5B和Sigma-10B。在训练过程中，为了平衡模型性能和KV缓存成本，对Sigma-1.5B和Sigma-10B的K头维度不进行压缩，仅减少K头的数量。具体来说，K头设置为4，而V头的数量是Q头数量的一半，设置为16。对于Sigma-1.5B，Q头维度设置为3072，对于Sigma-10B，则设置为6144，相当于隐藏状态的1.5倍，以扩展Q的表示空间。

3. DiffQKV注意力机制的理论与实证分析

3.1 理论分析

KV缓存和注意力计算是注意力层中的两个关键组件。通过减少K头的数量，Sigma模型能够显著降低KV缓存操作的成本。与标准的分组查询注意力（GQA）设计相比，Sigma模型在KV缓存操作上的成本降低率理论上可达到37.5%。尽管注意力计算是计算密集型的，但由于KV缓存操作是I/O密集型的，因此KV缓存成本的降低对整体推理效率的提升具有显著影响。

3.2 实现

尽管减少K头数量在理论上能够显著提高推理效率，但在实际部署中却面临挑战。为了解决这个问题，本文提出了几种临时解决方案，并强调了需要更广泛的支持来部署DiffQKV。其中，包括KV缓存的加载和存储策略以及灵活的注意力计算方法。

3.3 实证分析

通过一系列实验，验证了Sigma模型在推理效率上的显著提升。实验结果表明，与标准模型相比，Sigma模型在内核执行时间（KET）和CUDA事件经过时间（CEET）上均表现出显著优势。特别是在处理长上下文场景时，Sigma模型能够实现高达33.36%的推理速度提升。

4. 系统领域预训练与AIM ICIUS基准

4.1 系统领域数据收集

为了训练Sigma-System模型，本文精心收集了系统领域数据。通过识别来自120多个系统相关网站的15个主要源类别，收集了约195亿个令牌的数据。这些数据涵盖了学术论文、技术博客、开发者论坛、Stack Overflow等多个来源，确保了系统领域知识的全面性和多样性。

4.2 AIM ICIUS基准

为了评估系统领域任务的性能，本文构建了AIM ICIUS基准，包括CMDGen、Infrawise、Optiflow和NL2KQL四个主要任务。这些任务分别评估了模型在GPU相关命令生成、基准测试结果检索、网络拓扑优化和基础设施问题分析等方面的能力。实验结果表明，Sigma-System模型在AIM ICIUS基准上的性能显著优于其他基线模型，包括GPT-4等最先进的模型。

5. 性能评估

5.1 预训练设置

Sigma模型的预训练数据包括通用领域数据和系统领域特定属性数据，总计约6万亿个令牌。预训练过程分为多个阶段，逐步引入不同领域的数据，并采用退火策略进行微调。实验结果表明，经过充分预训练的Sigma模型在通用领域任务上也表现出色，与其他最先进的模型相比具有可比的性能。

5.2 系统领域性能

在AIM ICIUS基准上的实验结果表明，Sigma-System模型在所有任务上均表现出色，显著优于其他基线模型。特别是在CMDGen和Optiflow任务上，Sigma-System模型的优势更为明显，这得益于其在系统领域数据上的充分预训练和针对系统任务的微调。

5.3 通用领域性能

除了系统领域任务外，本文还对Sigma模型在通用领域任务上的性能进行了评估。实验结果表明，Sigma模型在常识推理、阅读理解、文本理解、语言熟练度、一般知识、编码和数学问题解决等方面均表现出色，与其他最先进的模型相比具有可比的性能。

6. 相关工作

近年来，为了提高LLMs的推理效率，研究人员探索了多种方法。其中，针对KV缓存优化的研究尤为突出。然而，大多数先前的研究倾向于对K和V向量进行均匀压缩，并且很少考虑对Q的优化。相比之下，DiffQKV注意力机制通过差分缩放Q、K、V组件，为提高LLMs的推理效率提供了新的思路。

7. 结论与未来工作

本文介绍了Sigma模型，一个专为系统领域设计的高效大型语言模型。通过采用DiffQKV注意力机制，Sigma模型在推理效率上实现了显著提升。实验结果表明，Sigma模型在处理长上下文场景时能够显著优于其他基线模型。尽管Sigma模型已经取得了显著进展，但仍然存在许多改进空间。未来的工作将集中在进一步优化Sigma模型的架构、扩展AIM ICIUS基准的任务范围以及提高预训练数据的质量等方面。

8. 具体实验细节与案例分析

8.1 实验细节

在实验过程中，本文采用了多种基准和评估指标来全面评估Sigma模型的性能。例如，在CMDGen任务上，采用了CMD分数、输出分数、校准分数、精确匹配、成功率和准确率等多个指标来衡量模型的性能。此外，为了验证DiffQKV注意力机制的有效性，本文还进行了大量消融实验，通过逐步引入差分压缩的KV和增强的Q等策略，分析了它们对模型性能的影响。

8.2 案例分析

以CMDGen任务为例，本文详细分析了Sigma模型在处理GPU相关命令生成任务时的表现。实验结果表明，Sigma模型能够准确生成符合要求的GPU命令，并且在执行结果上与真实命令高度相似。此外，Sigma模型在处理不同平台和不同类型的GPU命令时也表现出良好的泛化能力。这些实验结果充分证明了Sigma模型在系统领域任务上的有效性和实用性。

9. Sigma模型的应用前景与挑战

9.1 应用前景

Sigma模型在系统领域具有广泛的应用前景。例如，在AI基础设施的自动化优化方面，Sigma模型可以自动诊断和优化硬件、配置、云服务、数据库和工作负载等方面的问题，从而提高AI系统的效率和性能。此外，Sigma模型还可以应用于自然语言处理、知识图谱构建、智能问答等多个领域，为人工智能技术的发展提供有力支持。

9.2 挑战

尽管Sigma模型已经取得了显著进展，但在实际应用中仍然面临许多挑战。例如，如何进一步提高模型的推理效率以满足实时性要求？如何扩展模型的应用范围以覆盖更多领域和任务？如何提高模型的鲁棒性和可解释性以增强用户的信任感？这些问题都是未来研究需要重点关注的方向。

综上所述，Sigma模型是一个专为系统领域设计的高效大型语言模型，通过采用DiffQKV注意力机制显著提高了推理效率。实验结果表明，Sigma模型在处理系统领域任务时表现出色，并且在通用领域任务上也具有可比的性能。未来的工作将集中在进一步优化Sigma模型的架构、扩展AIM ICIUS基准的任务范围以及提高预训练数据的质量等方面，以推动Sigma模型在实际应用中的广泛部署和应用。