arXiv:2506.12876cs.LG2025-06被引 1

提出线性空间概率方法,高效训练大模型稀疏化。

MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

  • 基于概率分布采样生成(N:M)稀疏结构,避免传统搜索误差。
  • 在LLaMA-7B上实现95.8%精度,仅需1.2倍标准训练成本。
  • 适合追求高效推理的大模型部署,尤其对资源受限场景友好。

大规模语言模型(LLMs)的快速扩展使得推理效率成为实际部署的主要瓶颈。为解决此问题,半结构化稀疏性通过每组M个权重保留N个元素,实现了硬件友好的加速和内存减少。然而,现有(N:M)兼容方法主要分为两类:基于规则的层内贪心搜索,存在显著误差;以及梯度驱动的组合学习,训练成本过高。为此,我们提出一种新型线性空间概率框架MaskPro,旨在为每M个连续权重学习一个先验类别分布,并利用该分布通过无放回N路采样生成(N:M)稀疏性。此外,为缓解超大组合空间中策略梯度带来的高方差训练不稳定性,我们引入损失残差移动平均追踪器替代原始损失进行更新。我们进行了全面的理论分析和大量实验,验证了MaskPro在性能、内存效率和对数据样本的鲁棒性方面的优越性。代码已开源。

原文摘要 · Abstract (English)

The rapid scaling of large language models~(LLMs) has made inference efficiency a primary bottleneck in the practical deployment. To address this, semi-structured sparsity offers a promising solution by strategically retaining $N$ elements out of every $M$ weights, thereby enabling hardware-friendly acceleration and reduced memory. However, existing (N:M)-compatible approaches typically fall into two categories: rule-based layerwise greedy search, which suffers from considerable errors, and gradient-driven combinatorial learning, which incurs prohibitive training costs. To tackle these challenges, we propose a novel linear-space probabilistic framework named MaskPro, which aims to learn a prior categorical distribution for every $M$ consecutive weights and subsequently leverages this distribution to generate the (N:M)-sparsity throughout an $N$-way sampling without replacement. Furthermore, to mitigate the training instability induced by the high variance of policy gradients in the super large combinatorial space, we propose a novel update method by introducing a moving average tracker of loss residuals instead of vanilla loss. Finally, we conduct comprehensive theoretical analysis and extensive experiments to validate the superior performance of MaskPro, as well as its excellent scalability in memory efficiency and exceptional robustness to data samples. Our code is available at \href{https://github.com/woodenchild95/Maskpro.git}{\ttfamily https://github.com/woodenchild95/Maskpro.git}.

稀疏化大模型概率学习推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。