[ACM SAC 2025] Advanced Knowledge Transfer: Refined Feature Distillation for Zero-Shot Quantization in Edge Computing

https://arxiv.org/abs/2412.19125 (Accepted)

Abstract

기존 Zero-Shot Quantization(ZSQ, Data-Free Quantizaton) 분야에서는 full-precision(FP) Model로부터 높은 quality의 데이터를 생성하는 데 초점을 두는 연구가 진행되고 있음. 하지만, low-bit(높은 압축률) 환경에서 Quantized Model을 학습할 때는 Quantized Model이 정보 수용량 관련 한계를 갖기 때문에 데이터를 생성하는 기법만으로는 적절한 학습이 이루어지지 않음. 이러한 한계를 개선하기 위해 본 논문에서는 Quantized Model을 효과적으로 학습하기 위한 AKT(Advanced Knowledge Transfer) Method를 제안함. AKT Method는 feature distillation과정에서 spatial 단위 정보와 channel 단위 정보를 융합하여 distillation 함으로써 low-bit Quantized Model 학습에서 gradient exploding 문제를 해결함.

실험결과 Cifar-10 및 Cifar-100 데이터셋에서 우수한 성능을 보였으며, 특히 cifar-10 데이터셋에서 3,5bit에서 State-of-the-Art 성능을 달성함.

 

 

Contributions

1. Zero-shot Quantized Model 학습 시 channel 및 spatial information을 모두 활용하는 것은 학습에 매우 중요한 영향을 미치며, 이계도함수 기반 곡률(second-order culvature)을 통해 quantization과정에서 성능을 유지하는 데 매우 효과적인 방식임을 분석함.

 

2. Channel information과 Spatial information을 융합하는 것은 ZSQ 분야에서 최초의 접근이며, channel information과 spatial information 모두 quantization에서 매우 중요한 요소임을 처음으로 강조함.

 

3. 위 분석을 통해 AKT(Advanced Knowledge Transfer)라는 method를 제안함으로써 low-bit quantized model에서 SOTA 성능을 달성함.

 

4. AKT method는 기존에 널리 연구되던 data generation방식이 아닌 모델의 training을 개선한 방식이기 때문에 다른 data generation방식에 범용적으로 활용될 수 있음.

 

 

Introduction

Deep neural network(DNN)의 발전에 따라 최근에는 다양한 분야에서 이를 응용하고 적용하고 있음. 하지만, 높아진 성능에도 불구하고, DNN은 막대한 양의 컴퓨팅자원을 필요함에 따라 실제 적용을 위해서는 많은 제약이 존재함. 특히, IoT나 모바일 디바이스와 같은 edge computing환경에서는 이러한 한계가 특히 치명적으로 작용함. 이를 개선하기 위해 AI Model을 경량화하는 많은 기법들이 연구되고 있음 (Ex. 지식증류, Pruning, Quantization 등)

본 논문에서는 이러한 기법 중 Quantization에 대해 주목하며, 특히 학습데이터의 접근이 불가능한 매우 제한적인 환경에서의 Quantization (ZSQ)에 대해 새로운 학습방안을 제안함.

 

기존 ZSQ(Zero-Shot Quantization) 연구에서는 학습이 완료된 Full-Precision(FP) 모델의 정보를 Quantized 모델로 전이하는 knowledge distillation 방식을 주로 활용함. 하지만, knowledge distillation 방식은 학습 데이터가 필요한데, ZSQ 환경에서는 학습 데이터에 접근할 수 없으므로 이를 활용할 수 없음.

이를 해결하기 위해, FP 모델의 정보를 활용해 임의로 학습 데이터를 생성하고, 생성된 데이터를 사용하여 Quantized 모델을 학습하는 방식을 주로 활용함.

최근 연구로는 이러한 데이터 합성분야에 초점을 두어 활발한 연구가 진행되고 있으며, ZSQ분야의 Quantization 연구는 표 1과 같이 정리할 수 있음.

 

 

표1을 보면, training 관련한 접근보다 data generation을 위한 연구가 훨씬 활발히 이루어지고 있음을 확인할 수 있음.

 

