🤖 machine learning

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

MaskPro 引入了一种新颖的线性空间概率框架，该框架通过学习分类分布，经由无放回 N 路采样高效生成大语言模型中的 (N:M) 稀疏性，同时采用损失残差的移动平均来稳定训练，从而相较于现有的贪婪或梯度驱动方法实现更优越的内存效率和鲁棒性。

原作者： Yan Sun, Qixin Zhang, Zhiyuan Yu, Xikun Zhang, Li Shen, Dacheng Tao

发布于 2026-05-14

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yan Sun, Qixin Zhang, Zhiyuan Yu, Xikun Zhang, Li Shen, Dacheng Tao

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你拥有一个庞大且细节极其丰富的图书馆（大型语言模型），里面藏书数十亿本（参数）。为了让这座图书馆更便于携带且阅读速度更快，你想要扔掉一些书。然而，你不能随意丢弃书籍；你必须遵循一条严格的规则：对于书架上每 4 本书，你必须保留恰好 2 本，并扔掉另外 2 本。 这就是论文中所称的 (N:M) 稀疏性（具体为 2:4）。

挑战在于，如何确定在每一组 4 本书中究竟该保留哪 2 本，以确保图书馆依然能讲述最精彩的故事。如果选错了，图书馆就会变得毫无意义。

以下是论文 MaskPro 如何利用简单的类比来解决这个问题：

1. 问题：“选择过多”的噩梦

人们此前尝试过两种主要方法来解决这个问题，但两者都存在重大缺陷：

“规则手册”方法： 这种方法使用简单的数学规则（例如“保留最重的书”）来猜测该保留哪些。它速度很快，但往往出错，因为它没有纵观全局。这就像试图仅通过身高来挑选最佳队员，而忽略了他们的实际技能。
“试错”方法： 这种方法试图通过测试数百万种可能性来学习最佳组合。问题在于，组合的数量如此庞大（就像试图在沙漠中找到一粒特定的沙子），导致计算机内存耗尽并崩溃。这就像试图记住一个 100 位数字密码锁的所有可能组合；这需要太多的脑力。

2. 解决方案：MaskPro（“智能彩票”）

作者提出了 MaskPro，这是一种在学习保留哪些书时既不会耗尽内存也不会做出错误猜测的新方法。

“线性空间”技巧（简化地图）
MaskPro 不再试图记忆每一组 4 本书中所有可能组合的概率（这是不可能的），而是为每组 4 本书创建一张简单的“彩票”。

旧方法： 想象一场拥有数十亿张彩票的抽奖，每一张彩票对应从 4 本书中选出 2 本的一种可能方式。你需要一个仓库来存储所有这些彩票。
MaskPro 方法： 相反，你只需要 4 个小盒子（每本书对应一个）。你在每个盒子里放一张票，上面写着：“这本书被选中的可能性有多大？”然后，你运行一场特殊的抽奖，从这 4 个盒子中选出 2 个获胜者，并确保不重复选择同一本书。
结果： 这将所需的内存从“仓库”缩减到了“背包”。它呈线性扩展，意味着即使图书馆变得巨大，你的背包也不会变重。

“平滑追踪器”（拥有记忆的教练）
在教导计算机挑选正确的书时，你会给它展示示例（数据）。有时，一组“糟糕”的书可能仅仅因为示例很简单而表现良好，而一组“优秀”的书可能因为示例很难而表现糟糕。这会令计算机感到困惑。

问题： 如果计算机看到一组“糟糕”的书在一次简单的测试中表现良好，它可能会想：“太棒了！我会继续这样做！”即使这只是一个偶然。
解决方法： MaskPro 引入了一个“拥有记忆的教练”（移动平均追踪器）。教练不再仅仅关注当前测试的分数，而是关注当前分数与过去几次测试平均分数之间的差异。
类比： 想象一个学生参加考试。如果他得了满分，教练会问：“这次考试容易吗？你通常能得这么高吗？”如果学生通常得 80 分，却突然在一次超简单的考试中得了 100 分，教练会说：“这不是真正的进步；我们暂时不要改变你的学习习惯。”这阻止了计算机被随机运气所迷惑，并保持训练的稳定。

3. 结果：快速、廉价且可靠

论文声称 MaskPro 是一个游戏规则改变者，原因有三：

内存高效： 它能运行在标准计算机上，而其他方法会导致崩溃。这就像打包一个能放进口袋的帐篷，而不是需要一个卡车才能运走的帐篷。
数据高效： 你不需要庞大的训练数据库来教导它。论文表明，即使只有一个单一示例，它也能有效学习（尽管更多会更好）。这就像一位厨师尝一次就能学会新食谱，而不是需要尝一千次。
性能： 它保持了模型的智能。当他们在著名的 AI 模型（如 LLaMA 和 Gemma）上测试时，MaskPro 比旧的“规则手册”方法更好地保持了模型的智能，并且表现几乎与昂贵的“试错”方法一样好，但无需付出高昂代价。

总结

MaskPro 是一个新工具，它通过教导模型剪除哪些部分来帮助缩小巨大的 AI 模型。它通过以下方式实现这一目标：

简化数学计算，使其不需要超级计算机来记住选择。
使用“智能教练” 来忽略随机运气并专注于真正的改进。
仅需极少数据即可工作，使其使用起来既实用又廉价。

作者已公开其代码，供他人尝试，证明你可以在不损失其“脑力”的情况下，使 AI 模型变得更小、更快。

技术摘要：MaskPro——面向大语言模型严格 (N:M) 稀疏性的线性空间概率学习

