Feature-based Low-Rank Compression of Large Language Models via Bayesian Optimization
Background and Motivation
随着现有的模型越来越大,当部署LLMs时,有必要考虑efficiency and performance. Such as pruning, quantization and knowledge distillation. 这些方法中unstructured pruning and quantization可以减少一半甚至更多的存储空间,但是他们需要特定的GPU内核函数来实现这样的加速过程。相反, structured pruning 可以生成不依赖专门的硬件的轻量级模型,但这种方法的性能会远远落后于原始模型。LRC是另一种方法,将权重矩阵分解为两个稠密low-rank矩阵乘积,舍弃不重要的参数。
LRC的关键在于:low-rank decomposition methods and low-rank dimension allocation. 现有的LRC可以分为两类:weight-based and feature-based decomposition. 前者最小化截断SVD或者权重SVD产生的权重矩阵的重构误差。大多数基于 Transformer 的语言模型的权重具有高秩甚至接近满秩的典型特征;因此,直接分解可能会导致显着的误差。相反,模型的特征通常表现出低秩特征。因此,更多的工作集中在feature-based decomposition来最小化features的重构误差。另一方面,根据目标压缩率给不同的权重矩阵分配合适的低秩维度也可以减少模型整体性能的下降,因为它们对低秩压缩表现出不同的敏感度。
LRC的challenges:
- hard to maintain LLMs generality while achieving feature-based low-rank compression。这是因为 LLM 的特征空间维度极高,使得特征分布更加复杂,异常特征的存在可能会干扰准确的分布估计。因此,我们使用池化协方差矩阵代替样本协方差矩阵,这样可以更准确地估计特征分布.
- manual design for low-rank dimension allocation很达到理想的效果,且需巨大的搜索空间,grid 搜索需要相当长的时间。经过实验验证,不同类型参数的low-rank sensitivity存在显著的差异。参数分为几组,让每组共享相同的低秩维度。这种方法有效地缩小了搜索空间,此外,我们利用样本高效的贝叶斯优化来确定最佳低秩分配。
Main Contributions
- 我们分析了 LLM 在低秩压缩方面面临的挑战,并通过实证研究证明以 LLaMA 为代表的 LLM 对各种参数的低秩压缩表现出截然不同的敏感性。
- 我们提出了一种基于贝叶斯优化的新型特征低秩压缩 (Bolaco)。
- 大量实验表明,我们的 Bolaco 优于 LLM 中现有的强结构化剪枝和 LRC 方法。
Theoretical Framework/Algorithm
weight-based factorization is one naive method. For a linear layer $W\in \mathbb{R}^{d_2\times d_1}$, the objective can be formulated as:
如果$r<d_1d_2/(d_1+d_2)$,这个分解就可以减少总体的参数量。但是在大量的研究中发现,对权重矩阵进行直接的截断SVD会造成巨大的截断误差,这些weights通常有high rank。相反的是PLMs的representation space展现出了low-rank property。因此,考虑feature-based decomposition
Experimental Design and Results
Comparative Analysis
Discussion and Limitations
Conclusion
在论文的3.2章节“基于贝叶斯优化的低秩分配”中,具体通过以下步骤确定rank(秩):
低秩压缩率分配的优化目标:
- 定义了一个优化目标函数 ( H(\lambda) ),该函数评估压缩模型在特定低秩压缩率 ( \lambda ) 下的性能。这个优化问题的目标是在总压缩率 ( \rho ) 的约束条件下,最小化 ( H(\lambda) )。
贝叶斯优化的应用:
- 使用贝叶斯优化来寻找最优的低秩压缩率分配 ( \lambda )。贝叶斯优化通过一个随机代理模型(通常是高斯过程)来估计目标函数 ( H(\lambda) ),并根据每次搜索步骤的结果更新对 ( H(\lambda) ) 的后验估计。
- 给定前 ( t-1 ) 步的搜索结果 (\{ \lambda_1, \lambda_2, …, \lambda_{t-1} \}) 及其评价 ( H_{t-1} ),更新代理模型的均值 ( \mu(\lambda) ) 和方差 ( \sigma^2(\lambda) )。
获取函数(Acquisition Function):
- 使用期望改进(Expected Improvement, EI)作为获取函数来决定下一步的低秩压缩率分配状态。获取函数通过期望改进值 ( \alpha(\lambda) ) 来选择具有最大期望改进的点,从而引导搜索过程。
- 公式为:[
\alpha(\lambda) = \mathbb{E}_{H(\lambda)} \left[ \max \{0, H’ - H(\lambda) \} \right]
]
其中 ( H’ ) 为目前观察到的最小值。贝叶斯优化选择具有最大期望改进的点来进一步探索。
最终低秩确定:
- 通过贝叶斯优化过程获得的最优压缩率分配 ( \lambda^* ),计算最终的低秩。为了充分利用GPU矩阵乘法的加速效果,将低秩维度四舍五入为八的倍数:
- 公式为:[
r_i = \left[ \frac{(1-\lambda_i) d_1 d_2}{d_1 + d_2} / 8 \right] \times 8
]
正则化项:
- 为了防止贝叶斯优化在较小的验证集上过拟合,在优化目标中引入了反向KL散度(Reverse KL Divergence, RKL)作为正则化项。这个正则化项量化了压缩模型与原始模型预测分布之间的差异,从而平滑了目标函数,使高斯过程代理模型更准确地逼近真实的黑箱目标函数。
通过上述步骤,结合贝叶斯优化和低秩特征分解的方法,有效地确定了不同层和参数的最优低秩分配,从而在保持模型性能的同时实现了高效的低秩压缩。