KDST

이번 글에서는 CVPR 2021에 oral로 선정된 “Representative Batch Normalization with Feature Calibration”이란 논문을 소개합니다. 본 논문은 기존의 Batch normalization의 문제를 분석하고 이를 보완하는 Representative Batch normalization을 제안합니다. 

Introduction

Batch Normalization의 효과를 감소시키는 요인

 Batch Normalization(BN)은 convolution 연산의 결과로 나온 feature들을 mini-batch의 통계 정보를 사용하여 normalize 된 분포로 제한하는 과정을 수행합니다. 이 과정은 학습의 어려움을 낮춰 CNN 모델의 성능을 가져왔습니다. 하지만, BN이 가정하는 "서로 다른 인스턴스 feature의 채널별 분포가 모두 같다"는 2가지 문제점이 존재합니다.

1. Training 단계에서의 mini-batch statistics와 testing시 사용하는 running statistics 사이의 불일치.
   • training 단계에서 mini-batch 단위로 누적 계산되는 statistic은 추론 시에 사용할 running statistics에 완벽하게 적합하지 않습니다.
2. testing dataset의 인스턴스가 항상 훈련 세트의 분포에 속하지는 않을 수 있음.

이러한 두 가지 불일치 문제는 BN의 효과를 감소하게 합니다.

기존의 연구들

몇몇 연구들에서는 이를 해결하기 위해 normalize 단계에서 mini-batch 단위의 statistic 정보 대신 instance 단위의 statistic 정보를 활용하였습니다. 하지만 이는 batch 정보의 부족으로 학습의 불안정을 발생시켜 많은 경우에서 기존의 BN보다 안 좋은 성능을 보였습니다. 또 다른 연구에서는 여러 개의 normalize 기법을 조합하거나 attention 메커니즘의 도입으로 해결하였으나 많은 overhead를 도입하여 실용성을 떨어트리게 됩니다. 따라서 본 논문은 적은 비용으로 feature의 normalize 과정에서 mini-batch 이점과 instance 별 representation을 향상시키는 방법을 제안합니다.

Statistic의 불일치로 인한 문제 

Batch Normalization(BN)은 feature standardization과 affine transformation 동작으로 이루어져 있습니다. 본 논문은 centering과 scaling 동작으로 이루어진 standardization에 집중하였으며, centering과 scaling 과정은 각각 mini-batch의 statistics(mean, variance)를 기반으로 feature들이 zero-mean property(평균 값이 0)와 unit-variance(분산 값이 1)를 가지도록 합니다. 이때 testing instance와 training 시에 계산된 running mean, variance의 불일치는 몇 가지 문제를 발생시킵니다.

1. running mean의 불일치
   • running mean > instance mean: 유익한 representation의 손실.
   • running mean < instance mean: feature에 noise의 증가.
2. running variance의 불일치
   • 특정 채널의 값의 scale을 너무 크게 혹은 작게 만들어 feature의 분포를 불안정하게 만듦.
   • (예를 들어, 특정 channel에서의 running variance의 불일치는 해당 채널의 scale을 다른 channel에 비해 크게 만들어 다음 convolution에서 연산 결과가 그 채널에 지배적(dominant)이게 만듭니다.)

제안하는 Representative Batch Normalization은 feature의 채널별로 3개의 weight를 추가하여 centering, scaling calibration 과정을 수행합니다. 이는 기존의 BN 보다 training loss와 test accuracy를 작게 만들며 기존 BN을 손쉽게 대체 할 수 있습니다.

Proposed Method

Revisiting Batch Normalization

제안하는 Representative Batch Normalization(RBN)에 앞서 기존의 Batch Normalization(BN)의 standardization, affine transformation 과정을 수식화하면 아래와 같습니다.

Centering : $ X_m =X-E(X) $

Scaling : $ X_s= \frac{X_m}{\sqrt{Var(X)+\epsilon}} $

Affine : $ Y = X_s\gamma + \beta $

$X\in \mathbb{R^{N\times C\times H\times W}}$는 input feature이며 $E(X), Var(X)$는 각각 running mean과 variance를 의미합니다. training 과정에서 running mean과 variance는 dataset으로 부터 mean과 variance를 누적하여 계산됩니다. 이 과정은 아래와 같습니다.

$\mu _B=\frac{1}{NHW}\sum_{n=1}^{N}\sum_{h=1}^{H}\sum_{w=1}^{W}X_{(n,c,h,w)}$

$\sigma^{2}_B=\frac{1}{NHW}\sum_{n=1}^{N}\sum_{h=1}^{H}\sum_{w=1}^{W}(X_{(n,c,h,w)}-\mu_B)^2$

$E(X)\Leftarrow mE(X)+(1-m)\mu_B$

$Var(X)\Leftarrow mVar(X)+(1-m)\sigma^2_B$

