跳转至

PATCH: Learnable Tile-Level Hybrid Sparsity for LLMs

会议: NeurIPS 2026
arXiv: 2509.23410
代码: https://github.com/Paramathic/patch
领域: 模型压缩
关键词: 大语言模型剪枝, 混合稀疏度, 块级稀疏, 2:4半结构化稀疏, 硬件加速推理

一句话总结

PATCH 提出了一种块级混合稀疏框架，在冻结模型权重的前提下将权重矩阵划分为硬件友好的分块，通过 Gumbel-Softmax 可微学习为每个分块自适应分配密集或 2:4 稀疏模式，实现了 0% 至 50% 连续可控的全局稀疏率与实测 1.18×–1.38× 的端到端推理加速。

研究背景与动机

随着大语言模型（LLM）参数量迈入数十亿乃至千亿级别，推理部署时的显存占用与计算开销变得极其沉重。模型剪枝作为压缩的核心支柱之一，一直面临着精度保留与硬件加速难以兼得的困境。非结构化稀疏允许非零权重自由分布，能够高度保真地保留模型精度，但其极度不规则的内存访问模式无法利用现代 GPU 的 Tensor Core 进行高效加速；而以 NVIDIA 和 AMD 硬件为代表的 2:4 半结构化稀疏虽然具备原生硬件加速支持，却强行施加了每个 4 元素块严格剪去 2 个权重的刚性约束（固定 50% 稀疏率），导致一阶/二阶单次剪枝方法（如 Wanda、SparseGPT、Thanos、ProxSparse）甚至端到端可微掩码学习方法（如 MaskLLM）在模型各层均匀剪枝时出现显著的精度劣化。

更为关键的是，深层 Transformer 各层乃至层内各子模块的敏感度具有极强的不均匀性。现有理论和经验研究表明，自适应分配不同层的稀疏预算对维持整体模型能力至关重要。然而，纯 2:4 结构化稀疏无法对冗余模块与敏感模块进行差异化分配，必须承受全局统一的 50% 稀疏度，因而限制了模型在精度与速度之间的权衡空间。同时，现有的层间自适应剪枝方法大多局限于非结构化稀疏，无法转化为实际的吞吐提升。

针对这一核心矛盾，本文提出了一种能够打通硬件加速与精细化分配壁垒的混合稀疏思路：既然整矩阵统一 2:4 过于严苛，而非结构化又无法加速，那么可以将矩阵划分为硬件友好的二维分块（Tile），让每个分块自主选择是保持密集（0% 稀疏）还是执行 2:4 剪枝（50% 稀疏）。核心 idea：将 LLM 权重矩阵划分为硬件友好的块级网格，通过端到端可微 Gumbel-Softmax 采样联合学习每个块的「密集/稀疏」状态与 2:4 内部模式，实现 0%–50% 连续平滑可控的全局自适应混合稀疏，并借助 STOICC 编译器实现真实端到端推理加速。

方法详解

整体框架

PATCH 在完全冻结预训练模型权重的条件下，通过端到端反向传播优化轻量级的掩码分布 logits。输入为模型的各层可剪枝权重矩阵与目标全局密度 \(\rho\)；输出为二值化的块级混合稀疏掩码与剪枝后模型。整个方法由块级分块与可微状态选择、细粒度 2:4 候选模式生成、软掩码插值融合以及目标稀疏度惩罚正则化四个阶段构成，最终经二值化后由混合稀疏编译器 STOICC 编译部署。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入预训练模型权重 W<br/>设定目标全局密度 ρ"] --> B["块级分块与可微状态采样<br/>P_tile 经 Gumbel-Softmax 生成软分块掩码"]
    A --> C["细粒度 2:4 候选模式生成<br/>P_2:4 经 6 种候选模式插值出软 2:4 掩码"]
    B --> D["双层混合掩码融合插值<br/>M_soft = M_tile + (1 - M_tile) ⊙ M_2:4"]
    C --> D
    D --> E["目标全局稀疏度引导训练<br/>LM 任务损失 + 目标密度惩罚 + 权重模长正则"]
    E --> F["符号二值化与 STOICC 编译部署<br/>输出硬件友好的混合 Dense/2:4 算子"]

