Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learnable Permutation for Structured Sparsity on Transformer Models

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

TL;DR

This work tackles the challenge of achieving hardware-friendly structured sparsity in large Transformer models by addressing the inefficiency of traditional permutation methods. It introduces an end-to-end learnable permutation framework comprising a permutation cost predictor, a differentiable bipartite matching solver based on Sinkhorn iterations, and an end-to-end sparsity loss with knowledge distillation to align pruned networks with dense teachers. The method is validated across vision, language, and multimodal transformers, achieving state-of-the-art permutation performance under $N:M$ sparsity with limited accuracy loss and fast convergence. Its practical impact lies in enabling scalable, hardware-efficient pruning for large models without expensive weight updates or ad hoc heuristics, while remaining compatible with existing pruning pipelines like Wanda.

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.

Learnable Permutation for Structured Sparsity on Transformer Models

TL;DR

This work tackles the challenge of achieving hardware-friendly structured sparsity in large Transformer models by addressing the inefficiency of traditional permutation methods. It introduces an end-to-end learnable permutation framework comprising a permutation cost predictor, a differentiable bipartite matching solver based on Sinkhorn iterations, and an end-to-end sparsity loss with knowledge distillation to align pruned networks with dense teachers. The method is validated across vision, language, and multimodal transformers, achieving state-of-the-art permutation performance under sparsity with limited accuracy loss and fast convergence. Its practical impact lies in enabling scalable, hardware-efficient pruning for large models without expensive weight updates or ad hoc heuristics, while remaining compatible with existing pruning pipelines like Wanda.

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.
Paper Structure (31 sections, 8 equations, 4 figures, 7 tables)

This paper contains 31 sections, 8 equations, 4 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methods
  4. Preliminaries
  5. Optimize N:M Sparsity via Channel Permutation.
  6. Channel Permutation for Linear Layers.
  7. Channel Permutation for Transformer Layers.
  8. Learnable Channel Permutator
  9. Permutation Cost Predictor.
  10. Differentiable Bipartite Matching Solver.
  11. End-to-End Optimization Objectives.
  12. Training and Inference Details
  13. Training.
  14. Inference.
  15. Experiments
  16. ...and 16 more sections

Figures (4)

  • Figure 1: The channel permutation process enhances the friendliness of the 2:4 sparsification, making the overall saliency of the pruning metric more preserved.
  • Figure 2: Channel permutation for Linear layers. To guarantee output consistency, after the input channel of the $i$-th layer weight is permuted, the input activation of that layer should also be permuted accordingly, which can be realized by permuting the output channel of the previous ($i-1$)-th layer weight accordingly.
  • Figure 3: An overview of channel permutation for Transformer layers. The alignment of the input channel permutation of the current layer's weights to the output dimension of the previous layer's weights reflects a structural coupling.
  • Figure 4: Overview of our learnable permutation framework. A permutation cost predictor generates cost matrices for each linear layer, which are converted into permutation matrices via a differentiable bipartite matching solver. The original weights are permuted accordingly and then sparsified using an N:M mask generator. During training, the pruned weights are inversely permuted and used for loss computation. The entire process supports end-to-end optimization while maintaining gradient flow through the binary permutation matrix generated by the differentiable solver.