Learnable Permutation for Structured Sparsity on Transformer Models

📄 arXiv: 2601.22980v1 📥 PDF

作者: Zekai Li, Ji Liu, Guanchen Li, Yixing Xu, Ziqiong Liu, Xuanwu Yin, Dong Li, Emad Barsoum

分类: cs.LG, cs.CL

发布日期: 2026-01-30


💡 一句话要点

提出可学习的置换框架,用于Transformer模型结构化稀疏化

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: Transformer模型 结构化稀疏化 模型剪枝 权重置换 可学习置换

📋 核心要点

  1. Transformer模型剪枝中,权重置换能提升剪枝后性能,但搜索空间随模型增大而指数增长,现有方法难以有效重排序。
  2. 论文提出端到端可学习置换框架,通过可学习的置换代价矩阵、可微分二分图匹配求解器和稀疏性优化损失函数,直接优化置换算子。
  3. 在视觉和语言Transformer上的实验表明,该方法在结构化稀疏化任务中取得了最先进的置换结果。

📝 摘要(中文)

结构化稀疏化已成为一种流行的模型剪枝技术,被广泛应用于各种架构中,包括CNN、Transformer模型,尤其是在近年来大型语言模型(LLM)中。一个有希望进一步提高剪枝后性能的方向是权重置换,它将模型权重重新排序为更适合剪枝的模式。然而,随着Transformer架构规模的扩大,置换搜索空间的指数增长迫使大多数方法依赖于贪婪或启发式算法,限制了重排序的有效性。本文提出了一种新的端到端可学习置换框架。我们的方法引入了一个可学习的置换代价矩阵来量化给定权重矩阵的任意两个输入通道交换的代价,一个可微分的二分图匹配求解器来获得给定代价矩阵的最优二元置换矩阵,以及一个稀疏性优化损失函数来直接优化置换算子。我们在视觉和语言Transformer上广泛验证了我们的方法,证明了我们的方法实现了结构化稀疏化的最先进的置换结果。

🔬 方法详解

问题定义:论文旨在解决Transformer模型结构化稀疏化中,权重置换搜索空间过大,导致现有贪婪或启发式算法无法有效重排序的问题。现有方法无法充分利用权重置换的潜力,限制了剪枝后的模型性能。

核心思路:论文的核心思路是将权重置换问题转化为一个可学习的优化问题。通过学习一个置换代价矩阵,并利用可微分的二分图匹配求解器,找到最优的置换方案。这种方法避免了直接搜索庞大的置换空间,从而提高了置换的效率和效果。

技术框架:该方法主要包含三个模块:1) 可学习的置换代价矩阵:用于量化交换任意两个输入通道的代价。2) 可微分的二分图匹配求解器:根据代价矩阵,求解最优的二元置换矩阵。3) 稀疏性优化损失函数:用于直接优化置换算子,使其更适合结构化稀疏化。整体流程是,首先计算置换代价矩阵,然后使用二分图匹配求解器得到置换矩阵,最后通过稀疏性优化损失函数更新置换代价矩阵。

关键创新:最重要的技术创新点在于将离散的置换问题转化为连续的可学习问题。通过引入可学习的置换代价矩阵和可微分的二分图匹配求解器,实现了端到端的优化。与现有方法依赖贪婪或启发式算法不同,该方法能够直接优化置换算子,从而获得更好的置换效果。

关键设计:关键设计包括:1) 置换代价矩阵的初始化和更新策略。2) 二分图匹配求解器的选择和实现,需要保证可微分性。3) 稀疏性优化损失函数的构建,需要能够引导置换算子产生更适合结构化稀疏化的模式。具体的参数设置和网络结构细节在论文中应该有更详细的描述(未知)。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在视觉和语言Transformer上进行了广泛的实验验证,结果表明该方法在结构化稀疏化任务中取得了最先进的置换结果。具体的性能数据和对比基线需要在论文中查找(未知),但摘要明确指出该方法优于现有的贪婪或启发式算法。

🎯 应用场景

该研究成果可广泛应用于各种Transformer模型的压缩和加速,尤其是在资源受限的场景下,如移动设备和边缘计算。通过结构化稀疏化和权重置换,可以显著减小模型大小和计算复杂度,从而实现更高效的模型部署。未来,该方法有望进一步推广到其他类型的神经网络和更复杂的模型压缩任务中。

📄 摘要(原文)

Structured sparsity has emerged as a popular model pruning technique, widely adopted in various architectures, including CNNs, Transformer models, and especially large language models (LLMs) in recent years. A promising direction to further improve post-pruning performance is weight permutation, which reorders model weights into patterns more amenable to pruning. However, the exponential growth of the permutation search space with the scale of Transformer architectures forces most methods to rely on greedy or heuristic algorithms, limiting the effectiveness of reordering. In this work, we propose a novel end-to-end learnable permutation framework. Our method introduces a learnable permutation cost matrix to quantify the cost of swapping any two input channels of a given weight matrix, a differentiable bipartite matching solver to obtain the optimal binary permutation matrix given a cost matrix, and a sparsity optimization loss function to directly optimize the permutation operator. We extensively validate our approach on vision and language Transformers, demonstrating that our method achieves state-of-the-art permutation results for structured sparsity.