현재 data generation방식은 5-bit환경과 4-bit환경에서는 FP model(32-bit)과 비교했을 때에도 성능하락이 거의 일어나지 않으면서 Quantized Model을 성공적으로 학습할 수 있음. 하지만. 3-bit (low-bit) 환경에서는 성능이 매우 급격하게 감소하면서 학습에 어려움을 겪고 있음.

 

이는 단순히 data를 generation 하는 접근만으로는 Quantized Model에 충분하지 않으며, 학습과 관련된 접근이 필요함을 암시함.

 

따라서, 본 논문에서는 low-bit quantization에서 성능하락에 영향을 미치는 요인들에 대해 다음과 같이 분석하였음.

1) low-bit quantized model은 정보 수용량이 매우 적기 때문에 기존의 dsitillation기법을 활용하는 경우 효율적으로 정보를 받아들이지 못함. 이는 학습 중 gradient exploding issue로 이어짐.

2) 현재의 knowledge distillation기법은 low-bit quantized model학습에 최적의 방식이 아님. 

 

이러한 분석을 토대로 knowledge distillation 기반 학습을 진행 시, 제한된 정보수용량을 갖는 quantized model을 효과적으로 학습하기 위한 AKT method를 제안함.

 

 

 

Related Work

Quantization

Quantization은 연속적인 실수형 값들을 이산적인 integer value로 변환하는 과정을 의미함. Quantization은 변환기준에 따라 uniform or non-uniform 혹은 symmetric or asymmetric으로 구분할 수 있음. 본 논문에서는 image proccessing에 적합한 asymmetric quantization을 활용하였음.

Quantization 수식은 아래와 같음.

 

 

 

1번 수식은 zero-point $z$로부터 최솟값과 최댓값을 뺀 범위(range)를 scaling 하는 수식이며, 이 과정을 거친 후, 수식 2에서는 floating value를 integer valu $q$로 변환하는 과정을 거침. 마지막으로, quantized integer $q$를 다시 reconstruction 하기 위해 수식 3과 같은 dequantization과정을 거침.

 

Zero-Shot Quantization (ZSQ)

Introduction에서 언급한 대로, ZSQ는 현재 2가지의 주된 연구방향으로 구분되어 있음.

첫 번째로, Data Generation Methodology에서는 최근까지 정교한 데이터를 합성하기 위한 연구가 진행 중이며, 대표 방식으로는 DSQ부터 Qimera, AdaDFQ 등 많은 접근이 존재함. 이러한 방식은 low-bit quantization에서는 성능 관련한 한계가 존재함.

 

두 번째로, Training Methodology에서는 SQuant, AIT 등의 방식이 대표적인 접근으로 여겨지고 있음. 하지만 이 방식들 역시 low-bit quantization에서의 성능하락 원인에 대해서는 직접적으로 분석하지 못하였으며, 더욱 체계적인 분석이 필요함.

 

 

 

Preliminaries

Knowledge Distillation

Knowledge distillation은 학습에 2가지 모델을 활용하는데, 이미 잘 학습되어 있고 모델 크기가 매우 큰 모델을 teacher model로 두고 사용자가 실질적으로 활용하는 경량화된 가벼운 모델을 student model로 구분함.

teacher model의 output은 student model을 학습할 때 label로 활용하는 방식임.

 

ZSQ분야에서는 Quantization이 되어 있지 않은 본래 모델 (FP model)을 teacher model로 활용하며, Quantized Model을 student model로 활용함.

 

Quantized model은 표현범위가 축소된 상태이기 때문에 그대로 inference 하게 된다면 매우 낮은 성능을 보이며, 추가적인 재학습이 필요함. 따라서 Knowledge distillation으로 추가학습을 진행

 

Knowledge distillation은 teacher model의 output과 student model의 output을 가지고 loss를 구하는 logit distillation과 모델의 feature map(activation map)으로 loss를 구하는 feature distillation으로 구분할 수 있음.

기본적으로 knowledge distillation은 logit distillation만 활용하며, feature distillation은 logit distillation에 추가적으로 활용함. logit distillation과 feature distillation 수식은 아래와 같음.

 

 

 

Logit distillation에서의 $T(x)$는 input data $x$에 대한 teacher model output을 의미하며, $S(x)$는 student model의 output을 의미함. $\tau$는 knowledge distillation에 활용되는 temperature로 smooting을 위한 hyper-parameter임.

