🤖 machine learning

SHUFFLESPARSE: Learned Shuffles for Structured Sparse Networks

本文介绍了 SHUFFLESPARSE，这是一种通过学习单个置换矩阵，在多种稀疏模式和训练范式下显著缩小结构化稀疏与非结构化稀疏网络之间的精度差距，同时保持极低推理开销的方法。

原作者： Abhishek Tyagi, Arjun Iyer, Liam Young, William H Renninger, Christopher Kanan, Yuhao Zhu

发布于 2026-07-22

📖 1 分钟阅读☕ 轻松阅读

原作者： Abhishek Tyagi, Arjun Iyer, Liam Young, William H Renninger, Christopher Kanan, Yuhao Zhu

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

想象一下，你正试图建造一辆速度最快的赛车。你拥有一个强大的引擎（一个深度神经网络），它能够解决复杂的难题，但它很重且极其耗油。为了让它变得更快，工程师们经常尝试拆除那些并非严格必要的零件——这个过程被称为“剪枝”（pruning）。如果你只是随机拆除螺栓和电线，赛车可能会运行得很糟糕，因为剩余的部件无法以符合工厂工具的方式进行连接。然而，如果你以一种非常特定、有组织的方式移除部件（比如按照整齐的网格每隔一个拆除一个螺栓），工厂的机器就能工作得快得多。这就是“结构化稀疏性”（structured sparsity）的世界：通过迫使 AI 模型遵循整齐、可预测的模式，使其变得更小、更快。

但问题在于，过于整齐会让赛车变得笨拙。如果你强迫引擎遵循僵化的网ب格，你可能会不小心切掉那个在特定转弯时至关重要的零件，导致赛车无法应对复杂的弯道。这就是研究人员面临的问题：结构化模式虽然快，但与“非结构化”方法（即可以在任何地方移除部件的方法）相比，往往会损失精度。核心疑问是：我们能否在保持整齐网格的速度的同时，又不失去应对任何转弯的灵活性？这篇论文深入探讨了这一谜题，探索我们是否可以教会 AI 仅仅通过微调其内部连接的排列方式，就让这些僵化的模式完美运作。

来自罗彻斯特大学的研究团队为这项研究引入了一个聪明的技巧，叫做 SHUFFLESPARSE。把神经网络层想象成一个巨大的房间，里面挤满了试图与一群专家（权重）交流的人（数据）。在标准的“结构化”设置中，专家们被排列成僵硬的行和列，就像阅兵式上的士兵。这让管理者可以快速地大声下达指令（计算速度快），但如果房间里的人需要与站在错误行里的专家交谈，就会产生问题。

通常，解决方案是让专家们随心所欲地站立（非结构化），但这样管理者就会感到困惑并减慢速度。SHUFFLESPARSE 提供了一个折中方案。它在对话开始前增加了一个神奇的“洗牌”（shuffle）步骤。想象一个舞池，在音乐响起之前，每个人都被告知要根据一个特定的、学习到的模式交换座位。这种洗牌的设计初衷是：当人们最终坐下来与那些排列僵硬的专家交谈时，他们实际上是在与“正确”的人交谈，即便专家们的位置并未改变。

研究发现，通过教会 AI 在僵化模式之外同时学习这种特定的“洗牌”（一个置换矩阵），模型可以找回由于被迫进入网格而损失的几乎所有精度。这就像是意识到士兵们不需要改变阵型；他们只需要让新兵以不同的顺序在他们面前排好队即可。

团队在两种截然不同的场景下测试了这个想法。首先，他们在图像识别任务（如识别猫和狗）和语言任务上从头开始训练模型（动态稀疏训练）。他们发现 SHUFFLESPARSE 始终能缩小僵化、快速的模型与灵活、缓慢的模型之间的差距。例如，在名为 ViT-B/16 的流行图像模型上，加入此洗牌功能使得在极高稀疏度（90–95%）下的精度差距从接近 2% 降低到了不到 1%。在像 GPT-2 这样的语言模型中，它同样提升了 AI 对文本的理解能力，显著降低了“困惑度”（衡量混乱程度的指标）。

