让变换模型权重自动排列,提升剪枝效果
Learnable Permutation for Structured Sparsity on Transformer Models
- 设计可学习的权重排列矩阵,量化任意通道交换代价
- 通过可微二分匹配求解最优排列,实现端到端优化
- 在视觉与语言模型上均达到当前最好剪枝性能
结构化稀疏性已成为主流模型压缩技术,广泛应用于卷积神经网络、Transformer 模型及大型语言模型。为提升剪枝后性能,权重排列是一种有前景的方向,它通过重排权重使其更易剪枝。然而,随着 Transformer 架构规模增大,排列搜索空间呈指数增长,现有方法多依赖贪心或启发式算法,限制了重排效果。本文提出一种端到端可学习的排列框架:引入可学习的排列代价矩阵,量化任意输入通道交换的代价;采用可微二分匹配求解器,根据代价矩阵获得最优二值排列矩阵;并设计稀疏性优化损失函数,直接优化排列操作。我们在视觉和语言 Transformer 上广泛验证该方法,结果表明其在结构化稀疏性剪枝方面达到当前最优水平。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。