关键设计

1. 块级分块与可微状态采样：将离散分块决策转化为连续概率优化 针对传统分块剪枝由于组合爆炸无法全局搜索最优块配置的难题，PATCH 将权重矩阵 \(\bm{W} \in \mathbb{R}^{d_1 \times d_2}\) 划分为 \(b_1 \times b_2\)（默认 \(128 \times 128\)）的分块网格。每个分块分配一个可学习的未归一化对数概率 \(\bm{P}_{\text{tile}} \in \mathbb{R}^{\frac{d_1}{b_1} \times \frac{d_2}{b_2}}\)，代表该分块保留完全密集的得分。为了在保留密集与采用稀疏两类状态之间进行可微选择，算法将固定的 0 标量与对数概率拼接构成两类分布 \([\bm{P}_{\text{tile}}, 0]\)，利用带有温度退火 \(\tau\) 和缩放因子 \(\kappa\) 的 Gumbel-Softmax 进行松弛采样：

\[\tilde{\bm{M}}_{\text{tile}} = \mathrm{GS}([\bm{P}_{\text{tile}}, 0]; \tau, \kappa)_{:,:,0} \otimes \mathbf{1}_{b_1 \times b_2}\]

其中 \(\mathbf{1}_{b_1 \times b_2}\) 为全 1 矩阵，\(\otimes\) 为克罗内克积（Kronecker product）。通过广播机制，同一分块内的所有权重元素共享相同的连续概率值。当概率偏向 1 时该块表现为密集，偏向 0 时则转向 2:4 稀疏。

2. 双层混合掩码融合插值：平滑连接局部结构与全局硬件友好度 为了在同一个权重矩阵中同时容纳密集区域与 2:4 稀疏区域，PATCH 将块级掩码与微观 2:4 掩码进行层级融合。对于 2:4 稀疏模式，遵循 MaskLLM 的可微参数化方案，将每 4 个连续元素分配一个 6 维分布向量 \(\bm{P}_{\text{2:4}} \in \mathbb{R}^{\frac{d_1 d_2}{4} \times 6}\)，从 6 种合法的 2:4 二进制排列矩阵 \(\bm{S} \in \mathbb{R}^{6 \times 4}\) 中通过 Gumbel-Softmax 加权得到软掩码 \(\tilde{\bm{M}}_{\text{2:4}}\)。随后，PATCH 通过凸组合将宏观块掩码与微观 2:4 掩码进行元素级融合：

\[\tilde{\bm{M}} = \tilde{\bm{M}}_{\text{tile}} + (1 - \tilde{\bm{M}}_{\text{tile}}) \odot \tilde{\bm{M}}_{\text{2:4}}\]

这种级联设计保证了反向传播时梯度的顺畅流动：当 \(\tilde{\bm{M}}_{\text{tile}}\) 接近 1 时，分块不受 2:4 掩码的噪声扰动，完整保留原始密集表示能力；当 \(\tilde{\bm{M}}_{\text{tile}}\) 接近 0 时，分块严格蜕化为 2:4 结构化稀疏。对于显存受限场景（如 8B 模型在单张 80GB GPU 上训练），PATCH 提供了 PATCHTile 变体，将 \(\tilde{\bm{M}}_{\text{2:4}}\) 冻结为 SparseGPT 或 MaskLLM 预先生成的固定二值掩码，仅优化参数量极小的分块 logits \(\bm{P}_{\text{tile}}\)。

3. 显式全局稀疏度引导训练：打破层间均匀剪枝的刚性约束 传统的权重衰减或范数惩罚难以精确控制最终剪枝后的整体稀疏度，导致压缩目标具有随机性。PATCH 设计了一种带绝对值偏差惩罚的全局稀疏约束，直接在所有可剪枝线性层之间联合计算当前全局密度与预设目标 \(\rho\) 的偏离程度，其完整的联合优化目标为：

