通过联合优化架构与量化策略实现 LLM 压缩

# 联合优化架构与量化选择的LLM压缩方法
来源：https://arxiv.org/html/2606.04063
11institutetext：挪威北极大学（UiT）22institutetext：挪威奥斯陆大学
22email：\{hoang\.l\.la,phuong\.ha\.hoai\}@uit\.no
\{truongl,amirhost\}@ifi\.uio\.no###### 摘要

由于大语言模型（LLM）对内存和计算资源需求巨大，其部署面临重大挑战。虽然部分方法通过从头开发小型或微型语言模型来应对这一问题，但这些方法需要大量GPU训练资源。将预训练LLM压缩后部署至边缘设备是一种极具吸引力的替代方案。除剪枝和量化之外，神经架构搜索（NAS）也是一种有效的压缩手段，但现有NAS方法往往限制了搜索空间，并将架构搜索与量化解耦处理。我们提出了一种可微分NAS框架，能够探索完整搜索空间，并对LLM线性层的架构配置与混合精度量化进行联合优化。实验结果表明，该框架在精度-延迟权衡方面表现出色：在精度相当的条件下，我们的模型推理速度比顺序执行NAS再量化的基线方法最高快1.4×；或在等效延迟下，在七项推理任务上的平均精度最高提升6%。

## 1 引言

近年来，大语言模型（LLM）受到广泛关注，但其高计算量和大内存需求使得在笔记本电脑、智能手机等资源受限设备上的部署十分困难。云端LLM带来的隐私顾虑日益增加，推动了本地推理需求的增长，而内存需求仍是主要瓶颈。