Representative Batch Normalization

• Statistics for Calibration

제안하는 RBN에서 Batch Normalization(BN)의 running statistic을 보정하기 위해 사용한 instance-specific feature의 statistic은 channel 별 정보를 사용하였으며 이는 다음과 같습니다.

$\mu_c = \frac{1}{HW}\sum_{h=1}^{H} \sum_{w=1}^{W} X_{(n,c,h,w)}$

$\sigma ^2 _c = \frac{1}{HW} \sum_{h=1}^{H} \sum_{w=1}^{W}(X_{(n,c,h,w)}-\mu _c)^2$

• Centering Calibration

BN의 Centering 과정에서  feature 값 중 running mean보다 큰 부분은 유지되어 다음 레이어로 넘어가게 됩니다. 이때 잘못된 running mean은 feature의 유용한 정보를 지우거나 추가적인 noise를 추가하게 됩니다. 이를 해결하기 위해 centering calibration은 BN의 centering 과정 전에 feature의 추가적인 statistic($K_m$)을 채널 별 learnable weight($w_m$)을 통해 feature를 calibration 합니다. 수식으로 표현하면 다음과 같습니다. (논문의 실험에서 $K_m$은 $\mu_c$를 기본값으로 사용하였습니다.)

$X_{cm(n,c,h,w)}=X_{(n,c,h,w)} + w_m\odot K_m$ 

$\odot$: $w_m$과 $K_m$을 같은 shape로 만들고 dot product 수행.

Centering calibration이 적용된 X와 적용되지 않은 X의 centering operation 결과 차이를 수식화하면 아래와 같습니다. $w_m$ > 0 그리고 $K_m > E(X)$와 같은 조건에서 centering calibration은 feature의 유용한 정보를 살려주며 $w_m$ < 0 그리고 $K_m > E(X)$와 같은 조건에서는 feauture의 noise를 줄여주는 방식으로 동작하게 됩니다.

$X_{cal}-X_{no} = w_m\cdot (K_m-E(X))$

• Scaling Calibration

유지할 특징을 조절하는 centering과 다르게 scaling 과정에서는 feature의 명암을 조절합니다. 이때 부정확한 running variance는 특정 channel의 강도를 크게 만들게 됩니다. 이를 해결하기 위해 scaling calibration은 scaling operation 과정 후에 instance statistic($K_s$)을 채널별 learnable weight($w_v,w_b$)를 사용하여 feature를 calibration합니다. 

$X_{cs(n,c,h,w)}=X_{s(n,c,h,w)}\cdot R(w_v\odot K_s+w_b)$

$R()$: out-of-distribution feature 억제용 restricted function($0<R()<1$)

Scaling calibration 후 feature의 variance는 아래와 같이 표현됩니다.

$Var(X_{cs})=Var(X\cdot R(w_v\cdot K_s+w_b))$

Scaling calibration은 feature의 variance를 작게 만들고 이를 통해 채널 간의 안정적인 distribution을 유지하게 합니다. $w_v$가 작을수록 $Var(X_{cs})$는 작아지며 $w_b$를 통해 restricted function의 제한 위치를 조정하게 됩니다.

Experiments

제안하는 RBN은 분류에서의 다른 방식들에 비해 더 나은 성능을 보였으며 scalling calibration이 적용되었을 때 채널간의 분포를 보다 안정적이게 만들어주는 것을 확인 할수 있습니다.

Conclusion

Representation Batch Normalization(RBN)은 기존 BN의 장점인 Batch statistic 정보를 유지하며 적은 추가적인 파라미터를 통해 효과적으로 instance의 statistic을 적용합니다. centering calibration은 feature의 유용한 representation을 강화하고 noise를 감소시킵니다. Scaling calibration은 feature의 강도를 제한함으로써 채널간의 안정적인 feature 분포를 형성하게 합니다. RBN은 BN의 centering/scaling operation을 사용하는 기존의 방식들에도 쉽게 적용 가능하다는 장점 또한 가집니다.

이상으로 논문 리뷰를 마칩니다.

자세한 내용이나 궁금하신점은 full paper를 참고하시고나 댓글로 남겨주세요.

감사합니다.

Dongjin Kim

이번에는 CVPR 2020에 Oral session에 발표된 논문인 Revisiting Knowledge Distillation via Label Smoothing Regularization을 리뷰하려고 합니다. 논문은 링크를 참조해주세요

Background

Geoffrey Hinton 이 2014년 NeurIPS에 발표한 Distilling the Knowledge in a Neural Network 이후로 분야가 개척되어왔습니다.

Knowledge Distillation에서 Hinton이 발표한 방식은 pretrained teacher model의 logits과 student모델의 logits의 soft-target을 구해서 KL Divergence를 취해서 서로 분포간의 거리를 좁히는 방식으로 학습이 진행되어 왔습니다.