\[\mathcal{L} = \mathcal{L}_{LM}\left(x; \{\tilde{\bm{M}}_{i} \odot \bm{W}_{i}\}_{i}\right) + \lambda_{1} \left\lvert \frac{\sum_{i} \|\tilde{\bm{M}}_{i}\|_{1}}{\sum_{i} \|\bm{W}_{i}\|_{0}} - \rho \right\rvert - \lambda_{2} \frac{\sum_{i} \|\tilde{\bm{M}}_{i} \odot \bm{W}_{i}\|_{2}^{2}}{\sum_{i} \|\bm{W}_{i}\|_{2}^{2}}\]

其中 \(\lambda_1\) 强力驱动全局保留权重比例收敛至目标密度 \(\rho\)，允许重要层保留更多密集块，冗余层分配更多 2:4 块；\(\lambda_2\) 则鼓励保留权重的幅值最大化以保障数值稳定性。训练过程中温度 \(\tau\) 从高到底衰减、\(\kappa\) 逐渐增大，使得软掩码在训练末期自然退火为明确的离散极值。

4. 硬件友好的分块执行与 STOICC 编译器协同 即使数学上得到了混合掩码，主流推理库（cuBLAS 仅支持稠密，cuSPARSELt 仅支持全局固定 2:4）亦无法直接加速包含混合块的单一 GEMM。PATCH 与 STOICC 编译器紧密协同，将 \(128 \times 128\) 的掩码分块（Mask Tile）映射为自适应调优的执行分块（Execution Tile，如 \(128 \times 64\) 或 \(64 \times 64\)）。由于自回归生成的瓶颈在于逐 Token 的解码阶段（Decoding），STOICC 首先在解码维度对候选调度配置进行探测，固定最优执行分块尺寸后复用于 Prefill 阶段，彻底解决了混合稀疏 GEMM 的访存调度难题。

一个完整示例

以一个前馈网络（FFN）的线性投影权重 \(\bm{W} \in \mathbb{R}^{4096 \times 4096}\)、目标全局稀疏度 35%（即全局密度 \(\rho = 0.65\)）为例： 1. 网格划分：矩阵被划分为 \((4096 / 128) \times (4096 / 128) = 32 \times 32 = 1024\) 个大小为 \(128 \times 128\) 的分块。对应初始化 1024 个标量 logits \(\bm{P}_{\text{tile}}\)。 2. 状态前向传播：在训练步中，若第 10 个分块的 logit 为正（如 \(+1.8\)），经过 Gumbel-Softmax 计算后其稠密保留概率接近 0.98；若第 100 个分块 logit 为负（如 \(-2.1\)），其稠密概率接近 0.05，转而激活其内部的 2:4 采样掩码。 3. 全局自适应平衡：损失函数中的绝对值正则项持续监测整网所有层密度的加权平均。若当前模型整体密度为 0.72，超出目标 0.65，正则项梯度促使注意力层等关键位置维持正 logit，而促使 FFN 中间层冗余分块的 logit 变负。 4. 推理二值化：训练结束时，直接根据 \(\bm{P}_{\text{tile}} > 0\) 判断分块类别。假设该层 1024 个分块中有 717 个被判定为 2:4 稀疏，307 个被判定为密集，该矩阵实现的稀疏度为 \((717 \times 0.5 + 307 \times 0) / 1024 \approx 35.0\%\)，并由 STOICC 打包转换为硬件加速内核。

损失函数 / 训练策略

  • 训练数据与步数：在 SlimPajama 数据集上训练 2000 步，序列长度为 4096 tokens，全局有效 batch size 为 256。
  • 权重保持冻结：训练期间所有基础模型权重 \(\bm{W}\) 保持固定不变，只更新掩码参数 \(\bm{P}_{\text{tile}}\) 与 \(\bm{P}_{\text{2:4}}\)，大幅降低了过拟合与显存开销。
  • 参数退火策略：Gumbel-Softmax 温度 \(\tau_t\) 随训练步数逐步下降，缩放因子 \(\kappa_t\) 逐步上升，使得掩码在早期保持充分的全局探索性，在训练末期收敛为准 one-hot 离散选择。

实验关键数据

主实验

