2025-11-17T08:49:21.061208

PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models

Zou, Yin, Pei et al.

Channel permutation is a powerful technique for enhancing the accuracy of N:M sparse models by reordering the channels of weight matrices to prioritize the retention of important weights. However, traditional channel permutation methods rely on handcrafted quality metrics, which often fail to accurately capture the true impact of pruning on model performance. To address this limitation, we propose PermLLM, a novel post-training pruning framework that introduces learnable channel permutation (LCP) for N:M sparsity. LCP leverages Sinkhorn normalization to transform discrete permutation matrices into differentiable soft permutation matrices, enabling end-to-end optimization. Additionally, PermLLM incorporates an efficient block-wise channel permutation strategy, which significantly reduces the number of learnable parameters and computational complexity. PermLLM seamlessly integrates with existing one-shot pruning methods to adaptively optimize channel permutations, effectively mitigating pruning-induced errors. Extensive experiments on the LLaMA series, Qwen, and OPT models demonstrate that PermLLM achieves superior performance in optimizing N:M sparse models. The code is available at https://github.com/lanchengzou/PermLLM.

academic

PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models

基本信息

论文ID: 2510.10136
标题: PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models
作者: Lancheng Zou, Shuo Yin, Zehua Pei, Tsung-Yi Ho, Farzan Farnia, Bei Yu (香港中文大学)
分类: cs.LG cs.AI
发表会议: NeurIPS 2025 (39th Conference on Neural Information Processing Systems)
论文链接: https://arxiv.org/abs/2510.10136
代码链接: https://github.com/lanchengzou/PermLLM

摘要

通道置换是一种通过重新排序权重矩阵的通道来优先保留重要权重，从而提高N:M稀疏模型准确性的强大技术。然而，传统的通道置换方法依赖于手工设计的质量度量，往往无法准确捕捉剪枝对模型性能的真实影响。为了解决这一局限性，本文提出了PermLLM，一个引入可学习通道置换(LCP)的N:M稀疏性后训练剪枝框架。LCP利用Sinkhorn归一化将离散置换矩阵转换为可微分的软置换矩阵，实现端到端优化。此外，PermLLM采用了高效的分块通道置换策略，显著减少了可学习参数数量和计算复杂度。PermLLM与现有的一次性剪枝方法无缝集成，自适应优化通道置换，有效缓解剪枝引起的误差。

研究背景与动机

问题定义

核心问题: 传统通道置换方法使用手工设计的质量度量(如保留权重重要性之和)来评估置换方案，但这些度量与实际剪枝误差之间存在差距。
重要性: 随着大语言模型规模的快速增长，模型压缩技术(如剪枝)对于高效部署至关重要。N:M稀疏性因其硬件友好性(NVIDIA Sparse Tensor Core支持)而备受关注。
现有局限:
- 手工设计的质量度量无法准确反映剪枝对模型性能的真实影响
- 传统方法无法充分捕捉复杂的层间交互
- 优化空间巨大(对于Cin个输入通道，有Cin!种可能的置换)

研究动机

论文通过一个具体例子(图1)展示了问题：最大化重要性分数的通道置换可能导致更大的输出误差，说明手工度量与实际性能之间存在本质性差异。

核心贡献

首次提出可学习通道置换(LCP): 将离散的通道置换问题转化为可微分的优化问题，实现端到端学习。
Sinkhorn归一化技术: 利用Sinkhorn归一化将离散置换矩阵松弛为软置换矩阵，解决了置换矩阵的非可微性问题。
分块通道置换策略: 显著降低参数复杂度从O(C²ᵢₙ)到O(Cᵢₙ×B)，计算复杂度从O(C³ᵢₙ)到O(Cᵢₙ×B²)。
通用框架设计: 可与现有一次性剪枝方法(Wanda、RIA等)无缝集成。
优异的实验表现: 在LLaMA系列、Qwen、OPT等多个模型上验证了方法的有效性。

方法详解

任务定义

给定预训练的权重矩阵W ∈ R^(Cout×Cin)，目标是找到最优的置换矩阵P，使得重排序后的权重矩阵Ŵ = WP在应用N:M稀疏性后能够最小化与原始密集模型的输出差异。

核心技术架构

1. 软置换矩阵松弛

将硬置换矩阵P松弛为软置换矩阵P̂：

S₀(X) = exp(X)
Sᵢ(X) = Tc(Tr(Sᵢ₋₁(X)))
S(X) = lim(l→∞) Sl(X)
P̂ = SL(WP/τ)

其中Tr和Tc分别表示行归一化和列归一化操作，τ为温度参数控制软置换矩阵的硬度。

2. 硬化过程与梯度近似

前向传播时，通过匈牙利算法将软置换矩阵硬化为严格的置换矩阵：

P = argmax P∈P Tr(P⊤P̂)

反向传播时使用直通估计器(STE)近似梯度：∂P/∂P̂ = 1。

3. 分块通道置换

为降低计算复杂度，将通道划分为多个大小为B的块，每个块内独立进行置换：

PB = diag(P₁, P₂, ..., PNB)
ŴB = WPB

参数数量从C²ᵢₙ降低到Cᵢₙ×B，计算复杂度从O(C³ᵢₙ)降低到O(Cᵢₙ×B²)。

优化目标

PermLLM直接最小化密集模型和稀疏模型输出之间的余弦相似度损失：

Lcosine(y, ỹ) = 1 - (y·ỹ)/(||y||·||ỹ||)

与现有剪枝方法的集成

PermLLM可以与任何基于重要性度量的一次性剪枝方法集成。对于给定的重要性矩阵S，置换后的重要性矩阵为Ŝ = SPB，掩码通过以下方式获得：

argmax M ∑∑ (M⊙Ŝ)i,kM:(k+1)M

使用STE处理argmax的非可微性问题。

实验设置

数据集与模型

模型: LLaMA 7B-13B, LLaMA-2 7B-13B, LLaMA-3.1 8B, Qwen-2.5 7B, OPT 6.7B
校准数据: 从C4数据集随机选择128个样本，每个包含1024个token
评估任务:
- 语言建模：Wikitext2 (困惑度)
- 零样本任务：HellaSwag, ARC-Easy/Challenge, OpenBookQA, RTE

对比方法

基线方法: SparseGPT, Wanda, RIA
传统通道置换: Wanda+CP, RIA+CP
本文方法: PermLLMWanda, PermLLMRIA

实现细节

优化器: AdamW
学习率: {1e-3, 5e-3}
Sinkhorn迭代次数: 5
温度参数: 从1线性衰减到0.1
块大小: 64
训练时间: 7B模型约2.5小时(4 GPUs)，13B模型约5.5小时(8 GPUs)

实验结果

主要结果

语言建模性能(Wikitext2困惑度)

方法	LLaMA 7B	LLaMA-2 7B	LLaMA-3.1 8B	Qwen-2.5 7B
Dense	5.68	5.47	6.24	7.74
Wanda	11.59	12.16	23.42	24.44
Wanda+CP	11.07	11.00	21.09	18.76
PermLLMWanda	9.41	9.39	14.03	13.58
RIA+CP	10.99	10.26	19.80	17.58
PermLLMRIA	9.95	9.60	15.79	15.93