[논문 리뷰] Masked Distillation with Receptive Tokens

Tao Huang, Yuan Zhang외 5명|arXiv (Cornell University)|2022. 05. 29.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 밀도 높은 예측 작업을 위한 새로운 지식 증류 프레임워크인 MasKD를 제안한다. 이 프레임워크는 특징 맵 내 관심 영역 픽셀(PoIs)을 동적으로 식별하고 우선순위를 정하는 학습 가능한 수용성 토큰을 사용한다. 어텐션 기반의 적응형 증류 마스크를 생성하고, 다수의 다양한 마스크를 딱도 손실을 통해 활용함으로써, 작업 전용 사전 지식이 필요 없이도 객체 검출 및 세분화 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Distilling from the feature maps can be fairly effective for dense prediction tasks since both the feature discriminability and localization priors can be well transferred. However, not every pixel contributes equally to the performance, and a good student should learn from what really matters to the teacher. In this paper, we introduce a learnable embedding dubbed receptive token to localize those pixels of interests (PoIs) in the feature map, with a distillation mask generated via pixel-wise attention. Then the distillation will be performed on the mask via pixel-wise reconstruction. In this way, a distillation mask actually indicates a pattern of pixel dependencies within feature maps of teacher. We thus adopt multiple receptive tokens to investigate more sophisticated and informative pixel dependencies to further enhance the distillation. To obtain a group of masks, the receptive tokens are learned via the regular task loss but with teacher fixed, and we also leverage a Dice loss to enrich the diversity of learned masks. Our method dubbed MasKD is simple and practical, and needs no priors of tasks in application. Experiments show that our MasKD can achieve state-of-the-art performance consistently on object detection and semantic segmentation benchmarks. Code is available at: https://github.com/hunto/MasKD .

연구 동기 및 목표

• 모든 픽셀을 동일하게 취급하는 표준 특징 증류의 비효율성을 해결하기 위해, 모델 성능에 기여도가 다른 픽셀에 대한 차별화된 처리가 가능하도록 하기 위해.
• 밀도 높은 예측 작업에서 증류 영역을 선택하기 위해 정답 바운딩 박스나 작업 전용 사전 지식에 의존하지 않도록 하기 위해.
• 어텐션 기반 수용성 토큰을 통해 특징 맵 내 세밀한 동적 픽셀 수준의 중요도 패턴을 학습하기 위해.
• 다양한 마스크를 활용하면서도 복잡한 픽셀 간 의존성을 포착하고 학습 안정성을 유지함으로써 증류를 향상시키기 위해.
• 다양한 마스크의 중요도를 고려한 마스크 가중치 기반 메커니즘을 통해 학생 모델의 성능을 향상시키기 위해.

제안 방법

• 교수 특징 맵 내 관심 영역 픽셀(PoIs)을 식별하기 위해 픽셀 단위 어텐션을 계산하는 학습 가능한 수용성 토큰을 도입한다.
• 수용성 토큰에서 유도된 픽셀 단위 어텐션을 통해 증류 마스크를 생성하여 특징 재구성에서 우선순위를 둘 픽셀을 나타낸다.
• 특징 맵 내 다양한 복잡한 픽셀 간 의존성을 포착하기 위해 다수의 수용성 토큰을 활용한다.
• 교수 네트워크를 고정한 채로 정규 작업 손실을 사용하여 수용성 토큰을 훈련함으로써 엔드 투 엔드 최적화를 가능하게 한다.
• 다양한 마스크 간의 다양성을 유도하기 위해 딱도 손실을 적용한다. 이를 통해 중복을 방지한다.
• 다양한 마스크의 기여도에 따라 중요도를 적응적으로 할당하는 마스크 가중치 모듈을 도입한다.

실험 결과

연구 질문

• RQ1 특징 맵 내 픽셀에 대한 동적이고 어텐션 기반의 선택 방식이 밀도 높은 예측 작업에서 지식 증류를 향상시킬 수 있는가?
• RQ2 정답 애너테이션이나 작업 전용 사전 지식에 의존하지 않고도 정보가 풍부하고 다양한 증류 마스크를 학습할 수 있는가?
• RQ3 다수의 수용성 토큰과 적응형 마스크 가중치 기반 기법이 균일하거나 히وري스틱 샘플링 대비 증류 성능 향상에 어느 정도 기여하는가?
• RQ4 제안된 방법이 객체 검출 및 세분화와 같은 다양한 아키텍처와 작업에 일반화되는가?
• RQ5 학습된 모델이 교수 및 학생 특징 표현과 일치하는 의미 있는 픽셀 간 의존성을 학습할 수 있는가?

주요 결과

• COCO 객체 검출에서 MasKD는 Faster R-CNN-R50 대비 2.4 AP 향상을 기록하며, 표준 증류 기반 베이스라인을 크게 앞서 간다.
• 세분화 작업에서는 COCO에서 DeepLabV3-R18의 mIoU를 2.79 포인트 향상시켜 다양한 작업에 대한 강력한 일반화 능력을 입증한다.
• 작업 전용 사전 지식이 필요 없이도 다양한 백본 아키텍처와 데이터셋에서 일관되게 성능 향상을 이룬다.
• 절단 실험 결과, 다수의 마스크 조합, 딱도 손실을 통한 다양성 확보, 마스크 가중치 기반 기법이 증류 효율을 크게 향상시킨다.
• 시각화 결과, 학습된 마스크가 전체 바운딩 박스가 아닌 개별 객체 내 분류에 유의미한 국소화된 영역에 집중하고 있음을 확인하여 세밀한 어텐션 기반 동작을 입증한다.
• 제안된 방법은 예측되지 않은 아키텍처와 데이터셋으로도 잘 일반화되어 실세계 적용 환경에서의 강건성과 실용성을 보여준다.