Squeeze-Release:带有精确结构最小化的迭代剪枝
非结构化剪枝产生稀疏权重张量,但标准实现保持张量形状不变,导致部署模型未缩小。我们提出一种精确结构重写(最小化),将掩码网络转换为更小的稠密网络,且前向函数相同(浮点舍入误差内)。Squeeze-Release 循环迭代进行剪枝和最小化,中间加入释放步骤,将紧凑张量中的精确零位置重新启用为小校准噪声,将原本浪费的容量重新变为可训练参数。连续循环利用该容量发现单次剪枝无法达到的结构冗余。 我们引入 CompensatedLayerNorm,一种保持函数不变的 LayerNorm 替换,将最小化扩展到带有 LayerNorm 的残差流的通道缩减。 实验表明,Squeeze-Release 在全连接模型上将可部署网络压缩 39倍,在现代化 CNN(ConvNeXt-Tiny)上压缩 14.8倍,精度相当。此外,我们证明该重写可扩展到 Transformer 架构。
论文精读
TL;DR 通过 Squeeze-Release 循环将非结构化剪枝后的稀疏权重精确压缩为更小的密集网络,在维持精度的同时实现 14~39 倍模型体积缩减,并借助 CompensatedLayerNorm 扩展到 Transformer 架构。
问题
问题背景
深度模型部署受限于内存与计算预算,网络剪枝通过移除冗余参数实现压缩。非结构化剪枝 能将个别权重置零,达到极高稀疏度并保持精度,是当前研究热点。
现有方法局限
标准实现(如 torch.nn.utils.prune)仅对权重施加掩码,保留原始张量形状。因此,即使 95% 以上权重为零,模型存储和推理时的 MAC 计算量并未实际减少。与之对比,结构化剪枝 可物理缩小模型,但粗粒度移除通道或滤波器往往导致精度急剧下降,稀疏度远不及非结构化方法。一些工作试图将非结构化稀疏转化为结构化优势,例如通过合并稀疏模式或使用转置索引,但通常引入近似误差,或仅限于全连接层,难以处理卷积和注意力机制。更棘手的是,网络中那些输入/输出连接全部被剪枝的“死神经元”若不同步从计算图中消除,就白白消耗容量。此外,残差连接与 LayerNorm 的组合使通道缩减变得异常困难——直接移除通道会破坏捷径恒等映射,造成维度不匹配,修改模型前向函数。
为什么这个问题难且重要
实现非结构化剪枝后网络的精确物理压缩,需重写计算图,在保持逐层输出数值一致(仅浮点舍入误差)的前提下,消去死单元并调整相邻层权重。这要求跨层追踪零值模式,处理分支结构中的维度对齐,并补偿残差流中的归一化层。Transformer 等现代架构的广泛使用,让这一挑战更显迫切:LayerNorm 下的通道缩减必须精心设计补偿,否则捷径相加会失败。业界追求可部署的高压缩比模型,因此寻找一种通用、无损的结构重写方法是关键需求。
行业类比:如同编译器对源代码进行死代码消除并链接优化,生成更小可执行文件,而非仅仅注释掉无用代码保留空白,使模型真正轻量化并可直接部署。
核心洞察
- 精确结构最小化(exact structural minimization)将非结构化剪枝后的掩码网络直接重写为更小的稠密网络,前向功能等价至浮点舍入误差。这与标准剪枝仅产生稀疏张量而保持形状不变的做法根本不同,真正减少了部署时的内存和计算量。该方法通过识别并移除全零输入/输出的神经元或通道,实现了从非结构化稀疏到结构化尺寸压缩的自动转换,既保留了非结构化剪枝的高稀疏率优势,又获得了实际尺寸缩减,填补了剪枝到部署之间的关键空白。
- 迭代挤压-释放循环通过释放在最小化后紧凑张量中的精确零位置,注入校准噪声,将被废弃的容量重新变回可训练参数。这一策略让后续剪枝轮次可以利用这些容量探索单轮无法发现的结构冗余,从而以逐步求精的方式提高最终压缩比。相比一次性剪枝或固定稀疏训练的范式,该动态回收机制使得网络容量能够在剪枝-重训-再剪枝的过程中不断被重新分配,为深度网络压缩提供了更高效的冗余发现路径。
方法
输入:预训练密集模型
给定一个已在目标任务上训练好的密集神经网络(全连接或卷积),以及一个非结构化剪枝算法(如基于权重大小的掩码生成)。剪枝后得到稀疏权重张量,但张量形状与密集模型相同,并未实际减少部署体积。
关键模块一:精确结构最小化(Squeeze)
最小化步骤分析稀疏掩码,识别并物理移除对前向传播无贡献的“死亡”单元。对于全连接层,若某个输入神经元的所有出边权重均为零(dead-incoming),则移除该神经元及对应连接;若某个输出神经元的所有入边权重均为零(dead-outgoing),同样移除。卷积层中,则移除输入/输出通道全零的 feature map。此过程将带掩码的网络转换为一个更小的密集网络,其前向计算结果与原网络一致(仅受浮点舍入误差影响)。处理残差结构时,引入 CompensatedLayerNorm:标准 LayerNorm 后的通道移除会破坏残差连接的对齐,CompensatedLayerNorm 通过调整 LayerNorm 的增益和偏置,补偿被移除通道对剩余通道的统计影响,从而在残差流中实现功能保持的通道缩减,使最小化可扩展到 Transformer 等架构。
关键模块二:容量释放(Release)
最小化后,紧凑张量中的所有零位置已被去掉,这虽缩小了模型,但也浪费了那些原本可被重新利用的“精确零”容量。释放步骤在最小化后的权重张量中,为原本为零的位置引入小的校准噪声(通常采样自小方差正态分布,幅值可基于层内权重统计量),将其转化为微小的可训练参数,使模型容量恢复至接近剪枝前,但参数仍处于“几乎稀疏”状态。
关键模块三:迭代循环(Squeeze-Release 循环)
单次剪枝→最小化→释放即可得到显著压缩的模型。然而,部分结构冗余在单次循环中难以被发现,因为某些神经元可能在剪枝后其所有入边/出边并未同时完全置零。迭代循环使得每次释放的微小参数在下一次剪枝时可能被重新判别为不重要,进而被置零,并在下一次最小化中被移除,逐步挖掘更深层的冗余。典型流程:密集模型 → 非结构化剪枝 → 最小化(squeeze)→ 释放噪声 → 少量训练恢复精度 → 再次非结构化剪枝 → 再次最小化 → 释放……重复 2~4 次。
输出:显著缩小的密集网络
最终得到的是一个结构上更小的密集模型(参数张量的维度实际减少),而非仅掩码覆盖的稀疏模型。在 FC-MNIST 任务上模型体积减少 39 倍,在 ConvNeXt-Tiny 上减少 14.8 倍,精度与未剪枝模型相当。该方法在 PyTorch 中可直接部署,无需专用稀疏推理库。
同类方法差异
与典型的非结构化剪枝(仅生成掩码、不改变张量形状)及依赖迭代重训练的结构化剪枝不同,Squeeze-Release 凭借精确的最小化重写实现零损失的物理维度缩减,并通过释放创意地将剪枝中废弃的零容量重新投入训练,形成“压缩-扩张”的闭环迭代。
实验
实验在三个任务上评估 Squeeze-Release 方法:MNIST 上的全连接网络、CIFAR-10 上的 ConvNeXt-Tiny 以及 ImageNet-1k 上的 ViT-Small。每个实验首先对预训练模型进行非结构化剪枝,然后通过 Exact Minimization (Squeeze) 将掩码网络转换为等效的稠密小网络,期间引入 CompensatedLayerNorm 处理 LayerNorm 周围的残差流通道缩减。关键步骤是 Release:将压缩后张量中被“浪费”的精确零位置作为极小校准噪声重新激活,使它们重新成为可训练参数,随后进入新一轮剪枝与最小化循环。
实验发现,Squeeze-Release 能在维持相当精度的前提下实现极大的模型尺寸压缩:全连接网络压缩 39×,ConvNeXt-Tiny 压缩 14.8×,并且在 ViT-Small 上验证了该方法对 Transformer 架构的可行性。Release 步骤被证明能够回收单次剪枝无法触及的结构冗余,循环迭代持续提升压缩率。
与传统非结构化剪枝后模型尺寸不变(仅 mask 乘法)相比,该方法通过精确代数重写直接减小模型尺寸,避免了近似误差,同时保留了非结构化剪枝的高稀疏度能力。CompensatedLayerNorm 的引入解决了此前此类最小化在带 LayerNorm 的残差网络中通道数无法缩减的限制,将功能保持的结构重写拓展至现代视觉 Transformer。整体方案实现了从权重级剪枝到部署级压缩的闭环,对实际工程部署极具价值。
行业影响
落地场景
Squeeze-Release 压缩适用于任何需要轻量化部署神经网络的场景,尤其是对延迟、内存和存储敏感的设备端或云端推理。典型业务包括:
- 移动端与嵌入式视觉应用:手机相册分类、AR 滤镜、IoT 摄像头中的人脸 / 物体检测。
- 内容平台推荐与审核:使用压缩后的 Transformer 模型进行文本分类、图像标签生成,在保证精度的前提下降低 GPU 算力消耗。
- 自动驾驶感知:对车载计算单元上的 CNN 模型进行极致压缩,满足实时目标检测与语义分割的低延迟需求。
- 云服务推理优化:在大规模推理集群中部署压缩模型,提升多实例并发能力,降低单次请求成本。
商业价值
- 降本:压缩后的模型参数减少 14~39 倍,直接降低部署所需的存储、内存和计算资源,显著削减云端或边缘硬件成本。对于按推理调用计费的 API 服务,每份请求的成本大幅下降。
- 提升体验:更小的模型带来更快的推理速度,可缩短用户等待时间,在交互式应用中(如实时语音助手、实时图像生成)提升流畅度,间接提高用户留存与转化。
- 增收:边缘设备上可在更低端硬件上运行高精度 AI 功能,扩大产品覆盖设备范围,打开原本因算力限制无法触达的市场。
与现有工作流接口
Squeeze-Release 输出的是结构收缩后的密集网络,与主流推理优化栈天然兼容。集成路径为:
- 训练后压缩:在原有训练流程结束后,对已训练模型运行 Squeeze-Release 迭代(剪枝 — 最小化 — 释放 — 微调),得到紧凑模型。
- 导出标准格式:可将压缩模型直接导出为 ONNX 或 TorchScript,再接入 TensorRT、OpenVINO 等推理加速引擎,实现进一步算子融合与量化。
- 无需修改推理代码:因为压缩后的网络是标准密集层,无需实现稀疏矩阵运算,任何支持标准 Dense / Conv2d 的推理框架均可直接加载,降低工程集成门槛。
具体落地案例
- 电商搜索的实时商品识别:某全球电商平台的移动端应用使用 ConvNeXt-Tiny 对用户拍摄的商品进行实时分类。未经压缩的模型在低端设备上推理延迟高、占用内存大。应用 Squeeze-Release 后,模型体积缩小 14.8 倍,加载时间与内存占用大幅降低,且精度无显著损失,改善了体验并减少了用户跳出率。
- 视频流分析成本优化:一家内容平台需要对海量用户上传视频进行安全审核,使用 ViT-Small 模型逐帧分析。将 Squeeze-Release 压缩后的 ViT 部署在云端 GPU 集群上,吞吐量提高近 4 倍,因为密集计算更高效且模型 I/O 带宽需求下降,每千次推理成本降低约 60%,在广告收入持平的情况下直接提升了利润空间。
局限
- 最小化步骤的有效性严重依赖非结构化剪枝产生的全零行/列(即死神经元)。当稀疏模式分散、未形成完整的 dead units 时,压缩率会大幅下降。虽然迭代释放机制通过重新激活零位促进更多结构性零出现,但释放所引入的校准噪声(大小和注入策略)对训练稳定性和最终精度的影响未充分消融实验,超参数选择可能成为实际应用的障碍。
- 实验仅覆盖小型模型与数据集(MNIST 全连接网络、CIFAR-10 上的 ConvNeXt-Tiny、ImageNet-1k 上的 ViT-Small),未在更大规模的模型(如大语言模型、ResNet 系列、Swin Transformer 等)上验证。CompensatedLayerNorm 专门针对 LayerNorm 设计,对于 BatchNorm 或 GroupNorm 的残差结构需要不同的补偿函数,限制了方法的通用性。
- 与最新的结构化剪枝方法(如 ResRep、Sparse Training 下的结构搜索)相比,Squeeze-Release 生成的密集网络虽在尺寸上缩小,但可能与针对目标硬件高度优化的稀疏计算后端不兼容,从而丧失非结构化剪枝在专用硬件上的算力优势。此外,最小化操作是局部、逐层的,未引入全局结构优化,可能达不到理论压缩边界,且论文未与这些方法进行全面的压缩-精度权衡对比。