$\sigma$는 softmax함수를 통과시킨 logit을 의미하며, $KL$은 Kullback-leibler divergence loss를 의미함.

 

Feature distillation 수식에서의 $fT(x)$는 teacher model의 feature map이며,  $fS(x)$는 student model의 feature map을 의미함. $d$는 각 student와 teacher feature map 간 distance를 의미함. 이 distance를 줄이려고 노력함으로써 두 모델 간의 feature map 차이를 최소화시키는 것이 feature distillation의 목적임.

 

 

 

 

Methodology (AKT)

논문에서 제안하는 AKT method는 아래 3가지 과정으로 구분됨.

1) Dual-Information Decomposition

2) Refined Feature Distillation

3) Advanced Knowledge Transfer

 

AKT method의 overview는 다음과 같음.

 

AKT Overview

 

1) Dual-Information Decomposition

이 과정은 모델로 구성되어 있는 여러 layer에 대한 feature map들을 추출하고, 이를 spatial 한 측면의 정보와 channel 기준의 정보로 분할하는 과정임. feature map을 분할하는 과정은 아래 수식과 같음.

 

 

$x, C, H, W$는 각각 feature map, channel, height, width를 의미하며, $x^2$는 feature map을 제곱한 값을 의미함. 제곱을 하는 이유는 feature 값을 강조하여 반영하기 위함임.

 

spatial 한 정보는 말 그대로 공간적인 정보를 의미하기 때문에 image의 pixel단위로 구분함.

따라서, 수식 6에서는 필요 없는 정보인 channel 정보($C$)는 평균값으로 병합함으로써 channel정보는 한 개의 값으로 축소시키면서 width와 height 정보를 유지함.

 

반면 channel attention을 하는 수식 7에서는 Channel 정보를 유지하되, width와 height정보는 평균값으로 병합함. 이를 통해 spatial attention과는 반대로 channel정보를 중점적으로 반영하는 feature map을 얻게 됨.

 

위처럼 하나의 feature map $x$를 2가지 기준으로 attention 된 map으로 분해하는 과정이 Step 1에 해당하는 dual-information decomposition과정임.

 

2) Refined Feature Distillation

Step 1에 해당하는 과정을 토대로 본 논문에서는 분해된 feature map을 서로 다른 관점에서 통합하는 새로운 loss function을 제안함. 논문에서 제안하는 loss는 refined feature distillation loss이며 $L_RFD$로 표기함. $L_RFD$는 아래 그림과 수식 8과 같이 계산할 수 있음.

 

 

 

 

수식 8과 그림 2는 각 layer의 feature map마다 $L_RFD$를 계산하는 과정에 해당함. 수식에서 첫 번째 $KL()$은 teacher model과 student model로부터 분해된 spatial attention map을 비교해서 차이를 최소화하는 과정이며, 두 번째 $KL()$은 channel 단위의 attention map을 비교함. 이러한 과정을 $N$개의 feature map ($N$개의 layer) 만큼 반복하고 평균을 구함으로써 최종적인 $L_RFD$를 계산할 수 있음.

 

$L_RFD$는 feature distillation을 개선하기 위한 새로운 loss이며, 뒤에 이어지는 step 3(Advanced Knowledge Transfer) 과정에서 활용됨.

 

3) Advanced Knowledge Transfer

Preliminaries에 언급한 대로 knowledge distillation은 logit distillation과 feature distillation이 구성될 수 있다고 설명함.

본 논문에서는 quantized model 학습을 위해 logit distillation과 feature distillation을 모두 활용하며, 본 논문에서 제안하는 최종 loss 수식은 수식 9와 같음.

 

 

$L_AKT$ 수식에서 $\alpha$는 feature distillation과 logit distillation을 어느 정도 비율로 반영할지 결정하는 hyper-parameter이며, 실험에서는 default로 0.5로 지정하였음.$L_RFD$는 step 2에서 언급한 새로운 feature distillation loss를 활용하였으며, $L_KL$은 기존의 일반적인 logit distillation에 해당함.

 

 

 

 

Experimental Results

실험은 ZSQ에서 base-line으로 활용되는 GDFQ부터 최신 기법에 해당하는 HAST, AdaDFQ 등을 활용하였으며, Cifar-10 및 Cifar-100 데이터셋에서 Resnet-20 모델을 기준으로 실험함.

 

