DeepMind发布Matryoshka（套娃）量化：利用嵌套表示实现多精度LLM的低比特深度学习

本文将介绍 Google DeepMind 提出的 Matryoshka 量化技术，该技术通过从单个大型语言模型 (LLM) 实现多精度模型部署，从而革新深度学习。我们将深入研究这项创新技术如何提高 LLM 的效率和准确性。

随着深度学习模型在规模和复杂度上持续增长，效率和灵活性变得至关重要。量化作为一种成熟的技术，通过降低权重和激活的比特精度来减少模型大小和推理时间。传统的量化方法通常需要为每个精度级别训练单独的模型，导致资源效率低下和维护成本增加。Matryoshka 量化 (MatQuant) 通过训练一个能够满足多个精度要求的统一模型来克服这些挑战，从而简化部署和模型管理流程。

Matryoshka 原理

术语 "Matryoshka" 源自俄罗斯套娃，其核心思想是一个对象可以嵌套在另一个对象内部。在量化上下文中，该原理体现为：

• 嵌套表示： 较高比特的整数（例如 int8）自然包含较低比特的表示（例如 int4、int2）。这种固有的嵌套是 MatQuant 的基础。
• 统一模型训练： MatQuant 并非为每个精度训练单独的模型，而是训练一个统一的模型，其中 int8 权重的最高有效位 (MSB) 可以直接提取以获得 int4 或 int2 模型。
• 无缝精度转换： 这种嵌套结构允许在无需额外训练的情况下提取中间精度（如 int6 或 int3），从而根据硬件能力和资源可用性提供灵活的精度选项。

LLM 中的 Matryoshka 量化：原理与优势

什么是 Matryoshka 量化？

Matryoshka 量化 (MatQuant) 是一种多尺度量化训练方法，旨在跨不同的整数精度优化单个神经网络。与将每个精度级别视为独立的优化问题不同，MatQuant 利用了整数固有的嵌套结构。具体实现方式如下：

• 嵌套权重切片： 模型的权重以全精度 int8 格式存储。通过提取最高有效位，可以直接从同一模型导出较低精度的版本（int4、int2）。
• 联合优化： 训练过程同时最小化跨多个精度的量化损失。这通过结合复合损失函数来实现，该函数考虑了 int8、int4 和 int2 级别的误差，并通过重新加权因子进行平衡。
• 协同蒸馏： 在训练期间，高精度 (int8) 输出作为低精度输出的指导 "教师"。这种协同蒸馏通过确保低精度模型继承高精度表示的鲁棒性来增强其保真度。

在大型语言模型 (LLM) 中的应用

大型语言模型（如自然语言处理 (NLP) 中使用的模型）对计算资源的需求尤为显著。当应用于 LLM 时，MatQuant 具有以下优势：

• 资源效率： 使用 MatQuant 训练的 LLM 可以在 int8、int4 和 int2 精度之间动态切换。这种灵活性允许在高性能服务器和资源受限的设备上进行高效推理。
• 提高的低精度性能： 低精度模型（尤其是 int2）在独立训练时通常会遭受显著的精度下降。与标准方法相比，MatQuant 的联合优化和协同蒸馏过程已显示出可将 int2 精度提高高达 10% 的潜力，从而确保即使是最节省资源的模型版本也能保持高性能。
• 插值和灵活性： 除了目标精度之外，MatQuant 还可以通过简单的切片提取中间模型（例如 int6 或 int3），从而在精度和计算成本之间实现灵活的权衡。
• 逐层精度定制： MatQuant 支持混合精度策略，其中 LLM 的不同层可以以不同的精度运行。这种逐层混合搭配方法允许关键层（例如处理注意力机制的层）以更高的精度运行，而不太关键的层可以使用较低的精度，从而优化整体性能，同时尽量减少精度损失。

MatQuant 如何改进 LLM

增强的低精度准确性：