其次，他们将此方法应用于已经训练好的、庞大的语言模型（如 LLaMA-2 和 Qwen），这些模型是冻结状态且无法重新训练的。他们使用“单次”（one-shot）剪枝法切除权重，然后应用了 SHUFFLESPARSE 洗牌。结果令人瞩目：在 LLaMA-2 7B 模型上，这种方法比表现最好的非洗牌方法在零样本准确率（zero-shot accuracy）上提升了 4.6 个点。这表明，即使对于巨大的、预制好的模型，一个简单的学习型重排也能在无需重新训练整个模型的情况下，释放出隐藏的潜力。

至关重要的一点是，作者展示了这不仅仅是关于随机洗牌。当他们使用随机、固定的洗牌而非学习到的洗牌来测试系统时，性能反而下降或保持不变。魔力在于 AI 学习到了最适合该任务的特定洗牌方式。论文还指出，虽然这种洗牌会带来微小的成本（增加约 3% 到 8.7% 的运行时间），但考虑到它带来的巨大精度提升，这个代价是非常值得的，尤其是因为它保留了结构化模式的核心速度优势。

简而言之，SHUFFLESPARSE 表明我们不必在速度和智能之间做选择。通过教会 AI 在撞击到僵化、快速处理的齿轮之前，先对其输入进行恰当的重排，我们可以实现两全其美：既拥有极快的速度，又拥有惊人的精度。作者总结道，这种学习型的置换是一个通用的工具，适用于不同类型的僵化模式和不同种类的 AI 任务，为实现高效、高性能的 AI 提供了一条充满前景的路径。

技术摘要：SHUFFLESPARSE：为结构化稀疏网络学习置换（Learned Shuffles for Structured Sparse Networks）

1. 问题陈述

结构化权重稀疏性（例如 N:M、块状、对角线模式）对于在现代 GPU 上加速训练和推理至关重要，因为它与厂商优化的算子相匹配。然而，结构化稀疏性的性能始终低于非结构化动态稀疏训练（DST）和一次性剪枝（one-shot pruning），尤其是在极高稀疏度（90–95%）的情况下。作者认为，这种性能差距的根本原因在于缺乏表达能力：虽然稠密层可以实现任何非零权重的模式，但结构化模式将网络限制在了一组固定且僵化的权重配置中。这种僵化性限制了网络学习最优表示的能力，导致其准确率相比非结构化对应方法出现下降。

现有的解决方案，如 AutoShuffleNet（密集网络中的通道置换）或 NVIDIA 的置换搜索（针对固定 2:4 权重的静态置换），都与特定的设置或稀疏模式绑定在一起。此前尚无研究能够展示一种统一的置换原语，既能同时适用于从头开始的训练（training-from-scratch），也能适用于一次性剪枝（one-shot pruning）下的各种典型结构化模式。

2. 方法论：SHUFFLESPARSE

作者引入了 SHUFFLESPARSE，这是一个统一的框架，它为每一层学习一个单一的置换矩阵，并将其与结构化权重矩阵共同学习。稀疏化层的核心公式为：
$y = (W\Pi)x$
其中 $W$ 是遵循固定模式 $S$ 的结构化稀疏权重矩阵， $\Pi$ 是学习到的列置换矩阵。

核心技术组件

可微置换估计器（Differentiable Permutation Estimator）： 由于置换矩阵是离散的，SHUFFLESPARSE 采用了软松弛（soft relaxation）方法。它学习一个双随机矩阵 $M$ ，并在学习阶段结束时解码出一个硬置换 $\Pi$ 。为了强制执行置换约束，作者在任务损失函数中添加了一个基于 $\ell_1 - \ell_2$ 范数差异的惩罚项 $P(M)$ 。该惩罚项驱动 $M$ 随着训练的进行向单热（one-hot）配置（即有效的置换）靠拢。
统一的应用场景： 该公式适用于两种不同的机制：
1. 从头开始的动态稀疏训练 (DST)： 稀疏掩码 $S$ 、稠密权重 $\Theta$ 和软置换 $M$ 被共同学习。剪枝与生长操作发生在置换后的坐标中，从而确保重要性得分与推理时的通道分组保持一致。
2. 预训练大语言模型 (LLM) 的一次性剪枝： 预训练权重被冻结。仅优化置换矩阵 $M$ ，以最小化校准集上稀疏输出与稠密输出之间的重建损失（余弦相似度）。
推理效率： 在推理阶段，学习到的置换被“硬化”（转换为单热矩阵）并固定下来。SHUFFLESPARSE 并非执行显式的矩阵乘法，而是通过**预计算的索引映射（precomputed index maps）**来实现置换。结构化稀疏 GEMM 算子通过置换后的地址读取输入，从而保留了原始算子的结构和 FLOP 数。由于存在不完美的内存访问合并问题，这会产生极小的开销（≤8.7%），但避免了单独的置换操作。

