提出可学习的2:4稀疏化方法,让大模型推理更快更省资源。
MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models
- 用可学习分布建模2:4稀疏模式,通过Gumbel Softmax实现端到端训练。
- 在Wikitext上达6.72困惑度,优于现有方法(最高10以上)。
- 稀疏模式可迁移,适合不同任务或领域快速部署。
大型语言模型(LLMs)参数量巨大,通常存在显著冗余。本文提出MaskLLM,一种可学习的剪枝方法,能在LLMs中建立半结构化(即“N:M”)稀疏性,以降低推理时的计算开销。不同于设计新的重要性准则,MaskLLM通过Gumbel Softmax采样,将N:M模式显式建模为可学习分布,支持在大规模数据集上端到端训练。该方法具备两大优势:1)高质量掩码——可有效扩展至大规模数据集并学习精准掩码;2)可迁移性——概率化掩码分布建模支持稀疏性在不同领域或任务间的迁移。我们在多个模型(包括LLaMA-2、Nemotron-4和GPT-3,参数量从843M到15B)上评估了2:4稀疏性,实验结果表明显著优于现有方法。例如,在Wikitext上,主流方法困惑度(PPL)达10或更高,而密集模型仅为5.12,MaskLLM仅通过学习掩码(权重冻结)即达到6.72的更低困惑度。此外,其可学习特性支持为下游任务或领域定制无损的2:4稀疏化掩码。代码已开源:https://github.com/NVlabs/MaskLLM。
原文摘要 · Abstract (English)
Large Language Models (LLMs) are distinguished by their massive parameter counts, which typically result in significant redundancy. This work introduces MaskLLM, a learnable pruning method that establishes Semi-structured (or ``N:M'') Sparsity in LLMs, aimed at reducing computational overhead during inference. Instead of developing a new importance criterion, MaskLLM explicitly models N:M patterns as a learnable distribution through Gumbel Softmax sampling. This approach facilitates end-to-end training on large-scale datasets and offers two notable advantages: 1) High-quality Masks - our method effectively scales to large datasets and learns accurate masks; 2) Transferability - the probabilistic modeling of mask distribution enables the transfer learning of sparsity across domains or tasks. We assessed MaskLLM using 2:4 sparsity on various LLMs, including LLaMA-2, Nemotron-4, and GPT-3, with sizes ranging from 843M to 15B parameters, and our empirical results show substantial improvements over state-of-the-art methods. For instance, leading approaches achieve a perplexity (PPL) of 10 or greater on Wikitext compared to the dense model's 5.12 PPL, but MaskLLM achieves a significantly lower 6.72 PPL solely by learning the masks with frozen weights. Furthermore, MaskLLM's learnable nature allows customized masks for lossless application of 2:4 sparsity to downstream tasks or domains. Code is available at https://github.com/NVlabs/MaskLLM.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。