解决LLM部署挑战的方法主要有两类：开发新型轻量语言模型，以及压缩现有预训练LLM。第一类方法是从头训练小型语言模型，例如参数量为10亿的TinyLlama\[31 (https://arxiv.org/html/2606.04063#bib.bib26)\]，需要在16块A100-40GB GPU上训练90天；类似地，参数量为27亿的Phi-2\[16 (https://arxiv.org/html/2606.04063#bib.bib27)\]需要在96块A100 GPU上训练14天。

第二类方法充分利用预训练LLM，避免了从头训练，大幅缩短了训练时间。在结构化剪枝\[1 (https://arxiv.org/html/2606.04063#bib.bib7)\]和量化\[10 (https://arxiv.org/html/2606.04063#bib.bib28)\]等技术之外，神经架构搜索（NAS）也已成为LLM压缩的重要方法。然而，当前面向LLM压缩的NAS应用面临诸多挑战，例如需要资源密集型超网训练\[6 (https://arxiv.org/html/2606.04063#bib.bib8),5 (https://arxiv.org/html/2606.04063#bib.bib5)\]，或在超网训练中仅更新少量候选子网络\[25 (https://arxiv.org/html/2606.04063#bib.bib3),19 (https://arxiv.org/html/2606.04063#bib.bib11)\]。

与仅探索有限搜索空间的已有方法\[25 (https://arxiv.org/html/2606.04063#bib.bib3),19 (https://arxiv.org/html/2606.04063#bib.bib11),5 (https://arxiv.org/html/2606.04063#bib.bib5)\]不同，我们提出的NAS框架通过松弛化处理，在完整定义的离散搜索空间上进行优化。通过在给定约束下直接优化架构参数——而非像\[25 (https://arxiv.org/html/2606.04063#bib.bib3)\]那样进行预筛选——我们的方法能够探索更广泛的搜索空间。

此外，与剪枝后统一应用量化的传统方法不同，我们的方法对架构配置和逐层非均匀量化策略进行联合优化，从而找到最优的剪枝-量化组合，与现有最优方法相比，可实现显著更低的内存占用和更高的精度。我们的主要贡献如下：

- •我们提出了一种新颖的差异化权重纠缠超网设计，结合约束差异化优化方法，用于高效压缩预训练LLM。与现有最优方法相比，我们的方法在精度和延迟方面均表现更优。
- •我们提出了首个统一NAS框架，能够同时优化LLM的模型架构和逐层量化精度，解决了长期以来将剪枝与量化视为独立步骤的局限性。通过联合优化方法得到的模型，推理速度比顺序执行NAS再量化流程所得模型最高快1.4×。
- •我们开发了一种新颖的向量化实现方案，大幅加速了LLM权重纠缠超网的训练，与原始方法\[24 (https://arxiv.org/html/2606.04063#bib.bib1)\]相比，训练时间最高缩短4×。

## 2 相关工作与我们的进展

### 2.1 权重纠缠NAS

Sukthanker等人\[24 (https://arxiv.org/html/2606.04063#bib.bib1)\]在TangleNAS中提出了权重叠加技术，将所有可能的权重矩阵配置整合为单一加权表示，并为每种配置分配一个可学习的重要性标量。这使得训练与架构搜索可在单阶段内同步进行，有望降低搜索成本。然而，将TangleNAS直接应用于大语言模型（LLM）面临重大挑战。

- •TangleNAS\[24 (https://arxiv.org/html/2606.04063#bib.bib1)\]中的混合操作方式不适用于压缩LLM的深度维度。它将深度削减限制为仅丢弃末尾连续的若干块，而预训练基础模型中各个块的重要性差异显著，移除较为关键的块可能导致精度大幅下降。
- •TangleNAS中的权重纠缠机制未针对GPU效率进行优化，使其在LLM所需规模下进行超网训练时不切实际（详见第3.3节 (https://arxiv.org/html/2606.04063#S3.SS3)）。

**我们的贡献：** 我们提出了一种专门针对大语言模型压缩的高效权重纠缠超网设计。我们的框架通过在各层引入多样化量化精度选项，大幅扩展了搜索空间。针对TangleNAS的第一个局限性，我们提出了感知重要性的深度剪枝方法，实现更灵活、更有效的深度压缩（详见第3.2.2节 (https://arxiv.org/html/2606.04063#S3.SS2.SSS2)）。针对第二个局限性，我们开发了软件层面的优化方案，大幅加速超网训练，使得仅用单块NVIDIA A100 80GB GPU即可在合理时间内完成LLM压缩。

### 2.2 面向LLM压缩的神经架构搜索技术

Transformer的一个显著特性是排列等变性，即可以对嵌入特征、MLP中间特征和注意力头进行重排，而不会显著影响模型精度\[27 (https://arxiv.org/html/2606.04063#bib.bib9)\]。利用这一特性，常见的预处理步骤是按重要性对预训练LLM的各组件进行排序\[25 (https://arxiv.org/html/2606.04063#bib.bib3)\]。在选择子网络时，可直接从超网中选取前若干个神经元或注意力头。这一预处理方法在相关工作\[5 (https://arxiv.org/html/2606.04063#bib.bib5),6 (https://arxiv.org/html/2606.04063#bib.bib8),25 (https://arxiv.org/html/2606.04063#bib.bib3)\]中被广泛采用，本研究也沿用了这一做法。

LoNAS\[19 (https://arxiv.org/html/2606.04063#bib.bib11)\]和subnet-selection\[25 (https://arxiv.org/html/2606.04063#bib.bib3)\]均采用两阶段神经架构搜索（NAS）方法，在预训练大语言模型中寻找最优子架构。第一阶段使用LoRA\[14 (https://arxiv.org/html/2606.04063#bib.bib53)\]训练超网；第二阶段采用多目标搜索，以超网作为子网络的精度估计器，寻找在精度和性能指标（如延迟、能效）方面最优的子架构。

与LoNAS\[19 (https://arxiv.org/html/2606.04063#bib.bib11)\]不同，subnet-selection方法\[25 (https://arxiv.org/html/2606.04063#bib.bib3)\]在NAS过程之前加入了预筛选步骤。该方法按重要性对模型的特征和块进行排序，并在采样子网络时选取排名靠前的神经元、注意力头或块。研究者发现，LoNAS中使用的随机采样会在搜索空间中引入偏差——较小的子架构比较大的子架构更新频率更高，从而增加了超网训练的难度。为解决这一问题并优化搜索空间，他们引入了基于网格的采样方法，将搜索空间划分为$K$个分区，并从每个分区中选取最优候选。在超网微调阶段，随机选取$k\ll K$个子网络，通过知识蒸馏将这些子网络与最大（原始）网络一同训练。然而，这种启发式方法可能基于初始选择标准引入较强偏差，从一开始就可能导致次优架构。

**我们的贡献：** 已有技术主要关注频繁更新的子网络，而忽略了绝大多数很少甚至从未被更新的子网络，这可能导致错过真正的最优解。在本工作中，我们提出了一种新颖的差异化神经架构搜索（NAS）方法，专门用于压缩预训练大语言模型。通过利用权重纠缠风格的超网，我们的差异化超网无需从搜索空间中随机采样架构，从而避免了LoNAS中架构分布的偏斜问题。同时，我们的方法不依赖任何启发式规则来预筛选架构（如subnet-selection所做的那样）。相反，我们的方法在微调过程中探索所有可能的候选子网络，并逐步收敛至最优结构。此外，与早期不支持量化或需要顺序应用压缩与量化技术的方法不同，我们的方法能够同步优化各层的架构参数和量化精度。

## 3 方法

### 3.1 约束差异化NAS

#### 3.1.1 问题建模

我们将大语言模型（LLM）的压缩问题建模为一个约束优化问题。我们的搜索空间$S$由$\zeta\in S$参数化，以完全可微分的方式（类似于\[24 (https://arxiv.org/html/2606.04063#bib.bib1)\]）控制架构结构。采样的候选网络记为$\hat{\zeta}\sim P_{\zeta}(S)$，其中$P_{\zeta}(S)$是搜索空间$S$上由$\zeta$参数化的概率分布。$\mathcal{L}_{\text{train}}$和$\mathcal{L}_{\text{val}}$分别表示训练损失和验证损失。因此，LLM压缩任务可建模为如下双层约束优化问题：

$$\zeta^{*}=\quad\min_{\zeta}\mathcal{L}_{\text{val}}(w^{*},\hat{\zeta})\tag{1a}$$

$$\text{s.t.}\quad w^{*}=\arg\min_{w}\mathcal{L}_{\text{train}}(w,\hat{\zeta}),\quad B_{\text{min}}<F_{\text{params}}(\zeta_{\text{discrete}})<B_{\text{max}}\tag{1b}$$

其中$w$为预训练模型的权重。$F_{\text{params}}(\hat{\zeta})$表示NAS过程中离散化步骤所得最优神经架构的参数量。该离散化步骤需对架构参数$\zeta$应用$\arg\max$函数以选定最终架构，得到$\zeta_{\text{discrete}}=\arg\max(\zeta)$。由于$\arg\max$操作的引入，$F_{\text{params}}$不可微分。一种近似$F_{\text{params}}(\zeta_{\text{discrete}})$的简单方法是通过计算期望值$\mathbf{E}_{\hat{\zeta}\sim P_{\zeta}(S)}[F_{\text{params}}(\hat{\zeta})]$对硬约束进行松弛化处理。

#### 3.1.2 约束优化

我们可以通过为约束1b (https://arxiv.org/html/2606.04063#S3.E1.2)添加正则化项，将约束优化问题转化为无约束优化问题。具体地，约束1b (https://arxiv.org/html/2606.04063#S3.E1.2)可通过一对ReLU函数形式化表示如下：

$$F_{\text{constraint}}=\;ReLU\!\left(\mathbf{E}_{\hat{\zeta}\sim P_{\zeta}(S)}[F_{\text{params}}(\hat{\zeta})]-B_{\text{max}}\right)+ReLU\!\left(B_{\text{min}}-\mathbf{E}_{\hat{\zeta}\sim P_{\zeta}(S)}[F_{\text{params}}(\hat{\zeta})]\right)\tag{2}$$

公式1a (https://arxiv.org/html/2606.04063#S3.E1.1)中的损失项可改写为：

$$\mathcal{L}_{\text{train}}(w,\hat{\zeta})=\;\mathbf{E}_{\hat{\zeta}\sim P_{\zeta}(S)}[\mathcal{L}_{\text{CE}}(w,\hat{\zeta})]\tag{3}$$

$$\mathcal{L}_{\text{val}}(w^{*},\zeta)=\;\mathbf{E}_{\hat{\zeta}\sim P_{\zeta}(S)}[\mathcal{L}_{\text{CE}}(w^{*},\hat{\zeta})+\eta F_{\text{latency}}(\hat{\zeta})]+\lambda F_{\text{constraint}}$$

其中$L_{\text{CE}}$为交叉熵损失，$F_{\text{latency}}(\hat{\zeta})$为采样架构的预期推理延迟。该延迟项由各选项延迟的概率加权平均计算得出，各选项延迟值从预先计算的查找表中获取。$\eta$为超参数，用于控制验证损失$L_{\text{CE}}$与推理延迟$F_{\text{latency}}$之间的权衡；$\lambda$为超参数，用于控制正则化项$F_{\text{constraint}}$的强度。值得注意的是，$F_{\text{latency}}$可替换为其他用户自定义指标，例如能耗或内存占用。

#### 3.1.3 超网微调过程中的剪枝

假设共有$D$个可配置的架构维度，如块的数量、神经元数量或每块的注意力头数量。对于每个架构维度，共有$C$种不同选择。记$p^{d}_{c}$为第$d$个架构维度中第$c$个选择的概率。架构熵定义为：

$$H=-\frac{1}{D}\sum_{d=1}^{D}\sum_{c=1}^{C}p^{d}_{c}\log p^{d}_{c}.\tag{4}$$

当$H<\epsilon$时，表明已收敛至单一子架构，此时剪除所有冗余分支，仅保留最优子架构以继续微调。

#### 3.1.4 知识蒸馏

在剪枝步骤之后，当超网收敛至最优子网络（即$H<\epsilon$）时，我们通过知识蒸馏方法继续对其进行微调，以最大子网络（原始模型）作为教师模型，以t