[논문 리뷰] DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

Minsik Cho, Keivan Alizadeh-Vahid외 2명|arXiv (Cornell University)|2021. 08. 28.

Advanced Image and Video Retrieval Techniques참고 문헌 39인용 수 5

한 줄 요약

이 논문은 신경망 가중치와 클러스터 중심점을 공동 최적화할 수 있도록 허용하는 미분 가능 k-means 클러스터링 레이어인 DKM을 제안한다. 이를 통해 종단 간 학습이 가능해져 최소한의 정확도 손실로 최신 기술 수준의 모델 압축을 달성한다. k-means를 미분 가능한 어텐션 메커니즘으로 재구성함으로써 DKM은 고정밀도, 저비트 양자화를 가능하게 하여 MobileNet-v1에 대해 22.4배의 압축률을 달성했으며, ImageNet의 top-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 클러스터링의 비가역성으로 인해 네트워크 가중치와 클러스터 중심점을 공동 최적화할 수 없는 문제를 해결한다.
• 기존 손실 함수나 모델 아키텍처를 수정하지 않고도 학습 중 태스크 인식 기반 클러스터링을 가능하게 한다.
• 컴퓨터 비전 및 NLP 모델에서 모두 뛰어난 압축-정확도 트레이드오프를 달성하기 위해 미분 가능한 클러스터링을 통해 초과 압축을 실현한다.
• 다양한 차원의 k-means 클러스터링을 지원하는 즉시 사용 가능한, 파라미터가 없는 레이어를 제공한다.

제안 방법

• 이산 클러스터 할당을 미분 가능한 어텐션 메커니즘으로 대체하는 미분 가능한 k-means 클러스터링 레이어인 DKM을 제안한다.
• 입력 가중치와 학습 가능한 클러스터 중심점 간의 도트 프로덕트 어텐션 메커니즘을 사용하여 k-means 클러스터링을 부드러운 할당 문제로 재구성한다.
• 온도 조절 소프트맥스를 사용하여 클러스터 할당의 날카움을 제어함으로써 백프로파게이션 중 기울기 전파를 가능하게 한다.
• 네트워크에 삽입 가능한 경량이며 파라미터가 없는 레이어로 DKM을 적용하여 네트워크 가중치와 클러스터 중심점을 공동 최적화한다.
• 학습 중에는 미분 가능한 어텐션 메커니즘을 통해 기울기 업데이트를 유지하면서 추론 시에는 가장 가까운 중심점으로 예측을 고정한다.
• 기존 DNN에 아키텍처 변경 없이 DKM을 통합함으로써 원래의 손실 함수와 학습 파이프라인을 그대로 유지한다.

실험 결과

연구 질문

• RQ1 k-means 클러스터링을 어떻게 미분 가능하게 만들 수 있을까? 이를 통해 딥 네트워크와 종단 간 학습이 가능해지는가?
• RQ2 미분 가능한 k-means 클러스터링은 기존의 k-means나 이전의 미분 가능한 접근 방식보다 더 나은 압축-정확도 트레이드오프를 달성하는가?
• RQ3 DKM은 기존에 압축이 어려운 것으로 알려진, 매우 파rameter 효율적인 모델인 MobileNet-v1을 효과적으로 압축할 수 있는가?
• RQ4 DKM은 NLP 모델, 특히 DistilBERT와 같은 모델에서 저정밀도 양자화 조건에서도 잘 작동하는가?
• RQ5 DKM은 임베딩 레이어와 트랜스포머 레이어 등 서로 다른 레이어 간에 압축률을 세밀하게 제어할 수 있으며, 정확도 저하 없이 이를 달성할 수 있는가?

주요 결과

• ResNet50를 사용해 ImageNet1k에서 74.5%의 top-1 정확도를 달성했으며, 모델 크기는 3.3MB로 압축률이 29.4배에 이르렀다.
• MobileNet-v1에 대해 DKM는 0.72MB의 모델 크기로 63.9%의 top-1 정확도를 기록했으며, 22.4배의 압축률을 달성했다. 이는 최신 기술 수준의 방법들보다 정확도에서 6.8% 향상되고 모델 크기 감소율이 33% 높았다.
• DistilBERT에 대해 DKM는 11.8배의 압축률을 기록했으며, GLUE 벤치마크에서 정확도 손실이 오직 1.1%에 불과했다. 이는 GOBO보다도 더 뛰어난 압축률과 정확도를 달성했다.
• DKM 기반의 압축은 학습과 추론 사이에 정확도 손실이 0.2%에서 1.7% 수준으로 발생했으며, 이는 미분 가능한 학습과 이산적 추론 간 격차로 인한 것으로, 매우 미미하고 관리 가능한 수준이었다.
• 4/4 비트 구성(트랜스포머 레이어에 4비트, 임베딩에 4비트)이 모든 BERT 변종에서 1비트 및 2비트 구성보다 뛰어난 성능을 보였으며, 특히 압축이 어려운 모델인 MobileBERT에서 두드러진 성능 향상을 보였다.
• DKM는 학습 중에 부드러운 수렴과 점진적인 가중치 클러스터링을 가능하게 하여, 맞춤형 학습 루프나 추가 정규화 기법이 필요로 하지 않았다.