$$ L=L_{CE}(y|x;\theta)+L_{KD}(y,y'|x;\theta_{stu},\theta_{teacher})$$

그 이후 많은 Knowledge Distillation Method(이하 KD)들이 발표되고 있으며, 기존의 logits(neural network output)들을 가지고 KD를 어떻게, 무엇을 하는 지에 대해서 연구가 진행되고 있습니다.

Knowledge Distillation에는 구성요소가 여러가지가 있는데, Knowledge 종류, learning scheme, architecture 등등이 있습니다. 간단히만 설명하자면, Knowledge 종류에는 logits을 활용한 response-based knowledge, intermediate layer output을 활용하는 feature-based knowledge, 그리고 instance를 활용한 relation-based knowledge가 존재합니다.

또한, learning scheme은 Distillation을 누구와 언제 하는지에 관한 것인데, pretrained model을 갖는 것은 offline KD, teacher와 student가 동시에 학습하는 것은 online KD, teacher와 student가 같은 것이 self KD라고 할 수 있습니다.

이중에서 논문은 response-based knowledge, 즉 logits을 활용하여 teacher와 student가 같은 self knowledge distillation에 대해서 집중적으로 다룰 예정입니다.

- Response based Knowledge Distillation

logits $z$ vector를 활용하여 $p(z_i,T)=exp(z_i/T)/{\sum_{j}exp(z_j/T)}$인 soft-target은 teacher로 부터 student에게 informative dark knowledge를 전달한다는 것이 일반적인 통념이었습니다.

하지만, 저자는 이러한 common belief에 대해서 혹시 regularization method중 하나인 label smoothing regularization과 유사하게 생겼다는 의문을 품고 2가지 실험을 구성하여 이를 검증해보려고 합니다. 만약 두 실험이 성립하는 경우, dark knowledge가 정보를 전달하기 보다는 그냥 label들을 정규화(regularize)하는 method중에 하나임을 실험적으로 증명할 수 있기 때문입니다.

1번째 실험은 student model로 사용되는 light model과 teacher model로 사용되는 deep model의 역할을 바꾸어 학습을 시켜보는 것입니다. (Reversed KD, Re-KD)

2번째 실험은 teacher model을 덜 학습시킨다음(low accuracy) student에게 knowledge transfer를 진행하는 것입니다. (Defective KD, De-KD)

Cifar100 data를 활용한 해당 Exploratory Experiment의 결과는 다음과 같습니다.

1번째 실험은 Re-KD와 일반적인 Normal KD의 결과 비교 입니다.

표를 보면 T->S 에서 Student의 성능향상이 뚜렷하게 높긴하지만, Re-KD에서 S->T에서의 학습 향상을 볼 수 있습니다.

2번째 실험은 De-KD의 결과입니다.

개인적으로는 이 실험결과가 엄청 놀라웠는데요. Teacher를 poorly-trained model로 구성하여 student를 학습시킨 결과 student가 눈에 띄게 증가하는 것을 확인할 수 있었기 때문입니다. 이를 통해서 

실험 1의 Normal KD와 실험2의 De-KD를 비교해보면 대체로 NormalKD만큼은 아니지만 De-KD에서도 성능향상이 이루어진다는 사실을 확인할 수 있었고, 심지어는 ResNet50을 teacher로 활용하는 MobileNetV2를 학습시키는 경우, normal KD보다 높은 성능을 보이기도 하는 등, 실제로 dark knowledge가 아니라 regularizer의 역할을 할 수도 있을 것 같다 라는 생각이 드는 실험이었습니다. 2번째 실험과 1번째 실험의 normal KD를 같이 비교하지 않은점은 좀 아쉬웠습니다.

Motivation

이러한 실험적 확인 이후, 본 논문에서는 KD를 regularization method중에 하나인 Label Smoothing Regularization(LSR)과 유사하기 때문에 이를 수학적으로 유사함을 증명하려고하였습니다. Normal Classification에서 output distribution은 $p(k|x)=p(k)=softmax(z_k)=exp(z_k)/{\sum_{i=1}^{k}exp(z_i)}$ 를 사용하여 ground truth에 대해서 cross-entropy loss를 활용합니다.

$$H(q,p)=-\sum^{K}_{k=1}{q(k|x)log(p(k|x))}$$

Label Smoothing Regularization (LSR)에서는 cross-entropy loss에서 사용되는 ground truth q(k|x)를 바꾸어 q'(k|x)를 사용합니다. 정답 label을 확률 1로 보는 것보다 낮은 값을 선택한 후 남은 값을 다른 label에 분배하는 것입니다.u(k)=1/K로 uniform distribution입니다.

$$q'(k|x)=(1-a)q(k|x)+au(k)$$

Knowledge distillation에서의 loss function은 다음과 같습니다.($\tau$는 temperature)

$$L_{KD}=(1-a)H(q,p)+aD_{KL}(p^t_{\tau},p_{\tau}) $$

LSR에서의 loss function은 다음과 같습니다.

$$H(q,p)=-\sum^{K}_{k=1}{q'(k)log(p(k))}$$

$$L_{LS}=(1-a)H(q,p)+aD_{KL}(u,p)$$

이 두식을 정리해서 풀면 다음과 같이 굉장히 유사한 식으로 구성이 됩니다.

q'(k) in LSR: $q'(k)=(1-a)q(k|x)+au(k)$

q'(k) in KD: $q'(k)=(1-a)q(k|x)+ap^{t}_{\tau}(k) when \tau =1$

이를 통해서 $\tau$를 조정하면 $p^{t}_{\tau}(k)$가 $u(k)$와 유사하기 때문에 이를 저자는 KD가 LSR의 특별한 케이스로 보고 있습니다. 왜냐하면 temperature $\tau$가 커지면 커질 수록 유사한 값이 나오기 때문입니다.

이러한 사실들을 기반으로 그렇다면 KD가 regularizer의 한 종류로 볼 수 있기 때문에 이를 통해서 Teacher-free Knowledge Distillation이라는 새로운 방법론을 제시합니다.

Proposed Method

1) Teacher-free Knowledge Distillation (self-training)

간단한 pipeline은 student를 학습시킨 후, 그 student model을 복사하여 student model을 teacher로 사용하는 방법입니다.  아래 사진에서 teacher model로 student 자신을 사용하는 방법입니다. 나머지는 normal KD와 동일합니다. 개인적으로는 이 부분은 2015년 이후 한 번쯤은 생각해볼 방법론이라 간단하게 넘어가겠습니다.

2) Teacher-free Knowledge Distillation (self-regularization)