실험결과는 Table2와 같음.

 

 

Data generation method인 GDFQ와 HAST, AdaDFQ에 본 논문의 제안방식인 AKT를 적용하여 학습하는 경우 거의 모든 경우에서 정확도 향상의 효과를 볼 수 있었으며, 특히 AdaDFQ에 AKT를 적용하였을 때 3-bit환경에서 매우 큰 성능향상이 일어났음을 확인할 수 있음.

 

 

 

Table 2에서 사용된 3가지 method 뿐 아니라, table 3에서처럼 다른 기존의 SOTA method들과 비교한 경우에도 AKT는 매우 우수한 성능을 보임. 실험조건에서는 코드가 공개되어 있는 AdaDFQ를 기준으로 실험을 하였으며, 코드가 공개되어 있지 않은 RIS나 TexQ에도 AKT를 적용한다면 더욱 큰 성능향상이 이뤄질 수 있을 것으로 예상함.

 

 

Ablation Study

RFD Loss의 효과를 분석하기 위해 논문에서는 추가적인 ablation study를 진행함. Table 4는 attention 정보에 따른 정확도임.

 

 

Channel 정보만 사용하거나, spatial 정보만 활용하는 기존 feature distillation과 달리 RFD loss를 활용하였을 때 가장 좋은 성능을 보임을 확인함.

 

추가적으로 quantized model의 학습 안정성을 측정하기 위해 논문에서는 이계도함수 기반의 hessian trace지표를 활용함.

 

Hessian Trace: gradient를 2번 미분한 값으로 모델의 training stablility를 측정할 수 있는 지표임. 낮은 값일수록 안정적으로 학습이 진행된다는 의미이며, 성능향상에 효과적임.

 

 

 

 

figure 3을 보면 논문에서 제안하는 RFD loss (blue graph)가 기존 방식에 비해 낮은 값을 보여 안정적으로 학습되고 있음을 확인할 수 있음.

특히, 3-bit 환경에서는 100 epoch 이후로 기존 방식들은 급격하게 불안정한 학습이 되는 경향을 보이지만, RFD loss를 활용하는 경우에는 매우 안정적인 학습이 진행됨을 확인할 수 있음.

 

 

Discussion

channel or spatial 중 하나의 정보에만 attention 하는(one-sided attention) 기존의 feature distillation은 학습에 상당한 loss를 불러일으킴. 논문에서 제안하는 AKT는 이러한 한계를 개선하기 위해 channel 정보와 spatial 정보를 효과적으로 mapping 하는 방식으로 (dual-perspective attention) ZSQ분야의 새로운 접근임. 이러한 방식은 low-bit quantization 환경에서 제한된 정보수용량에도 효과적인 지식전이(knowlegde transfer) 학습이 가능하게 함.

 

하지만, AKT 방식도 한계가 존재함.

첫 번째로, AKT 방식은 feature distillation을 활용하는 방식이기에 기존 학습방식들과 같이 input data에 대한 quality에 의존적임. 특히, feature distillation에 활용되는 feature map(activation map)은 모델의 고정정보인 weight와 데이터에 따른 변동정보인 input에 따라 다른 값들이 뽑히기 때문에 특히 input의 영향이 큼. 따라서, data의 quality가 어느 정도는 보장되어야 좋은 성능을 보인다는 한계가 있음,

 

결론적으로, 생성되는 데이터의 quality는 AKT 방식에서 중요한 영향을 미치며, 데이터의 qualtiy와 학습방식에 대한 융합연구가 중요함.

 

 

 

Conclusion

본 논문에서는 ZSQ의 knowledge distillation에서 teacher model과 student model 간 information gap (quantization noise)를 줄이기 위한 방식을 제안하였음. AKT는 3가지 step을 통해 quantized model의 제한적인 정보 수용량에도 불구하고, full-precision model의 정보를 효율적으로 quantized model에 전이시킴. 특히, low-bit quantization에서는 feature distillation이 좋은 효과를 보이며 논문에서 제안하는 방식은 학습방식이기 때문에 다른 data generation method보다 범용적으로 활용될 수 있음.

 

 

 

(제 논문이기 때문에 논문에 대한 review는 작성하지 않았습니다. 질문이 있으신 경우 댓글로 남겨주시면 답변드리도록 하겠습니다!)