This is a Korean review of
"CrossKD: Cross-Head Knowledge Distillation for Object Detecton"
presented at CVPR 2024.
TL;DR
- CrossKD라는 효과적인 prediction mimicking distillation scheme을 제안하여, student의 dtection head의 중간 feature를 teacher의 detection head에 전달함. 이러한 cross-head prediction은 teacher의 prediction을 모방하게 됨.
- 본 방식은 student head의 annotation과 teacher prediction으로부터의 contradictory supervision signal 전달을 완화하여 student의 성능을 향상함.
- CrossKD는 더 task-oriented information (classfication and localization)을 제공하는 prediction mimicking의 장점을 그대로 활용함.
Introduction
Limitation of Conventional Prediction Mimicking
- 기존의 prediction mimicking을 통해 student를 학습할 때는, ground-truth target과 teacher의 prediction을 동시에 모방하게 되는데, student로부터의 ground truth target과 teacher로부터의 distillation target 간의 큰 불일치가 존재함. 이는 그림 2 (a)와 (b)의 초록색 원 비교를 통해서도 확인할 수 있음.
- 결과적으로 student detector는 contradictory learning process를 경험하고, 이는 optimization을 심각하게 방해함.
- 이를 극복하기 위해, 이전 prediction mimicking method는 teacher와 student 간의 중간 정도 discrepancy를 포함한 영역에서 지식 증류를 수행함.
- 하지만, heavily uncertain region은 일반적으로 student에게 도움이 되는 정보를 많이 가지고 있기 때문에 이전 방법들은 이러한 특성을 활용하지 못함.
CrossKD: Cross-head Knowledge Distillation
- Student head의 중간 feature를 teacher head로 전달하여 cross-head prediction을 얻음.
- 얻어진 new cross-head prediction과 teacher prediction 사이에 knowledge distiilation을 수행함.
- $(+)$ cross-head prediction는 teacher detection head와 공유되기 때문에 teacher prediction과 상대적으로 일치함. 즉, discrepancy를 완화할 수 있음.
- $(+)$ 이론적으로 optimal 하고 task-oriented information (classification information과 localization information)을 제공할 수 있음.
Related work
Object Detection
- GFL은 bounding box representation을 probability distribution으로 표현해서, localization ambiguity을 포함할 수 있는 dynamic label assigment method임.
Knowledge Distillaiton Small Object Detection
- HEAD는 student feature를 독립적인 assistant head로 전달해서 heterogeneous teacher-student pair 간의 gap을 줄임. 이와 반대로, CrossKD는 student feature를 teacher에게 전달해서 성능을 개선시킴.
Method
Analysis of the Target Conflict Problem
- 모든 이미지에 특정 category를 할당하는 classification과는 다르게, 최신 detector에서의 label은 보통 dynamically 할당되고 deterministic 하지 않음. 각 location에서의 label을 결정하기 위해 assigner에 의존하지만, detector는 assigner의 label을 정확하게 재생산하지 못하며, 이는 teacher target과 ground truth target 간의 target conflict를 초래함.
- Target conflict를 측정한 그림 3을 통해, teacher와 student가 같은 label assignment strategy (ATSS & GFL - GFL)를 활용하더라도 ground-truth와 distillation target 간의 불일치가 상당히 존재함을 확인할 수 있고, 다른 assigner의 경우 (RetinaNet - GFL)에는 불일치가 더욱 커짐.
- "Distillation under Severe Target Conflict"에서 target conflict가 성능을 악화시킴을 보임.
- 이전의 prediction mimicking method는 teacher와 student prediction 간의 불일치를 최소화하기 위해 아래의 objective function을 활용함.
$$
\mathcal{L}_{\mathrm{KD}}=\frac{1}{|\mathcal{S}|} \sum_{r \in \mathcal{R}} \mathcal{S}(r) \mathcal{D}_{\text {pred }}\left(\boldsymbol{p}^s(r), \boldsymbol{p}^t(r)\right)
$$
- $\mathcal{S}$는 전체 이미지 영역 $\mathcal{R}$에 포함된 각 위치 $r$에 weight를 생성하는 region selection principle이며, teacher와 student 간의 차이가 큰 영역을 down-weighting 하여 target conflict문제를 어느 정도 완화할 수 있지만, heavily uncertain region은 student에게 도움이 되는 많은 정보를 가지고 있어, 이를 무시하면 성능이 크게 하락함.
Cross-Head Knowledge Distillation
- Detection head는 ${C}_{i}$로 표현되는 convolutional layer n개로 이루어져 있으며, ${C}_{i}$로부터의 feature map은 $\boldsymbol{f}_{i}, i\in {1,2,3,..., n-1}$로, $C_{1}$의 input feature map은 $ \boldsymbol{f}_{0}$로, 마지막 convolutional layer $C_{n}$에 의한 prediction은 $ \boldsymbol{p}$로 표현함.
- CrossKD는 student의 중간 feature $ \boldsymbol{f}_i^s, i\in{1,2,...,n-1}$를 teacher detection head의 convolutional layer $C_{i+1}^t$ 에 전달하고 cross-head prediction $\hat{\boldsymbol{p}}^{s}$을 만들어 distillation을 수행함.
- 기존의 KD가 $\boldsymbol{p}^{s}$와 $\boldsymbol{p}^{t}$간의 loss를 계산하는 것과 다르게, CrossKD는 cross head prediction $\hat{\boldsymbol{p}}^{s}$과 $\boldsymbol{p}^{t}$간의 loss를 계산하는 아래의 objective function을 활용함
$$\mathcal{L}_{\mathrm{CrossKD}}=\frac{1}{|\mathcal{S}|} \sum_{r \in \mathcal{R}} \mathcal{S}(r) \mathcal{D}_{\text{pred }}\left(\hat{ \boldsymbol{p} }^s(r), \boldsymbol{p}^t(r)\right)$$
- 복잡한 region selction principle $ \mathcal{S}(\cdot) $를 design 하는 것보다, 전체 prediction map에 대해서 동일한 weight를 취함. 즉, $ \mathcal{S}(\cdot) $ 는 1의 값을 가지는 constant function임.
- Detection loss의 gradient는 student의 전체 head를 통과(yellow의 supervised backprop path)하지만, distillation loss의 gradient는 teacher의 layer을 지나, student의 latent feature를 통과(green의 distillation backprop path)함. 이는 teacher와 student 간의 일관성을 증가시킴.
- 이를 통해, CrossKD는 student의 일부 detection head가 detection loss과만 관련되도록 해서 ground-truth target에 대한 optimization을 개선시킴.
Optimization Objectives
$$ \begin{aligned} \mathcal{L} &= \mathcal{L}_\text{cls}(\boldsymbol{p}_\text{cls}^{s}, \boldsymbol{p}_\text{cls}^{gt}) + \mathcal{L}_\text{reg}( \boldsymbol{p}_\text{reg}^{s}, \boldsymbol{p}_\text{reg}^{gt}) \\&+ \mathcal{L}_{\text{CrossKD}}^\text{cls}(\hat{ \boldsymbol{p}}_\text{cls}^{s}, \boldsymbol{p}_\text{cls}^{t}) + \mathcal{L}_{\text{CrossKD}}^\text{reg}(\hat{ \boldsymbol{p}}_\text{reg}^{s}, \boldsymbol{p}_\text{reg}^t) \end{aligned} $$
- 서로 다른 branch (classification, regression)에 대한 task-specific information을 전달하기 위해 classification에서는 GFL 논문에서 제안된 Quality Focal Loss (QFL)을 활용하고, regression은 형태에 따라 GIoU (RetinaNet과 ATSS 같이 anchor box로부터 bounding box를 regression 또는 FCOS와 같이 point로부터 regression) 또는 KL divergence (GFL은 box location의 distribution을 표현하기 위해 vector를 예측하는 regression)를 사용함.
Experiments
Method Analysis
Positions to apply CrossKD
- $i=0$일 때는, student feature가 student의 head를 거치지 않고, 바로 teacher의 head로 들어가기 때문에, 전체 student head가 detection loss에 의해서만 supervised 됨. (student head를 update 할 때 distillation loss의 영향이 없음.)
- $i$가 증가할수록, 더 많은 student head가 detection뿐만 아니라, distillation의 영향도 동시에 받게 되며, $i=n$일 때는 기존의 prediction mimicking의 방식과 유사해짐.
- 표 1을 통해서, 모든 $i$에 대해서 성능이 향상되며, $i=3$일 때 가장 좋은 성능을 기록함을 보임.
CrossKD vs. Feature Imitation
- 그림 5를 통해, PKD에 의해 만들어진 gradient는 전체 feature map에 큰 영향을 미쳐 효율적이지 않고, 대상 물체를 정확히 target 하지 못하지만, CrossKD에 의해 만들어진 gradient는 관심 물체가 있을 만한 잠재적인 sematic area에 집중하고 있음.
CrossKD vs. Prection Mimicking
- 표 4를 통해, CrossKD와 prediction mimicking을 동시에 사용하면 CrossKD만 사용할 때보다 성능이 더 악화됨. 이는 prediction mimicking으로 인해 target conflict가 발생하기 때문임.
- Prediction mimicking이 distillation loss 계산 시 student prediction을 바로 활용하기 때문에 CrossKD보다 더 작은 $L_{1}$ distance를 기록하지만 (a), ground truth에 대한 $L_{1}$ distance는 CrossKD보다 더 높음 (b). 이는 prediction mimicking 적용 시 target conflict로 인한 contradictory optimization 과정이 수반됨을 보여줌.
Comparison with SOTA KD Methods
CrossKD on Different Detectors
Distillation under Severe Target Conflict
- 서로 다른 assigner (GFL - Retinanet)를 적용하면 target conflict가 더욱 심해지고, 이는 모델의 optimization에 부정적인 영향을 미침 (40.2 → 30.3).
- Assigner가 같더라도 (GFL - ATSS) KD 적용 시, student만 사용할 때보다 성능이 하락함 (40.2 → 39.7).
- 반면, CrossKD는 심한 target conflict가 존재하더라도 뛰어난 성능을 기록함.
Distillation between Heterogeneous Backbones
Conclusions and Discussions
- CrossKD는 student head의 중간 feature를 teacher에게 전달하여 cross-head prediction을 만들어내고, 이를 distillation에 활용하여 supervised target과 distillation target 간의 conflict를 완화함.