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.

翻译：结构化稀疏性已成为一种广泛应用的模型剪枝技术，近年来被普遍采纳于包括CNN、Transformer模型，尤其是大语言模型（LLM）在内的多种架构中。权重置换是进一步提升剪枝后性能的一个有前景的方向，它通过重新排列模型权重，使其形成更易于剪枝的模式。然而，随着Transformer架构规模的扩大，置换搜索空间呈指数级增长，这迫使大多数方法依赖于贪心或启发式算法，从而限制了重排序的有效性。在本研究中，我们提出了一种新颖的端到端可学习置换框架。我们的方法引入了一个可学习的置换成本矩阵，用于量化交换给定权重矩阵任意两个输入通道的成本；一个可微分的二分图匹配求解器，用于在给定成本矩阵下获取最优的二进制置换矩阵；以及一个稀疏性优化损失函数，用于直接优化置换算子。我们在视觉和语言Transformer模型上广泛验证了所提方法，结果表明我们的方法在结构化稀疏性方面取得了最先进的置换效果。