通过联合优化多个精度并采用协同蒸馏，MatQuant 有助于减轻在低精度模型中常见的精度下降问题。这种改进对于 int2 模型尤其重要，因为它们可以从训练期间学习到的共享知识和精细的权重分布中受益。

可扩展性和灵活性：

MatQuant 简化了 LLM 的部署流程。无需为每个精度维护单独的模型，使用 MatQuant 训练的单个模型可以动态调整其精度。这种适应性对于硬件约束或计算负载随时间变化的环境至关重要。

优化的推理和能源效率：

通过减少操作的位宽，经过 MatQuant 训练的模型需要更少的计算资源，从而缩短了推理时间并降低了能耗。这种效率使其非常适合基于云的服务和边缘设备，在这些设备中，功率和内存资源受到限制。

跨各种任务的鲁棒性：

实验结果表明，经过 MatQuant 训练的模型在各种下游任务上表现出强大的性能。无论是语言理解还是生成，MatQuant 提供的多精度功能都确保 LLM 在各种应用中保持高水平的性能。

实验结果与性能分析

MatQuant 在 Gemma-2（2B 和 9B 参数） 和 Mistral 7B 等模型上进行了全面的评估，并与量化感知训练 (QAT) 和 OmniQuant 等传统量化方法进行了比较。

性能指标

• 精度提升： 与传统的量化方法相比，MatQuant 显著提高了 int2 的性能，在下游任务上的精度提高了高达 10%。
• 模型稳定性： 当应用于前馈网络 (FFN) 和注意力参数时，MatQuant 可增强模型稳定性，尤其是在低精度状态下。
• 精度灵活性： 除了提供 int8、int4 和 int2 模型之外，MatQuant 还允许创建中间精度模型（如 int6 和 int3），而无需额外的训练，从而在精度和计算成本之间实现灵活的权衡。

MatQuant 对模型量化的影响

Gemma-2 9B 上的多比特量化框架和性能

OmniQuant 多比特量化框架：

• 该图展示了 OmniQuant 如何将多个量化级别（2 比特、4 比特和 8 比特）组合成一个统一的损失函数 LMatQuant\mathcal{L}_{\text{MatQuant}}。
• 每个比特宽度处理模型权重的特定部分，并且它们的组合输出被聚合以优化整体量化性能。

Gemma-2 9B 上的任务平均性能：

• 折线图显示了基于每个 FFN（前馈网络）参数的有效比特数的不同量化技术下 Gemma-2 9B 的任务性能。
• MatQuant（蓝色三角形）优于其他方法（基线、MinMax 和切片 int8），即使在低比特宽度下也能保持性能。
• MatQuant-Interp.（蓝色虚线）显示了插值如何进一步提高性能。

不同比特宽度（int8、int4、int2）的权重分布：

• 右上角 (int8)： 显示了基线（红色）和 MatQuant（蓝色）的量化权重分布。MatQuant 提供了更平衡的分布。
• 左下角 (int4)： MatQuant（蓝色）比基线（红色）在量化桶中实现了更均匀的分布。
• 右下角 (int2)： 尽管比特范围有限，但 MatQuant（蓝色）实现了对量化桶的更好利用，从而提高了效率。

该图突出显示了 OmniQuant 如何提高 Gemma-2 9B 模型在多个比特宽度上的量化效率和性能。

Gemma-2 9B 在不同量化方法下的任务性能

任务平均值与每个 FFN 参数的有效比特数的比较

该图评估了 Gemma-2 9B 在不同量化技术下的性能，并根据每个 FFN（前馈网络）参数的有效比特数绘制。

• X 轴（每个 FFN 参数的有效比特数）： 表示用于量化的比特宽度。较低的比特表示更积极的压缩。
• Y 轴（任务平均值）： 衡量模型在任务基准套件上的性能。较高的分数表示更好的性能。

量化方法：

