[论文解读] DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

Minsik Cho, Keivan Alizadeh-Vahid+2 more|arXiv (Cornell University)|Aug 28, 2021

Advanced Image and Video Retrieval Techniques参考文献 39被引用 5

一句话总结

本文提出DKM，一种可微分的k-means聚类层，支持神经网络权重与聚类中心的端到端联合优化，实现最先进的模型压缩效果，且精度损失极小。通过将k-means重述为可微分的注意力机制，DKM实现了高精度、低比特量化——在MobileNet-v1上实现22.4倍压缩，Top-1 ImageNet精度达63.9%；在DistilBERT上实现11.8倍压缩，GLUE基准上仅损失1.1%精度。

ABSTRACT

Deep neural network (DNN) model compression for efficient on-device inference is becoming increasingly important to reduce memory requirements and keep user data on-device. To this end, we propose a novel differentiable k-means clustering layer (DKM) and its application to train-time weight clustering-based DNN model compression. DKM casts k-means clustering as an attention problem and enables joint optimization of the DNN parameters and clustering centroids. Unlike prior works that rely on additional regularizers and parameters, DKM-based compression keeps the original loss function and model architecture fixed. We evaluated DKM-based compression on various DNN models for computer vision and natural language processing (NLP) tasks. Our results demonstrate that DKM delivers superior compression and accuracy trade-off on ImageNet1k and GLUE benchmarks. For example, DKM-based compression can offer 74.5% top-1 ImageNet1k accuracy on ResNet50 DNN model with 3.3MB model size (29.4x model compression factor). For MobileNet-v1, which is a challenging DNN to compress, DKM delivers 63.9% top-1 ImageNet1k accuracy with 0.72 MB model size (22.4x model compression factor). This result is 6.8% higher top-1accuracy and 33% relatively smaller model size than the current state-of-the-art DNN compression algorithms. Additionally, DKM enables compression of DistilBERT model by 11.8x with minimal (1.1%) accuracy loss on GLUE NLP benchmarks.

研究动机与目标

• 解决压缩已高度紧凑的深度神经网络（如MobileNet）的挑战，这些网络缺乏显著的权重冗余。
• 克服k-means聚类在反向传播中不可微的问题，从而限制网络权重与聚类中心的联合优化。
• 在训练过程中实现任务感知聚类，而无需修改原始损失函数或模型架构。
• 通过可微分聚类在计算机视觉与自然语言处理模型中实现更优的压缩-精度权衡。
• 提供即插即用、无参数的层，支持多维k-means聚类，实现高效的权重共享。

提出的方法

• 提出DKM作为可微分的k-means聚类层，将离散的聚类分配替换为可微分的注意力机制。
• 通过输入权重与可学习聚类中心之间的点积注意力机制，将k-means聚类形式化为软分配问题。
• 使用温度控制的softmax来调节聚类分配的锐度，从而在反向传播中实现梯度流动。
• 将DKM作为轻量级、无参数的层插入神经网络中，以实现网络权重与聚类中心的联合优化。
• 在推理阶段将预测结果映射到最近的聚类中心，同时在训练过程中通过可微分注意力机制保留梯度更新。
• 将DKM无缝集成到现有DNN中，无需架构更改，同时保持原始损失函数与训练流程不变。

实验结果

研究问题

• RQ1 k-means聚类能否被设计为可微分，以支持与深度神经网络的端到端训练？
• RQ2 与标准k-means或先前的可微分方法相比，可微分k-means聚类是否能带来更优的压缩-精度权衡？
• RQ3 DKM能否有效压缩如MobileNet-v1这类高度参数高效的模型，而这类模型通常难以压缩？
• RQ4 在NLP模型（如DistilBERT）上，DKM在低精度量化下的表现如何？
• RQ5 DKM能否在不造成精度下降的前提下，对不同层（如嵌入层与Transformer层）实现细粒度的压缩率控制？

主要发现

• DKM在ImageNet1k上使用ResNet50实现了74.5%的Top-1精度，模型大小仅为3.3MB，压缩比达29.4倍。
• 对于MobileNet-v1，DKM实现了63.9%的Top-1精度，模型大小为0.72MB（压缩比22.4倍），在精度上比最先进方法高出6.8%，在模型尺寸缩减上高出33%。
• 在DistilBERT上，DKM实现了11.8倍的模型压缩，GLUE基准上仅损失1.1%精度，优于GOBO在压缩与精度两方面的表现。
• 基于DKM的压缩在训练与推理之间导致了0.2%至1.7%的精度下降，这是由于可微分训练与离散推理之间的差距所致，但该差距极小且可管理。
• 4/4比特配置（Transformer层使用4比特，嵌入层使用4比特）在所有BERT变体中均优于1比特与2比特配置，尤其在难以压缩的模型（如MobileBERT）上表现更优。
• DKM实现了平滑收敛与训练过程中的渐进式权重聚类，无需自定义训练循环或额外正则化。