想象一个大型语言模型(LLM)是一座拥有数十亿本书籍(参数)的庞大且高度有序的图书馆。为了让这座图书馆在标准计算机上快速运行,你需要移除一些书籍。然而,你面临着一个棘手的两难困境:
- “混乱”的方法(非结构化稀疏性): 你从书架的任何地方随机扔掉书籍。这种方法能保持图书馆知识的极高准确性,因为你可以进行非常精细的选择,但这会造成混乱。一位图书管理员(计算机的 GPU)在寻找书籍时不得不四处跳跃,导致过程缓慢且低效。
- “僵化”的方法(2:4 稀疏性): 你决定遵循一条严格的规则:“在每四本书中,必须恰好移除两本。”这使得图书管理员的工作变得轻松快捷,因为模式是可预测的。然而,这条规则过于僵化。有时,你必须移除的那两本书实际上是最重要的,导致图书馆失去其智能和准确性。
引入 PATCH:一位“智能瓷砖”图书管理员
这篇论文介绍了一种名为 PATCH(用于混合稀疏性的可学习瓷砖级配置剪枝)的新方法。PATCH 不像是在“混乱”和“僵化”的方法之间做选择,而是扮演一位智能图书管理员的角色,将图书馆划分为瓷砖(小块、可管理的书架区域)。
以下是其工作原理,使用一个简单的类比:
- 瓷砖系统: 想象图书馆被划分为方形瓷砖,就像马赛克一样。
- 决策: 对于每一块瓷砖,PATCH 系统学习做出选择:
- 选项 A(稠密): 保持这块瓷砖完全装满书籍。这是针对图书馆中准确性至关重要的“关键”区域。
- 选项 B(2:4 稀疏): 对这块瓷砖应用严格的“每四本移除两本”规则。这是针对图书馆拥有额外、冗余书籍的区域,可以安全地移除以节省空间并加快速度。
- 学习过程: 系统不会猜测。它通过“训练”自己来确定哪些瓷砖应该保持满员,哪些应该稀疏。它学会保持重要部分稠密,冗余部分稀疏,同时遵循对硬件友好的规则。
为什么这很重要?
- 兼得鱼与熊掌: PATCH 弥合了差距。它既保持了图书馆的准确性(像“混乱”的方法),又以计算机能够快速读取的方式组织它(像“僵化”的方法)。
- 灵活的速度: 你可以告诉 PATCH:“我希望图书馆缩小 25%",或者“缩小 50%"。它会调整“满员瓷砖”与“稀疏瓷砖”的数量,以完美达到该目标,而不是被困在固定的 50% 缩减率上。
- 现实世界的结果: 作者在从小型到超大型(高达 130 亿参数)的模型上测试了这种方法。
- 速度: 在标准消费级显卡(A6000 GPU)上,PATCH 使模型的运行速度比原始未剪枝模型快了 1.18 到 1.38 倍。
- 智能: 与其他在加速时使模型“变笨”的方法不同,PATCH 实际上使模型比当前最先进的刚性方法(MaskLLM)更准确(提高了 0.37% 到 2.96%)。
总结
将 PATCH 视为清理巨型仓库的一种方式。与其随机扔掉东西(这会减慢叉车的速度),要么遵循一条会扔掉重要物品的愚蠢规则,PATCH 智能地指定特定区域保持满员,而其他区域则按照叉车喜欢的模式进行清理。其结果是,仓库的导航速度更快,同时仍保留所有关键知识。
以下是论文"PATCH: LEARNABLE TILE-LEVEL HYBRID SPARSITY FOR LLMS"的详细技术总结。
1. 问题陈述
大型语言模型(LLMs)虽然提供卓越的性能,但在部署过程中却面临难以承受的记忆力和计算成本。模型剪枝是主要的解决方案,但现有方法面临一个根本性的权衡:
- 非结构化稀疏性:允许非零元素出现在任意位置,在保持高精度的同时,却造成了不规则的内存访问模式,阻碍了高效的 GPU 加速(例如 Tensor Cores)。
- 半结构化稀疏性(例如 2:4):强制实施刚性模式(每四个元素中有两个非零元素),这种模式对硬件友好并能加速推理。然而,固定的 50% 稀疏率以及在所有层中均匀应用,往往导致显著的精度下降,特别是在使用一次性剪枝(one-shot pruning)方法时。
当前方法难以弥合非结构化剪枝的灵活性与半结构化剪枝的硬件效率之间的差距,特别是在无法连续调整稀疏率或在各层之间非均匀分配稀疏性方面。
2. 方法论:PATCH
作者提出了PATCH(Pruning with a Learnable Tile-level Configuration for Hybrid Sparsity,基于可学习图块级配置进行混合稀疏剪枝),这是一个能够在保持硬件兼容性的同时,实现 0% 到 50% 连续稀疏率的框架。
核心机制
PATCH 将权重矩阵划分为图块(tiles,例如 b1×b2)。对于每个图块,模型学习决定其应保持稠密(0% 稀疏)还是变为2:4 稀疏(50% 稀疏)。
- 混合掩码生成:最终掩码是“稠密”选项和"2:4 稀疏”选项的加权组合。
- 可学习分布:
- 图块选择:一个可学习的 logit 分布决定了图块是稠密还是稀疏的概率。这通过Gumbel-Softmax进行采样,以支持可微分优化。
- 模式选择:在被指定为稀疏的图块内,一个独立的可学习分布从 6 种可能的排列中选择特定的 2:4 模式。
- 优化目标:训练目标结合了:
- 标准建模损失(例如,下一个词元预测)。
- 稀疏性正则化:一个惩罚项,用于惩罚与目标全局稀疏率(ρ)的偏差,从而实现对 0–50% 范围的精确控制。
- 权重正则化:鼓励更大的权重幅度,以辅助梯度传播。
变体
- PATCHJoint:联合优化图块级选择(稠密 vs. 稀疏)以及稀疏图块内的细粒度 2:4 模式。用于较小模型(<10 亿参数)。
- PATCHTile:一种节省内存的变体,其中 2:4 模式被冻结(从高质量的一次性方法如 MaskLLM 初始化),仅优化图块级的稠密/稀疏决策。这使得在有限的 GPU 内存上扩展到更大模型(高达 130 亿参数)成为可能。
推理与部署
- STOICC 集成:由于标准 GPU 库(cuBLAS, cuSPARSELt)不支持混合稠密/稀疏图块,PATCH 利用STOICC,这是一个基于 Triton 的编译器。STOICC 自动调整内核配置(图块大小,例如 128×128)以高效处理混合执行。
- 硬件兼容性:通过确保稀疏图块严格遵循 2:4 模式,PATCH 利用 NVIDIA/AMD Tensor Cores 进行加速,而稠密图块则利用标准稠密内核。
3. 主要贡献
- 混合稀疏性框架:提出了第一种通过在同一权重矩阵内动态混合稠密和 2:4 稀疏图块来实现连续稀疏率(0%–50%)的方法。
- 可学习的非均匀分配:与固定的 2:4 方法不同,PATCH 学习自适应地分配稀疏性。它在敏感层(例如,早期/晚期 Transformer 块、注意力机制)中保持稠密,并在冗余层(例如,中间块、MLP 的 up/gate/down 矩阵)中应用稀疏性。
- 端到端可微分训练:使用 Gumbel-Softmax 联合优化掩码选择和权重剪枝,避免了与一次性剪枝相关的精度下降。
- 实际部署:展示了混合图块级稀疏性与编译器(STOICC)的首次成功集成,从而在消费级 GPU 上实现了实际的速度提升。
4. 实验结果
作者在 0.5B 到 13B 参数的模型(Qwen-2.5, LLaMA-2/3, Gemma-3)上评估了 PATCH,使用了 8 个零样本下游任务和 WikiText2 困惑度。
精度:
- PATCH 始终优于最先进的 2:4 剪枝方法(MaskLLM, SparseGPT, Wanda, ProxSparse)。
- 在LLaMA-2 7B上,PATCH(在 25% 稀疏率下)实现了51.58%的平均精度,而 MaskLLM(50% 稀疏率)为48.62%,ProxSparse 为45.92%。
- PATCH 显著缩小了与稠密模型的差距;在 25% 稀疏率下,其精度恢复到与稠密模型相差约 2 个百分点的水平。
- 困惑度:PATCH 的困惑度低于所有基线。对于 25% 稀疏率下的 LLaMA-2 7B,PPL 为5.86,而 MaskLLM 为6.78。
加速与效率:
- 在NVIDIA A6000 GPU(LLaMA-2 7B)上,PATCH 实现了比稠密基线1.18 倍至 1.38 倍的端到端加速。
- 内存占用减少至稠密模型的0.59 倍–0.76 倍。
- PATCH 优于非结构化剪枝(不提供加速)和刚性 2:4 剪枝(通常因精度下降过多而无法证明加速的合理性)。
消融研究:
- 全局 vs. 逐层:全局稀疏率目标(允许每层具有可变稀疏率)显著优于均匀的逐层目标。
- 图块大小:性能在不同图块大小下表现稳健,尽管 4×4 提供了最精细的控制。硬件效率决定了使用更大的执行图块(128×128)。
- 初始化:该方法对初始化先验具有鲁棒性;全局稀疏率目标允许在训练期间进行动态重新分配。
5. 意义
PATCH 通过解决“精度与加速”的困境,代表了高效 LLM 推理的重要一步。
- 弥合差距:它成功地将半结构化稀疏性的硬件加速优势与非结构化分配的灵活性结合起来。
- 实用性:它超越了理论剪枝,利用现有的编译器基础设施(STOICC)在实际硬件上展示了真正的加速。
- 可扩展性:节省内存的变体(PATCHTile)证明,混合稀疏性可以应用于大模型(13B+),而无需为训练掩码投入巨大的计算预算。
- 未来方向:该框架与模式无关,理论上可以扩展到支持新硬件的其他 N:M 模式(例如 4:8),使其成为未来 LLM 压缩的通用解决方案。
总之,PATCH 提供了一种灵活、可学习且硬件高效的剪枝策略,允许从业者以前所未有的粒度调整模型质量与推理速度之间的权衡。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。