개인적으로는 이 부분이 엄청 많이 추가가 된 것은 아니지만 성능 향상이 나름 인상적이었습니다. LSR의 방식을 도입한 후, 해당 label을 soft-target을 취해서 student의 logits과 KD Loss를 취하는 방식입니다. 한마디로 ground truth에 대해서 fake teacher처럼 soft를 취해서 KD를 진행한다는 것이죠. 

KD에서 사용하는 fake teacher에 대한 output distribution은 다음과 같습니다.

$$ p^{d}(k)=\begin{cases}a &if     k = c \\ {1-a}/{k-1} & if     k \neq c \end{cases} $$

Pipeline을 보기 쉽게 설명하면 다음과 같습니다.

Teacher를 쓰는 것이 아니라 가상으로 두고 target값을 label smoothing(by softmax with temperature)하여 KD Loss를 취할 수 있게 하는 것이 이 논문의 핵심 방법론이라고 볼 수 있습니다. 

Experiment

cifar100에 대한 self-training Tf-KD의 결과는 다음과 같습니다.

사실 이 부분은 더 큰 모델을 사용할 필요 없이 Tf-KD_self가 유사하거나 높은 성능을 냄을 확인함으로서 더이상 big teacher가 필요 없다는 데 의의가 있는 실험결과입니다. 통념적으로 믿어온 잘 가르치는 선생님은 친구이기도 하다는 것과 같습니다.

두번째 Tf-KD_reg 방법에 대한 실험 결과는 다음과 같습니다.

Normal KD에 대해서 유사한 성능을 보여주고 있고, LSR에 대해서는 더 좋은 성능을 통해 Soft-target을 취했을 때 성능 향상에 도움이 됨을 확인할 수 있습니다. 이 방법의 의의는 선생님 조차도 없이 스스로 distillation을 통해 성능 향상을 꾀할 수 있다는 것을 확인하였고, 엄청 큰 모델을 스스로 학습시킬 때에도 자유롭게 사용가능하다는 점을 생각해보면 의미있는 실험결과임을 알 수 있습니다.

Conclusion

결론은 Dark Knowledge라고 부르는 부분은 사실 Regularization Term으로 볼 수 있다 라고 말하고 있습니다. 이를 통해서 스스로 distillation 하는 self-distillation이 teacher-student관계가 존재하는 learning scheme만큼의 성능을 보여줄 수 있다는 것을 실험, 수식으로 설명한 강력한 논문이라고 생각합니다. 하지만, De-KD에서 15%의 accuracy를 갖는 ResNet을 teacher로 하는 결과와 normal KD와 비교해볼 때, 눈에 띄게 성능저하가 존재하는 경우가 있다는 점을 고려해보면 Dark Knowledge가 존재할 수 있지 않을까? 라는 생각이 들었습니다. 

자세한 내용은 논문(full paper)을 참조하거나 댓글로 남겨주시기 바랍니다

Minsoo Kang