ROCKET:基于校准引导背包截断的快速优化模型压缩方法 ROCKET: Rapid Optimization via Calibration-guided Knapsack Enhanced Truncation for Efficient Model Compression
免训练的 LLM 压缩框架,通过单步稀疏字典分解和全局背包优化实现 SOTA 压缩效果
前置知识
截断 SVD(Truncated Singular Value Decomposition)
SVD 是将任意矩阵 $W \in \mathbb{R}^{d_1 \times d_2}$ 分解为 $W = U\Sigma V^T$ 的标准工具,其中 $U$ 和 $V$ 是正交矩阵,$\Sigma$ 是奇异值对角阵。截断 SVD 仅保留前 $k$ 个最大奇异值及其对应的向量,得到秩为 $k$ 的最优低秩近似 $W_k = U_k \Sigma_k V_k^T$。在 LLM 压缩中,它被用来将大权重矩阵分解为两个较小的矩阵相乘,从而减少参数量。但其核心局限在于:所有列被迫共享同一个低维子空间,无法捕捉权重矩阵内部的异构局部结构。
SVD 是 LLM 后训练压缩的最经典基线方法,理解它的局限性(刚性子空间约束)是理解 ROCKET 为何要引入字典学习的关键前提。
字典学习与稀疏编码(Dictionary Learning & Sparse Coding)
字典学习的核心思想是:用一组「原子」(basis vectors)构成字典 $D$,将每个数据向量表示为少量原子的稀疏线性组合 $x \approx D\alpha$,其中 $\alpha$ 是稀疏系数。与 SVD 的单一全局子空间不同,字典学习允许不同的向量选择不同的原子子集,形成「子空间并集」(union-of-subspaces)表示,从而捕捉数据内部的异构结构。传统字典学习通过 K-SVD 和 OMP(正交匹配追踪)等迭代算法交替优化字典和稀疏系数,计算代价高昂。
ROCKET 的核心贡献之一就是将字典学习从迭代优化改为单步闭式解,理解字典学习的表达优势和计算瓶颈是理解这一贡献的前提。
白化变换(Whitening Transform)
白化是一种数据预处理技术,通过线性变换使输入数据的协方差矩阵变为单位矩阵。给定校准数据 $X$,其 Gram 矩阵 $A = X^T X$ 的 Cholesky 分解为 $A = LL^T$,则白化输入为 $Y = XL^{-1}$,满足 $Y^T Y = I$。白化消除了特征之间的相关性,使得后续的低秩分解在几何上更优——因为正交的 $Y$ 列向量不会放大或缩放误差方向。在数据感知的 LLM 压缩中,白化将原始的「权重空间近似」问题转化为等价的「白化空间近似」问题,后者具有更干净的几何结构。
ROCKET 的整个分解流程都建立在白化权重空间 $W_L = LW$ 上进行,白化变换是连接校准数据和权重近似的桥梁。
多选背包问题(Multi-Choice Knapsack Problem, MCKP)
背包问题是组合优化中的经典问题:在总容量约束下选择物品以最大化价值。多选背包问题是其变体:物品被分成若干组,每组必须恰好选一个物品。在模型压缩的语境下,每一层是一「组」,该层的每种压缩配置(不同的秩和稀疏度组合)是一个「物品」,物品的「代价」是压缩后的参数量,「价值」是重建误差的负值。目标是在总参数预算约束下最小化所有层的总重建误差。MCKP 可通过动态规划在多项式时间内精确求解。
ROCKET 将逐层压缩配置的选择形式化为 MCKP,这是它实现全局最优参数分配的数学基础,也是区别于均匀压缩或启发式分配的核心创新。
激活感知压缩(Activation-Aware Compression)
传统权重压缩只关注权重矩阵本身的近似质量(如 Frobenius 范数误差),但实际上模型的输出是激活与权重的乘积 $XW$。激活感知压缩利用一小批校准数据 $X$ 来评估压缩对实际前向计算的影响:某些权重方向虽然范数小,但对应的激活方向幅度大,压缩它们会造成更大的输出误差。通过引入校准数据,压缩过程可以识别并保护这些「高敏感度」方向,从而在相同压缩比下保持更好的下游性能。
ROCKET 使用 256 条校准样本计算白化变换和重要性分数,其激活感知特性是它超越纯 SVD 方法的关键因素之一。
研究动机
当前大语言模型的后训练压缩面临一个根本性的效率-表达力两难困境。一方面,截断 SVD(如 SVD-LLM)是最高效的压缩方法——它将每个权重矩阵近似为两个较小矩阵的乘积,操作简单且支持标准密集线性代数运算。但 SVD 强制所有列共享同一个低维子空间,这种刚性约束在中高压缩比(30%–50%)下导致严重的性能退化:在 Qwen3-8B 上以 50% 压缩率应用 SVD-LLM,平均准确率从 70.5 骤降至 38.1,WikiText 困惑度从 12 飙升至 76。另一方面,基于字典学习的方法(如 CoSpaDi)通过子空间并集建模提供了更大的灵活性——允许权重矩阵的每一列独立选择不同的基向量子集。但这种灵活性的代价是巨大的计算开销:传统字典学习依赖 K-SVD 和 OMP 的迭代交替优化,在十亿参数级模型上需要数万秒的压缩时间(CoSpaDi 在 Llama3-1B 上需要 90081 秒),能耗高达 7.88 kWh,CO₂ 排放 0.782 kg,这使得大规模部署几乎不可行。
本文的目标是ROCKET 的目标是设计一种同时具备 SVD 的计算效率和字典学习的表达能力的训练无关压缩框架。具体而言,它追求三个可量化的目标:第一,在 20%–50% 压缩率范围内持续优于现有的 SVD、字典学习和预算分配方法;第二,压缩时间比 CoSpaDi 快两个数量级以上(目标:从数万秒降至千秒量级);第三,在 30% 压缩率下保留原始模型 90% 以上的性能,且无需任何微调或数据增强。此外,当施加轻量级恢复训练(仅使用 3000 万 token)时,压缩后的模型应能接近同等大小从头训练模型的性能。
与已有工作不同的是,ROCKET 抓住了一个被现有工作忽视的关键观察:权重矩阵的白化版本 $W_L = LW$ 的特征分解(EVD)天然等价于 SVD,但如果我们对分解后的系数矩阵施加结构化稀疏而非简单截断,就能在不增加计算复杂度的前提下实现字典学习级别的表达力。具体来说,现有字典学习方法(K-SVD + OMP)的瓶颈在于迭代优化——交替更新字典和稀疏系数,每次迭代都需要求解大规模稀疏编码问题。ROCKET 的核心洞察是:这个迭代过程可以被一个单步贪心算法替代——先通过 EVD 构建数据自适应的正交基,再根据融合了白化空间和原始空间信息的重要性分数进行结构化硬阈值稀疏化,最后通过闭式最小二乘更新字典矩阵。此外,在逐层压缩配置选择上,现有方法要么使用均匀分配(忽略层间差异),要么使用启发式梯度估计(需要训练),ROCKET 将其形式化为有约束的多选背包问题,通过动态规划在多项式时间内求解全局最优分配。
核心方法
ROCKET 的方法可以用一个图书馆整理的类比来理解:想象你要把一个巨大的图书馆(权重矩阵)缩减到更小的空间(压缩模型)。传统 SVD 方法相当于把所有书按一个固定分类体系排列——简洁但僵硬,很多书的细微差异被忽略了。字典学习方法则像请了一组图书馆员反复讨论每本书应该放在哪个书架上——精细但耗时。ROCKET 的做法更巧妙:先用一次快速扫描(EVD)确定书架的自然分组结构,然后根据每本书的「借阅热度」(激活-权重重要性)快速决定哪些书放在显眼位置(保留)、哪些收入仓库(稀疏化),最后微调书架布局(闭式最小二乘更新)以最小化总混乱度。同时,ROCKET 还有一个全局「预算管理员」(背包优化器),确保整个图书馆的总空间恰好符合目标,且没有哪个区域被过度压缩。技术路线上,整个流程分为两个阶段:(1)逐层压缩配置搜索——对每一层,枚举多种(秩, 稀疏度)配置并记录各自的参数量和重建误差;(2)全局最优分配——将逐层配置选择建模为多选背包问题,通过动态规划在总参数预算约束下最小化全局重建误差。
ROCKET 的核心创新在于用单步稀疏矩阵分解替代迭代字典学习,同时在重要性度量上融合了白化空间和原始空间的双重视角。具体来说,给定校准数据 $X$ 和白化权重 $W_L = LW$,ROCKET 计算 $W_L W_L^T$ 的前 $r$ 个特征向量 $B$(这等价于 $W_L$ 的 SVD 左奇异向量),得到系数矩阵 $C = B^T W_L$。与传统 SVD 不同,ROCKET 不是简单截断 $C$,而是对 $C$ 进行结构化稀疏化:为每个系数 $c_{ij}$ 定义融合重要性分数 $\text{imp}_{ij} = |c_{ij}| \cdot \|L^{-1}b_i\|_2^{\lambda}$,其中 $\lambda = 0.5$ 作为几何均值平衡白化空间的局部最优性($|c_{ij}|$)和原始空间的实际失真($\|L^{-1}b_i\|_2$)。这个双空间视角是关键:白化空间中两个相同大小的系数,由于对应的基向量被 $L^{-1}$ 缩放的程度不同,在原始空间中造成的重建误差可能差异巨大。稀疏化采用两阶段策略:先列级硬阈值保留每列前 $s$ 个最重要的系数,再全局重激活直到精确达到目标压缩比。最后,放弃正交性约束,通过闭式岭回归 $D_{\text{final}} = \arg\min_D \|W_L - DC_{\text{sparse}}\|_F^2 + \mu\|D\|_F^2$ 更新字典矩阵,获得更好的拟合而不增加推理开销。
方法步骤详情
ROCKET 的完整流程可分为以下几个步骤:第一步,计算白化变换。从 256 条校准样本 $X \in \mathbb{R}^{N \times d_1}$ 计算 Gram 矩阵 $A = X^T X$ 的 Cholesky 分解 $A = LL^T$。第二步,逐层配置搜索。对每个可压缩层,将其白化权重设为 $W_L = LW$,然后枚举预定义的(秩 r, 稀疏度 s)配置组合。对每个配置:(a)计算 $W_L W_L^T$ 的前 r 个特征向量得到正交基 B;(b)计算系数矩阵 $C = B^T W_L$;(c)计算融合重要性分数;(d)执行两阶段稀疏化——先过稀疏化到 $cr + \beta$($\beta = 5 \times 10^{-3}$),再全局重激活到精确目标比;(e)通过闭式最小二乘求解 $D_{\text{final}}$;(f)记录该配置的参数量、重建误差和稀疏度-截断比。第三步,全局最优分配。构建有约束的多选背包问题:每层是一个「组」,该层的每个配置是一个「选项」,目标是在总参数预算约束下最小化总重建误差,附加每层误差上限约束。通过自底向上的动态规划在 $O(LM\bar{B})$ 时间和 $O(\bar{B})$ 空间内精确求解。第四步,回溯恢复最优路径,将每层的两个因子 $U = L^{-1}D_{\text{final}}$ 和 $V = C_{\text{sparse}}$ 作为压缩后权重存储。
技术新颖性
ROCKET 的技术新颖性体现在三个层面。第一,算法层面:它将字典学习从迭代优化(K-SVD + OMP,通常需要数十次交替迭代)压缩为单步闭式操作——一次 EVD、一次硬阈值、一次最小二乘,没有反向传播、没有稀疏编码迭代、没有训练。这带来的不仅是速度提升(Llama3-1B 上压缩时间从 90081 秒降至 930 秒,快 96 倍),更是能耗和碳排放的两个数量级降低(从 7.88 kWh 降至 0.0765 kWh)。第二,重要性度量层面:融合白化空间和原始空间的双重要性分数是全新的设计,通过几何均值优雅地融合了两个视角。第三,优化层面:将逐层压缩配置选择形式化为有约束的多选背包问题并附加每层误差上限约束,既不是均匀分配,也不是可训练的梯度方法,而是纯粹的后训练组合优化。消融实验证明,仅用稀疏化就能将平均准确率从 37.6 提升到 45.4,再加背包分配进一步提升到 52.4。
实验结果
ROCKET 在多个维度的实验中展示了全面的优势。首先,在与最直接的竞争方法 SVD-LLM 和 CoSpaDi 的严格对比中(纯训练无关设置,无微调),ROCKET 在 Qwen3-8B 上从 20% 到 50% 四个压缩率上均显著领先:20% 压缩时平均准确率 68.7(SVD-LLM 62.5,CoSpaDi 65.7),30% 时 65.8(vs 55.8/60.0),40% 时 60.0(vs 45.4/50.3),50% 时 51.3(vs 38.1/42.0)。在更具挑战性的 Llama3.2-1B 上,优势更加明显——50% 压缩时 ROCKET 达到 31.9 平均准确率,而 SVD-LLM 和 CoSpaDi 分别仅为 25.4 和 25.8。困惑度方面,ROCKET 在 Qwen3-8B 50% 压缩时 WikiText 困惑度为 35,LAMBADA 困惑度为 24,远优于 SVD-LLM 的 76/88 和 CoSpaDi 的 59/41。在与其他预算分配方法的对比中,ROCKET 在 Qwen3-8B 的 20% 和 40% 压缩以及 LLaMA2-7B 的 40% 压缩三种配置下,均优于 ARS、Dobi-SVD 和 ARA。与深度/稀疏剪枝方法的对比同样令人印象深刻:在 Llama3.1-8B 上以 40% 压缩率,ROCKET(无量化)达到 0.56 平均准确率,优于 LLM-Pruner(0.46)、SliceGPT(0.46)、Bonsai(0.41)和 Wanda(0.39);加上量化后进一步提升至 0.65,甚至超过使用量化的 Dobi-SVD(0.63)。在 60% 压缩率下,ROCKET+量化达到 0.60,而 Dobi-SVD+量化仅为 0.52。跨模态实验进一步验证了通用性:在视觉语言模型 Qwen3-4B-VL 上以 20% 压缩,ROCKET 保持了 65.75 平均准确率(原始模型的 90% 以上)。在语音生成模型 VibeVoice 上,WER 仅从 0.148 微增到 0.149,UTMOS 从 3.52 降至 3.43。轻量恢复实验尤为亮眼:将 Qwen3-14B 压缩到 8B 参数(40% 压缩),仅用 3000 万 token 微调后,平均准确率从 63.56 恢复到 67.96,接近原始 Qwen3-8B 的 70.46,在 PIQA(78.51 vs 77.70)和 LAMBADA(65.55 vs 64.10)上甚至超越了从头训练的 8B 模型。环境效率方面,ROCKET 的压缩时间(930 秒 vs 90081 秒)、能耗(0.0765 kWh vs 7.88 kWh)和 CO2 排放(0.0337 kg vs 0.782 kg)均比 CoSpaDi 低两个数量级。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Qwen3-8B 零样本准确率 (20% 压缩) | 8 个基准平均准确率 | 68.7 | CoSpaDi 65.7 / SVD-LLM 62.5 | 较 CoSpaDi +3.0,较 SVD-LLM +6.2 |
| Qwen3-8B 零样本准确率 (50% 压缩) | 8 个基准平均准确率 | 51.3 | CoSpaDi 42.0 / SVD-LLM 38.1 | 较 CoSpaDi +9.3,较 SVD-LLM +13.2 |
| Llama3.2-1B 零样本准确率 (20% 压缩) | 8 个基准平均准确率 | 52.4 | CoSpaDi 42.7 / SVD-LLM 37.6 | 较 CoSpaDi +9.7,较 SVD-LLM +14.8 |
| Llama3.1-8B (40% 压缩, 含量化) | 5 个基准平均准确率 | 0.65 | Dobi-SVD 0.63 | +0.02 |
| Llama3.1-8B (60% 压缩, 含量化) | 5 个基准平均准确率 | 0.60 | Dobi-SVD 0.52 | +0.08 |
| VibeVoice 语音生成 (20% 压缩) | WER / UTMOS | 0.149 / 3.43 | 原始 0.148 / 3.52 | WER 仅增加 0.7%,UTMOS 下降 2.6% |
| Qwen3-14B→8B 恢复后 | 平均准确率 | 67.96 | 原始 Qwen3-8B 70.46 | 仅用 30M token 恢复到 96.4% 性能 |
| Qwen3-8B 推理吞吐 (30% 压缩) | tokens/s (batch=1, ctx=256) | 26.60 | CoSpaDi 25.76 / SVD-LLM 24.31 | 较 CoSpaDi +3.3%,较 SVD-LLM +9.4% |
局限与改进
ROCKET 的局限性主要体现在以下几个方面。首先,作者坦诚地指出,动态规划求解器虽然对于标准密集 Transformer 效率很高,但难以扩展到具有大量可压缩组件的架构,特别是现代 MoE(混合专家)模型——当每层有 128 个或更多专家时,压缩配置的组合空间会呈指数级增长,使得背包求解变得不可行。其次,恢复实验中的稀疏模式在压缩阶段就已固定且在微调过程中保持不变,这显然是次优的——联合学习自适应稀疏模式可能带来进一步改进。第三,ROCKET 当前仅在 Qwen、Llama 系列模型上进行了评估,未涉及其他主流架构如 Mistral、Gemma 等,其在不同注意力机制(如 GQA、MLA)下的表现尚不清楚。第四,虽然论文展示了在视觉语言和语音模型上的跨模态泛化,但评估规模较小(20% 压缩),在更激进的压缩比下跨模态表现如何仍待验证。第五,embedding 层和 lm_head 未被压缩,这意味着在小模型上实际压缩率可能与目标有偏差,且这部分参数的压缩潜力未被挖掘。最后,从我的观察来看,论文的高级基准测试(MATH、GPQA 等)结果显示压缩后的能力退化远比基础基准严重——Qwen3-8B 在 MATH 上从 52.57 降至 11.10(20% 压缩),这暗示压缩主要保留了「记忆型」能力而非「推理型」能力。
独立分析的弱点
从独立分析的角度,ROCKET 有几个值得关注的弱点。第一,重要性分数中的平衡参数 $\lambda = 0.5$(几何均值)是固定的超参数,但不同层、不同模型的最优 $\lambda$ 可能不同——例如,注意力层和 MLP 层在白化空间和原始空间的误差分布差异很大,统一使用 0.5 可能不是最优的。改进方向:可以将 $\lambda$ 也纳入背包问题的候选配置中,或通过小规模网格搜索逐层确定最优值。第二,两阶段稀疏化策略中 $\beta = 5 \times 10^{-3}$ 是固定值,这个「松弛量」的选择缺乏理论依据。在极端压缩率下,$\beta$ 过小可能导致重激活空间不足,过大则破坏了列级稀疏性带来的推理效率。改进方向:让 $\beta$ 随目标压缩率自适应调整。第三,ROCKET 的推理加速依赖于 MACKO 稀疏矩阵乘法内核,但在实际部署中,硬件特异性(不同的 GPU 架构、NPU 等)可能使这个内核无法直接使用,导致理论压缩收益无法转化为实际加速。改进方向:探索与更广泛支持的稀疏计算格式(如 N:M 稀疏、2:4 稀疏)的兼容性。第四,校准数据的选择对结果有影响——Alpaca(指令微调数据)比 RefinedWeb 好 1.4 个百分点(53.8 vs 52.4),但论文为了公平对比而统一使用 RefinedWeb。在实际应用中,应根据目标模型的使用场景选择校准数据。第五,论文声称「超过 90% 原始性能」是在 30% 压缩下的 Qwen3-8B 上成立,但这个说法不具普适性——在 Llama3.2-1B 上 30% 压缩后准确率仅为 45.9(原始 57.6),即 79.7% 性能。
未来方向
基于论文成果,有几个值得探索的未来方向。首先,作者明确提出将 ROCKET 扩展到 MoE 架构——这需要设计更高效的配置搜索策略,例如对同一层的所有专家使用共享配置、或利用专家间的结构相似性减少搜索空间。其次,作者提到在微调阶段联合学习自适应稀疏模式,而非固定压缩时确定的模式——这可以通过在微调时对每个系数引入可学习的二值掩码来实现。第三,ROCKET 的单步分解框架可以与量化结合,形成「分解+量化」的两级压缩流水线——论文中已初步验证了 ROCKET+量化的有效性,但未深入探索量化友好的分解策略。第四,当前的全局预算分配只优化了参数量约束,可以扩展为多目标优化,同时考虑 FLOPs、内存占用和延迟约束。第五,将背包分配从离散选择扩展为连续优化(通过松弛和投影),可能在更细粒度上实现更优的分配。最后,ROCKET 的快速压缩能力使其适合用于「压缩感知的模型选择」——在训练阶段就考虑最终的压缩需求,选择在压缩后表现最好的模型架构和超参数。
复现评估
论文明确声明将开源 ROCKET 的代码实现,这对复现非常有利。方法本身是完全训练无关的(纯后训练),不涉及随机性较大的训练过程,理论上复现一致性应该很高。算力需求方面:Llama3-1B 的压缩在单块 NVIDIA A100-SXM4-40GB 上仅需 930 秒(约 15 分钟),能耗 0.0765 kWh,这对大多数研究团队来说完全可承受。校准数据仅需 256 条样本,论文使用的是公开的 RefinedWeb 数据集。唯一的算力瓶颈可能是 Qwen3-14B 到 8B 的恢复实验,虽然只用 3000 万 token,但在 A100 上仍需要数小时的微调。消融实验提供了丰富的超参数参考:$\lambda = 0.5$、$\beta = 5 \times 10^{-3}$、$\alpha$ 自动确定、256 条校准样本等。总体而言,复现难度为中等偏易——核心算法(EVD + 硬阈值 + 最小二乘 + 动态规划)都有标准实现,但需要注意白化变换和重要性分数的数值稳定性。
论文图表