问题陈述
大语言模型（LLM）的快速扩展导致推理效率出现显著瓶颈，亟需有效的压缩技术。半结构化稀疏性，特别是 (N:M) 模式（即在每 M 个连续权重中保留 N 个元素），为加速和内存缩减提供了一种硬件友好的解决方案。然而，确定最优的 (N:M) 掩码是一个极具挑战性的组合优化问题。现有方法主要分为两类，且各自存在明显局限：

基于规则的方法：采用基于手工设计指标（例如最小化层重构误差）的贪婪式逐层搜索。虽然计算成本低廉，但存在固有偏差，且与端到端任务性能之间存在显著差距。
基于学习的方法（如 MaskLLM）：通过生成损失的梯度下降直接优化掩码。虽然效果显著，但训练成本高昂。具体而言，MaskLLM 需要存储 M 个权重组内所有可能掩码的 logits，导致内存复杂度为 $O(\binom{M}{N} \frac{d}{M})$ 。对于较大的 M，该复杂度呈指数级增长，使得其在大规模模型中不可行。此外，巨大的组合空间导致策略梯度方差极高，引发训练不稳定。

方法论：MaskPro
作者提出了 MaskPro，这是一种新颖的线性空间概率框架，旨在高效且稳定地学习 (N:M) 稀疏性。

线性空间概率松弛：
不同于为 M 个权重组中所有 $\binom{M}{N}$ 种可能的掩码定义分类分布，MaskPro 对问题进行了重构。它在每个组内的 M 个独立元素上建立分类分布，然后通过从该分布中进行 N 次不放回采样 来生成 (N:M) 掩码。
- 理论基础：作者证明，任何 (N:M) 掩码均可表示为 N 个不同基向量的概率和。
- 内存效率：这一重构将可学习参数（logits）的数量从 $O(\binom{M}{N} \frac{d}{M})$ 降低至线性 $O(d)$ ，其中 d 为权重总数。这消除了以往基于学习的方法所伴随的内存爆炸问题。
改进的策略梯度估计器 (PGE)：
为了优化 logits，MaskPro 利用策略梯度。作者指出，由于小批量差异导致的高方差（数据采样的损失变化往往远大于掩码采样的损失变化），在 LLM 中使用原始损失值的标准策略梯度是不稳定的。
- 损失残差：MaskPro 不直接使用原始损失 $f(m \odot w, \xi)$ ，而是基于相对于初始基线掩码 $m_0$ 的 损失残差 来更新 logits： $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$ 。这隔离了掩码变化的影响。
- 移动平均追踪器：为了进一步稳定训练并防止数值爆炸，引入移动平均追踪器 $\delta$ 来平滑基线。更新规则变为：
  $\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$
- 理论保证：论文提供了证明，表明这些改进的估计器仍然是真实策略梯度的无偏估计。此外，在采样掩码的损失小于基线损失两倍的情况下（实践中极易满足），改进估计器的方差严格低于普通策略梯度的方差。

主要贡献

线性空间框架：MaskPro 通过将选择过程建模为从 M 个元素的分类分布中进行 N 次不放回采样，将学习 (N:M) 稀疏性的内存需求从指数级 $O(\binom{M}{N} \frac{d}{M})$ 降低至线性 $O(d)$ 。
增强的策略梯度：作者引入了一种利用损失残差和移动平均追踪器的方差缩减更新机制，解决了大组合空间中策略梯度高方差导致的训练不稳定问题。
理论分析：论文提供了关于估计器无偏性以及所提更新规则方差缩减特性的全面证明。

实验结果
在多个大语言模型（Gemma-7B、Vicuna-7B、LLaMA-2-7B、DeepSeek-7B）上，针对各种下游任务（如 HellaSwag、RACE、PIQA）和稀疏模式（(2:4)、(4:8)、(8:16)）进行了广泛的实验。

性能：MaskPro 始终优于最先进的基于规则的方法（SparseGPT、Wanda、Pruner-Zero），并实现了与内存密集型 MaskLLM 相当的性能，准确率平均提升超过 2%。
内存可扩展性：当在 7B 模型上进行 (4:8) 稀疏性处理时，MaskLLM 在 8 张 A100 GPU 上遭遇内存不足（OOM）错误，而 MaskPro 无论稀疏模式如何，均保持稳定的内存使用量（约 36-37 GB）。
数据效率：MaskPro 对数据量表现出显著的鲁棒性。即使仅使用 1 个数据样本 进行训练，它也能实现稳定的性能，而 MaskLLM 则需要数十万样本才能收敛。
训练效率：与 MaskLLM 相比，该方法所需的训练样本和训练时间显著减少，同时保持了与基于规则的方法相当的内存使用量。

意义与主张
论文声称，MaskPro 通过解决基于学习的半结构化稀疏性的内存和稳定性瓶颈，为大语言模型的实际可扩展本地部署提供了一条可行路径。作者强调，他们的方法弥合了基于学习的方法的高性能与基于规则的方法的高效性之间的差距。通过将内存占用降低至线性空间，并通过改进的策略梯度稳定训练过程，MaskPro 使得在大规模模型上学习最优 (N:M) 掩码成为可能，而无需承担以往此类任务所伴随的过高计算开销。该方法被呈现为一种通用的概率掩码学习框架，对数据稀缺和模型规模具有鲁棒性。

1. 问题：“选择过多”的噩梦

2. 解决方案：MaskPro（“智能彩票”）

3. 结果：快速、廉价且可靠

总结

技术摘要：MaskPro——面向大语言模型严格 (N:M) 稀疏性的线性空间概率学习

类似论文