[論文レビュー] 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

ひとこと要約

本稿では、ニューラルネットワークの重みとクラスタ中心点の共同最適化を可能にする、微分可能k-meansクラスタリング層であるDKMを提案する。これにより、エンド・ツー・エンドの学習が可能となり、最小限の精度損失で最先端のモデル圧縮が達成される。k-meansを微分可能アテンション機構に再定式化することで、高精度で低ビット幅の量子化が可能となり、MobileNet-v1では22.4倍の圧縮率を達成し、ImageNetのトップ-1精度は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クラスタリングを効率的な重み共有に対応する、プラグアンドプレイかつパラメータフリーのレイヤーを提供する。

提案手法

• 離散的クラスタ割り当てを微分可能アテンション機構に置き換えることで、微分可能k-meansクラスタリングレイヤーDKMを提案する。
• 入力重みと学習可能なクラスタ中心点の間でドット積アテンション機構を用いてk-meansクラスタリングをソフト割り当て問題として定式化する。
• 温度制御付きソフトマックスを用いてクラスタ割り当ての鋭さを制御し、バックプロパゲーション中に勾配伝播を可能にする。
• ネットワークに組み込むことで、ネットワーク重みとクラスタ中心点の共同最適化を可能にする、軽量でパラメータフリーのレイヤーとしてDKMを適用する。
• 推論時には予測値を最も近い中心点にスナップするが、訓練中は微分可能アテンション機構により勾配更新を維持する。
• アーキテクチャの変更なしに既存のDNNにDKMを統合し、元の損失関数と訓練パイプラインを保持する。

実験結果

リサーチクエスチョン

• RQ1 k-meansクラスタリングを微分可能にすることで、深層ニューラルネットワークとのエンド・ツー・エンド学習が可能になるか？
• RQ2 微分可能k-meansクラスタリングは、標準k-meansや先行の微分可能アプローチと比較して、より優れた圧縮-精度トレードオフを達成するか？
• RQ3 DKMは、従来圧縮が難しいとされる、非常にパラメータ効率の良いモデル（例：MobileNet-v1）を効果的に圧縮できるか？
• RQ4 DKMは、自然言語処理モデル（例：DistilBERT）において、低精度量子化下でも優れた性能を示すか？
• RQ5 DKMは、異なる層（例：埋め込み層 vs. トランスフォーマー層）に対して、精度の劣化を伴わずに圧縮率の細かい制御が可能か？

主な発見

• ResNet50を用いたImageNet1kでは、3.3MBのモデルサイズでトップ-1精度74.5%を達成し、29.4倍の圧縮率を実現した。
• MobileNet-v1では、0.72MBのモデルサイズ（22.4倍圧縮）でトップ-1精度63.9%を達成し、最先端の手法と比較して精度で6.8%向上、モデルサイズ削減率で33%向上した。
• DistilBERTでは、11.8倍の圧縮率を達成し、GLUEベンチマークで精度低下はわずか1.1%にとどまり、GOBOを圧倒的に上回った。
• DKMベースの圧縮では、訓練と推論の間で精度低下が0.2%〜1.7%にとどまり、微分可能訓練と離散的推論のギャップによるものであり、最小限かつ管理可能であった。
• 4/4ビット構成（トランスフォーマー層に4ビット、埋め込みに4ビット）は、すべてのBERTバージョンで1ビットおよび2ビット構成を上回り、特に圧縮が難しいモデル（例：MobileBERT）で顕著な優位性を示した。
• DKMは、特別なトレーニングループや追加正則化なしに、滑らかな収束と段階的な重みクラスタリングを実現した。