• MatQuant（蓝色三角形）： 即使在低比特宽度下也能表现良好，并在较高比特宽度下接近基线饱和。
• Mix'n'Match（蓝色星形）： 由于混合量化策略，在中间范围比特宽度（3-5 比特）下显示出比 MatQuant 稍好的性能。
• MatQuant-Interp.（蓝色圆形）： 提供与 MatQuant 相似的一致性能，但在比特范围内略微平滑，受益于插值技术。
• 基线（红色三角形）： 在 2 比特时表现出最低的性能，并在 8 比特时与其他方法达到同等水平。

主要见解：

• 低比特效率： MatQuant 在 2 比特时显著优于基线，实现了近 10 个点的改进。
• 平台效应： 所有方法在超过 4 个有效比特时都在任务平均值上收敛于 75 左右，表明使用更多比特的回报递减。
• Mix'n'Match 优势： 在中间范围比特宽度上提供比 MatQuant 略微的优势，表明其在平衡压缩和性能方面的优势。

该图突出显示了 MatQuant 的效率和 Mix'n'Match 在 Gemma-2 9B 中间范围比特量化方面的适应性。

带有 OmniQuant 结果的 MatQuant

使用基于 OmniQuant 的训练比较 Gemma-2 和 Mistral 7B 模型的准确性和对数困惑度指标。

将 MatQuant 与 Gemma-2 (2B, 9B) 和 Mistral 7B 模型上的标准 OmniQuant 进行比较。

强调 MatQuant 优于基线方法，尤其是在 int2 精度方面，具有：

• Gemma-2 (2B) 上的准确度提高了 +4.37%
• Gemma-2 (9B) 上的准确度提高了 +8.01%
• Mistral 7B 上的准确度提高了 +6.35%

证明了 MatQuant 生成的插值模型（int6、int3）与显式训练的模型具有竞争力，突显了 MatQuant 在多尺度表示方面的效率。

带有 QAT 结果的 MatQuant

显示了相同模型但在 QAT 训练下的结果，表明 MatQuant 始终提高低精度性能：

• Gemma-2 (2B) 上 +4.69% 的准确度增益 (int2)
• Gemma-2 (9B) 上 +6.30% 的增益
• Mistral 7B 上 +6.34% 的增益

该表强调 经过 MatQuant 训练的 int2 模型甚至优于一些使用传统 QAT 训练的 int4 模型，突显了嵌套训练方法的有效性。

实际用例与应用

MatQuant 的灵活性和效率使其适用于各种场景：

边缘计算和物联网设备

• 挑战： 计算能力和内存有限。
• 解决方案： 部署从单个 MatQuant 训练的 LLM 派生的低精度模型（int2 或 int4），从而减少内存使用量和能耗。
• 示例： 智能相机、可穿戴设备和自主传感器受益于更快的推理和更低的功耗。

基于云的推理服务

• 挑战： 大规模数据中心中的高吞吐量需求。
• 解决方案： 使用来自 MatQuant 的 int8 精度模型可确保最大程度的准确性，同时优化资源利用率。
• 示例： 由 LLM 提供支持的虚拟助手和客户服务聊天机器人可以利用这些模型来高效地提供高质量的响应。

动态资源分配

• 挑战： 实时应用程序中资源可用性的波动。
• 解决方案： MatQuant 允许模型在精度之间动态切换 - 在资源充足时使用更高的精度，在资源不足时使用更低的精度。
• 示例： 流媒体服务和实时分析平台可以即时调整模型精度以保持最佳性能。

总结

Matryoshka 量化 是一种变革性的模型量化方法，它提供了一种统一的解决方案，可以在单个 LLM 中支持多个精度。通过利用整数表示固有的嵌套结构，MatQuant 不仅提高了低精度模型（尤其是 int2）的性能，还提供了适应各种硬件约束和动态资源条件的灵活性。

论文

https://avoid.overfit.cn/post/081ecf0860fa4611bd18ed4e78ca9212