[논문 리뷰] The Role of Masking for Efficient Supervised Knowledge Distillation of Vision Transformers

Seung‐Woo Son, Ryu, Jegwang외 2명|arXiv (Cornell University)|2023. 02. 21.

Advanced Neural Network Applications인용 수 5

한 줄 요약

MaskedKD는 시각 Transformer(ViT)의 지식 정복을 위한 계산 비용을 줄이기 위해 교사 추론 중 입력 이미지 패치의 25–50%를 마스킹하는 방법이다. 이는 학생 모델의 주의 메커니즘에서 유도된 시각성 점수를 기반으로 한다. 이 방법은 학생 모델의 정확도를 훼손하지 않으면서 지식 정복 FLOPs를 최대 50%까지 감소시키며, 총 학습 FLOPs를 최대 28% 감소시키고, 2.0×–3.7×의 처리량 향상을 달성한다.

ABSTRACT

Knowledge distillation is an effective method for training lightweight vision models. However, acquiring teacher supervision for training samples is often costly, especially from large-scale models like vision transformers (ViTs). In this paper, we develop a simple framework to reduce the supervision cost of ViT distillation: masking out a fraction of input tokens given to the teacher. By masking input tokens, one can skip the computations associated with the masked tokens without requiring any change to teacher parameters or architecture. We find that masking patches with the lowest student attention scores is highly effective, saving up to 50% of teacher FLOPs without any drop in student accuracy, while other masking criterion leads to suboptimal efficiency gains. Through in-depth analyses, we reveal that the student-guided masking provides a good curriculum to the student, making teacher supervision easier to follow during the early stage and challenging in the later stage.

연구 동기 및 목표

• 지식 정복 중 대규모 시각 Transformer(ViT) 교사 모델의 높은 계산 비용을 해결하기 위해.
• 학생 모델 성능을 떨어뜨리지 않으면서, 교사 모델을 실행하는데 필요한 FLOPs로 정의되는 정복 비용을 줄이기 위해.
• 교사 예측을 사전에 계산하고 저장하는 것을 피함으로써 메모리 집약적이고 데이터 증강 하에서 정확도가 저하되는 문제를 해결하기 위해.
• 자기지도 학습 및 오디오 기반 트랜스포머 설정에서도 효율적인 정복을 가능하게 하여, 감독 학습 기반 ViT 학습을 초월하기 위해.

제안 방법

• MaskedKD는 학생 모델의 마지막 멀티헤드 주의 레이어에서 계산된 패치 시각성 지표를 기반으로 마스킹할 패치를 선택하며, 기존 주의 점수를 재사용하여 거의 추가 비용 없이 구현된다.
• 이 시각성 지표는 클래스 토큰을 제외한 패치 간 주의 점수의 평균을 내어 학생 예측에 가장 관련성이 높은 이미지 특징을 식별한다.
• 가장 낮은 시각성 점수를 가진 패치가 마스킹되어 핵심 특징은 유지되면서 교사 모델에 입력되는 토큰 수가 줄어든다.
• 이 방법은 정복 중인 순방향 전파 중에 작동하며, 모델 아키텍처의 변경이나 교사 모델의 재학습이 필요하지 않다.
• 마스킹은 교사의 순방향 전파에만 적용되며, 학생은 전체 이미지를 처리하므로 특징 학습의 무결성이 유지된다.
• 이 방법은 바닐라 KD, DINO, 패치 수준 정복과 같은 다양한 ViT 정복 방법과 호환된다.

실험 결과

연구 질문

• RQ1 학생 정확도를 떨어뜨리지 않으면서 ViT 정복의 계산 비용을 줄일 수 있는가?
• RQ2 학생의 주의 메커니즘 정보만을 사용하여 교사 모델의 입력 패치를 마스킹함으로써 FLOP를 크게 줄일 수 있는가?
• RQ3 랜덤 또는 DINO 기반 시각성과 비교해 주의 기반 패치 시각성은 정복 성능 유지를 위해 얼마나 효과적인가?
• RQ4 MaskedKD는 자기지도 학습 및 비시각 트랜스포머 작업(예: 오디오 처리)으로 확장할 수 있는가?
• RQ5 성능을 유지하면서 최대 효율을 얻기 위해 마스킹할 패치의 최적 비율은 얼마인가?

주요 결과

• MaskedKD는 다양한 ViT 아키텍처와 정복 방법에서 학생의 top-1 정확도에 영향을 주지 않으면서 교사 모델의 입력 패치 수를 25–50% 감소시킨다.
• 이 방법은 정복 FLOPs(즉, 교사 추론 비용)를 최대 50% 감소시켜 DeiT-S를 DeiT-Ti로 정복할 경우 총 학습 FLOPs를 28% 감소시킨다.
• 교사 계산이 줄어들면서 처리량이 설정에 따라 2.0×에서 3.7×까지 향상된다.
• 패치 간 주의 점수를 사용한 마스킹은 성능을 저하시키며, 제안된 시각성 지표가 더 효과적임을 확인한다.
• DINO 주의 점수 기반 마스킹은 더 높은 계산 비용에도 불구하고 MaskedKD를 능가하지 못하며, 학생 주의 기반 마스킹의 우월성을 입증한다.
• 상위-k 시각성 패치(즉, 가장 중요한 패치)를 유지하는 것이 랜덤 또는 하위-k 마스킹보다 성능이 뛰어나며, 시각성과 성능 간의 상관관계를 검증한다.