3. 核心贡献

一种统一的原语： SHUFFLESPARSE 引入了一种单一的层级公式，既适用于从头开始的动态稀疏训练，也适用于预训练大语言模型（LLM）的剪枝，并能统一应用于各种结构化模式（块状、N:M、对角线）。
弥合准确率差距： 该方法一致地缩小了结构化方法与非结构化基准之间的性能差距。
- 在 ViT-B/16 (ImageNet-1K) 上，其相对于最佳非结构化方法的平均准确率差距从 1.92% 降低到了 0.97%（使用 DynaDiag）。
- 在 GPT-2 (WikiText-103) 上，困惑度（perplexity）差距在所有结构化方法中均显著降低（例如，对于 GPT-2 Medium 使用 DynaDiag，从 3.75% 降至 2.63%）。
提升剪枝性能： 在应用于 2:4 稀疏度的冻结预训练 LLM 时，与最强的非置换基准（Wanda）相比，SHUFFLESPARSE 使 LLaMA-2 7B 的零样本下游准确率提升了 +4.62 个百分点，Qwen-2.5 3B 提升了 +2.95 个百分点，超越了之前的最先进置换方法（PermLLM）。

4. 实验结果

视觉领域 (ViT-B/16)： 在四种结构化 DST 方法（SRigL, PixelatedBFly, DSB, DynaDiag）中，SHUFFLESPARSE 始终优于结构化基准和随机置换对照组。在高度稀疏（90–95%）时，改进最为显著。
语言领域 (GPT-2)： 在所有稀疏度和模型规模下，SHUFFLESPARSE 都降低了相对于非结构化 RigL 基准的困惑度差距。
LLM 剪枝： 在一次性 2:4 剪枝中，SHUFFLESPARSE 在 LLaMA-2 7B 和 Qwen-2.5 3B 上实现了与稠密模型性能之间最小的差距，证明了学习到的置换即使在权重冻结的情况下也能恢复出非结构化水平的准确率。
消融实验：
- 随机 vs. 学习： 用固定的随机置换代替学习到的置换会导致性能下降或收益微乎其微，这证实了性能提升源于学习特定任务的置换，而非任意的重排序。
- 敏感性： 该方法对置换惩罚强度 ( $\lambda$ ) 在一个数量级范围内都是鲁棒的。
- 置换强度： 分析显示，相比于早期层和 MLP 投影，更深层的网络和注意力投影学习到了更强的置换（相对于单位矩阵有更大的偏差），这一趋势在 DST 和剪枝两种机制中均保持一致。

5. 重要性与主张

本文声称，学习到的置换是恢复结构化稀疏模式中非结构化水平准确率的一种通用工具。 通过为每一层学习一个单一的置换矩阵，SHUFFLESPARSE 在不牺牲结构化稀疏带来的硬件加速优势的前提下，恢复了因僵化结构约束而损失的表达能力。

作者强调，这种方法：

保留了宿主结构（如 N:M 算子）提供的训练和推理加速。
仅增加极小的推理开销（≤8.7%）和适度的训练开销。
提供了一个跨不同网络架构（ViT, GPT, LLaMA）和稀疏模式（块状, N:M, 对角线）的统一解决方案。

这项工作确立了结构化稀疏中的准确率差距并非硬件友好型模式本身的固有局限，而是由于网络无法将其输入基（input basis）适配到这些模式中所导致的——而 SHUFFLESPARSE 通过学习到的置换解决了这一问题。