[논문 리뷰] Knowledge Distillation in YOLOX-ViT for Side-Scan Sonar Object Detection

Martin Aubard, László Antal외 2명|arXiv (Cornell University)|2024. 01. 01.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 측면 스캔 소나 이미지를 위한 시각 트랜스포머 기반 향상된 객체 검출기인 YOLOX-ViT을 제안하고, 지식 증류를 통해 이를 더 작은 효율적인 모델로 압축한다. 방법은 벽 검출에서 오진 양성률을 최대 20.35% 감소시키면서 mAP50를 0.13–0.16 향상시키며, 이는 ViT 통합과 KD가 함께 작용할 경우 수중 객체 검출의 정확도와 강인성을 향상시킴을 보여준다.

ABSTRACT

In this paper we present YOLOX-ViT, a novel object detection model, and investigate the efficacy of knowledge distillation for model size reduction without sacrificing performance. Focused on underwater robotics, our research addresses key questions about the viability of smaller models and the impact of the visual transformer layer in YOLOX. Furthermore, we introduce a new side-scan sonar image dataset, and use it to evaluate our object detector's performance. Results show that knowledge distillation effectively reduces false positives in wall detection. Additionally, the introduced visual transformer layer significantly improves object detection accuracy in the underwater environment. The source code of the knowledge distillation in the YOLOX-ViT is at https://github.com/remaro-network/KD-YOLOX-ViT.

연구 동기 및 목표

• 시각 트랜스포머 기반 YOLOX 아키텍처를 사용하여 측면 스캔 소나 이미지에서의 객체 검출 성능을 향상시키기.
• 모델 크기를 줄이되 검출 정확도를 유지하기 위해 지식 증류의 효과성을 조사하기.
• 수중 로봇을 위한 새로운 측면 스캔 소나 데이터셋을 제안하여 벽 검출에 집중하기.
• 비디오 환경에서 저조도 조건에서 시각 트랜스포머 레이어가 특징 추출 및 검출 강인성에 미치는 영향 평가하기.
• 큰 교사 네트워크에서의 증류를 통해 소형 모델의 오진 양성률을 감소시키기.

제안 방법

• 기본 구조의 SPPBottleneck 이후에 시각 트랜스포머(ViT) 레이어를 삽입하여 YOLOX를 향상시켰으며, 기존의 표준 컨볼루션 블록을 대체하였다.
• 표준 크로스 엔트로피 및 IoU 기반 손실 함수를 사용하여 더 큰 YOLOX-L 및 YOLOX-L-ViT 모델을 '교사'로 훈련시켰다.
• 지식 증류를 적용하여 복합 손실 함수를 사용: ℒ = λ·ℒhard + (1−λ)·ℒsoft, 여기서 ℒhard는 정답 레이블의 크로스 엔트로피이고, ℒsoft는 교사로부터의 소프트 레이블 증류이다.
• YOLOX-Nano 및 YOLOX-Nano-ViT 모델을 더 작은 '학생' 모델로 압축하기 위해 증류를 적용하였으며, 주로 특징 맵과 로짓 증류에 중점을 두었다.
• 교사 모델의 훈련 단계에서만 데이터 증강을 적용하였고, 학생 모델은 온라인 증강 유무에 따라 훈련시켰다.
• 기본 구조의 중간 특징 맵에 초점을 맞춘 힌트 기반 증류 전략을 사용하여 학생 및 교사 표현 간의 일치를 도모하였다.

실험 결과

연구 질문

• RQ1 지식 증류는 측면 스캔 소나 객체 검출을 위한 소형 YOLOX 모델에서 오진 양성률을 효과적으로 감소시킬 수 있는가?
• RQ2 YOLOX에 시각 트랜스포머 레이어를 통합하면 수중 환경에서 검출 정확도와 강인성이 향상되는가?
• RQ3 ViT와 지식 증류의 조합은 소형 모델에서 mAP50와 추론 효율성에 어떤 영향을 미치는가?
• RQ4 이러한 낮은 데이터 환경에서 온라인 데이터 증강은 큰 모델과 작은 모델의 성능에 어떤 영향을 미치는가?
• RQ5 ViT 레이어는 특히 벽과 같은 구조적 특징 검출에 있어 소나 이미지의 특징 추출을 향상시킬 수 있는가?

주요 결과

• 지식 증류를 통해 YOLOX-Nano-noAug 학생 모델의 오진 양성률이 L-ViT 교사 모델을 사용할 경우 6% 감소하였고, ViT 기반 교사 모델을 사용할 경우 20.35% 감소하였다.
• YOLOX-Nano-ViT-noAug 학생 모델은 오직 13.9개의 오진 양성과 함께 38%의 검출률을 기록하여 기본 Nano 모델보다 정밀도와 오진 양성률 통제에서 뛰어난 성능을 보였다.
• YOLOX-L-ViT는 온라인 데이터 증강 없이 훈련된 결과, 0.41 mAP50와 98.93%의 검출 지속 시간을 기록하였으며, YOLOX-L(0.33 mAP50, 87.05% 검출)보다 뚜렷이 뛰어난 성능을 보였다.
• ViT 레이어를 통한 향상 덕분에 L 모델의 검출 정확도가 8% 향상되었고, Nano 모델의 경우 23% 향상되었으며, 특히 장시간 비디오 추론에서 두드러졌다.
• 온라인 데이터 증강 없이 훈련된 모델은 증강을 적용한 모델보다 더 우수한 일반화 능력과 낮은 과적합 현상을 보였으며, 특히 L 및 L-ViT 변형에서 두드러졌다.
• ViT 기반 학생 모델(Nano-ViT-noAug)은 0.13 mAP50와 30.16%의 정밀도를 기록하여, ViT가 소형 모델에서도 특징 학습을 향상시킬 수 있음을 입증하였다.