Learnable Permutation for Structured Sparsity on Transformer Models

TL;DR

This paper introduces a learnable permutation framework for structured sparsity in Transformer models, enabling more effective weight reordering to improve pruning performance through an end-to-end differentiable approach.

Contribution

It presents a novel end-to-end learnable permutation method with a differentiable bipartite matching solver for Transformers, surpassing previous heuristic-based approaches.

Findings

Achieves state-of-the-art permutation results for structured sparsity.

Effective on both vision and language Transformer models.

Improves post-pruning performance significantly.

Abstract

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.