PATCH 在 0.5B 至 13B 参数量的多个模型家族（Qwen-2.5、LLaMA-2、LLaMA-3/3.1/3.2、Gemma-3）上进行了系统评估。评估涵盖 8 项零样本下游任务（MMLU, PIQA, ARC-Easy, ARC-Challenge, Winogrande, OpenBookQA, RACE, HellaSwag）的平均准确率，以及 WikiText2 上的语言建模困惑度（Perplexity，PPL）。

下表汇总了较大模型家族（LLaMA-2 7B, LLaMA-2 13B, LLaMA-3.1 8B）在不同稀疏度下的核心表现（对应原论文 Table 2）：

稀疏度 方法 模式 LLaMA-2 7B 均分 (%) LLaMA-2 7B PPL LLaMA-2 13B 均分 (%) LLaMA-2 13B PPL LLaMA-3.1 8B 均分 (%) LLaMA-3.1 8B PPL
0% Dense (稠密基线) - 54.61 5.12 58.38 4.89 60.31 5.84
50% Magnitude 2:4 43.44 54.39 45.94 8.89 35.93 765.92
50% Wanda 2:4 44.30 11.15 47.95 8.91 41.77 21.29
50% SparseGPT 2:4 45.09 10.12 49.67 8.86 45.53 15.11
50% Thanos 2:4 44.80 11.19 49.33 8.80 45.72 16.09
50% ProxSparse 2:4 45.92 9.18 50.80 7.11 45.14 15.17
50% MaskLLM 2:4 48.62 6.78 N/A N/A 52.80 8.58
45% PATCHTile Dense/2:4 48.99 6.55 53.24 5.85 53.60 8.20
35% PATCHTile Dense/2:4 50.08 6.18 54.60 5.44 55.28 7.89
25% PATCHTile Dense/2:4 51.58 5.86 56.31 5.00 56.48 7.34

消融与推理加速分析

为了验证全局自适应稀疏度相对层间均匀稀疏的优越性，以及冻结 2:4 模式对分块调优的影响，下表呈现了 Qwen-2.5 0.5B 上的关键消融（对应原论文 Table 4 与 Table 5），以及 LLaMA-2 7B 在 A6000 GPU 上的实测推理速度与显存对比（对应原论文 Table 7）：

表 1：全局稀疏 vs 层内均匀稀疏与 2:4 掩码冻结来源消融 (Qwen-2.5 0.5B PPL)

稀疏度 全局稀疏 (Global) 层内均匀 (Layer-wise) PATCHTile (冻结 MaskLLM) PATCHTile (冻结 SparseGPT) PATCHJoint (联合学习)
45% 14.57 15.17 15.06 21.84 14.57
35% 13.84 14.48 14.55 17.29 13.84
25% 13.47 13.95 14.17 14.89 13.47

表 2：LLaMA-2 7B 在 A6000 GPU 上的吞吐、加速比与显存占用 (BS=16, 长度 128+128)

方法与配置 实际稀疏度 8任务均分 (%) 端到端加速比 权重显存相对比 实测峰值显存 (GB)
Dense 基线 0% 54.61 1.00× 1.00× 16.721
MaskLLM (STOICC / cuSPARSELt) 50% (2:4) 48.62 1.47× / 1.40× 0.56× -
PATCH (45% 稀疏) 45% 48.99 1.38× 0.59× 11.769
PATCH (35% 稀疏) 35% 50.08 1.27× 0.68× 12.530
PATCH (25% 稀疏) 25% 51.58 1.18× 0.76× 14.189

关键发现

  • 打破 50% 刚性边界的高效性：在 LLaMA-2 7B 上，仅从 50% 稀疏度退让到 45% 稀疏度，PATCHTile 的下游平均精度便提升至 48.99%，超越 MaskLLM（48.62%），同时仍维持了 1.38× 的端到端真实加速；在 25% 稀疏度下，精度进一步提升至 51.58%（较 MaskLLM 暴涨 2.96%），仅牺牲少许加速比（1.18×）。
  • 层深与子模块间的自适应分布：深入分析稀疏度在模型中的分布发现，网络的前几层和最后几层被分配了更多密集块以保护表示能力，大部分剪枝由网络中间层吸收。在模块内部，注意力层被判定为高度关键，其中 Key 和 Value 矩阵几乎完全免于剪枝，Output 矩阵在极高稀疏时轻度剪枝，Query 吸收了注意力层的大部分稀疏；而 MLP 中的 Up、Gate、Down 投影矩阵承担了绝大部分的稀疏度。
  • 软分块粒度的鲁棒性：消融实验表明，掩码分块尺寸从 \(4 \times 4\) 增大到 \(128 \times 128\) 时，模型困惑度变化微乎其微（45% 稀疏下仅从 14.55 微变至 14.57），这证明选用硬件友好的 \(128 \times 128\) 大分块既保证了计算调度的硬件高效性，又未损失算法灵活性。

亮点与洞察

  • 填补硬件加速与灵活稀疏之间的空白：PATCH 首次开创性地在硬件友好的半结构化稀疏框架内构建了连续稀疏度空间，让系统部署人员能够在 0% 到 50% 的区间内自由调节精度与延迟的帕累托前沿。
  • 双层解耦与自适应梯度的精巧配合：通过将分块选择（宏观）与 2:4 内部模式（微观）拆解为两级级联的 Gumbel-Softmax 分布，既支持联合端到端寻优，又衍生出冻结细粒度模式、显存极度友好的 PATCHTile 方案，使得百亿级模型剪枝得以在单机低显存资源下完成。
  • 全栈式算法-编译器协同：不满足于理论理论 FLOPs 减少，PATCH 从算法设计阶段就以 STOICC 编译器的执行分块（Execution Tile）粒度为导向，彻底解决了混合稀疏度在现有深度学习计算框架中缺乏算子支持的技术卡点。

局限与展望

  • 端到端训练开销高昂：相比 SparseGPT、Wanda 等基于二阶近似或激活幅值的单次（one-shot）秒级剪枝方法，PATCH 依赖反向传播训练掩码 logits，在 0.5B 到 13B 模型上耗费 18 至 160 个 H100 GPU-hours。尽管这是一次性成本，但在更大规模模型（如 70B+）上的普及仍需进一步压缩训练开销。
  • 稀疏度上限受限于 50%：由于分块的基础稀疏单元为 2:4（50%），PATCH 的有效全局稀疏度被天然限制在 0%–50% 之间。未来若需突破 50% 迈向更高压缩比，需引入 1:4 等更稀疏的硬件结构化模板。
  • 硬件后端的平台依赖与上下文敏感性：实测显示在 A100 上因通信和 Tensor Core 饱和度原因加速比略低于 A6000；同时在长序列场景下（如 8K/16K tokens）注意力耗时上升，GEMM 加速被摊薄，需要进一步结合稀疏注意力算子进行端到端全链路加速。

相关工作与启发

  • vs MaskLLM: MaskLLM 首次将 2:4 剪枝建模为端到端掩码概率学习，但其受限于刚性的整网 50% 均匀稀疏；PATCH 继承了可微掩码优化的理念，引入块级选择器和全局目标正则化，实现了非均匀自适应与连续稀疏率控制。
  • vs Wanda / SparseGPT / Thanos: 本地单次剪枝方法计算极快，但在 2:4 强约束下难以恢复跨层依赖与全局精度；PATCH 通过全网端到端反向传播和分块自适应显著缩小了与 Dense 模型的精度差距。
  • vs OWL / AlphaPruning: OWL 等方法探索了层间非均匀稀疏分配，但其输出均为非结构化稀疏矩阵，在主流 GPU 上无法获得加速；PATCH 成功将非均匀自适应思想落地至结构化 2:4 硬件执行模式。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出硬件加速支持的块级连续自适应混合稀疏，概念优雅且突破了 2:4 的刚性约束。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 0.5B 到 13B 多模型家族、多维度消融、组合量化/LoRA 实验以及真实 GPU 端到端推理速度基准。
  • 写作质量: ⭐⭐⭐⭐⭐ 框架图示清晰，理论推导与硬件系统考量结合紧密，工程细节详实完备。
  • 价值: ⭐⭐⭐⭐⭐ 为 LLM 稀疏化在现代计算体系下的实用化落地提供了极具说服力的可行路径。

相关论文