[논문 리뷰] MT-TransUNet: Mediating Multi-Task Tokens in Transformers for Skin Lesion Segmentation and Classification

Jingye Chen, Jieneng Chen외 4명|arXiv (Cornell University)|2021. 12. 03.

Cutaneous Melanoma Detection and Management인용 수 15

한 줄 요약

MT-TransUNet는 작업별 토큰을 조율하고 주의 맵 및 예측 간 일관성을 강제하여 피부 병변 세분화와 분류를 동시에 수행하는 통합된 트랜스포머 기반 모델을 제안한다. 48M 파라미터와 0.17초의 추론 시간을 바탕으로 ISIC-2017 및 PH2 데이터셋에서 최신 기술 수준의 성능을 달성하며, 이전 방법들보다 정확도와 효율성 면에서 뛰어나면서도 픽셀 수준 및 이미지 수준의 레이블을 모두 활용한다.

ABSTRACT

Recent advances in automated skin cancer diagnosis have yielded performance on par with board-certified dermatologists. However, these approaches formulated skin cancer diagnosis as a simple classification task, dismissing the potential benefit from lesion segmentation. We argue that an accurate lesion segmentation can supplement the classification task with additive lesion information, such as asymmetry, border, intensity, and physical size; in turn, a faithful lesion classification can support the segmentation task with discriminant lesion features. To this end, this paper proposes a new multi-task framework, named MT-TransUNet, which is capable of segmenting and classifying skin lesions collaboratively by mediating multi-task tokens in Transformers. Furthermore, we have introduced dual-task and attended region consistency losses to take advantage of those images without pixel-level annotation, ensuring the model's robustness when it encounters the same image with an account of augmentation. Our MT-TransUNet exceeds the previous state of the art for lesion segmentation and classification tasks in ISIC-2017 and PH2; more importantly, it preserves compelling computational efficiency regarding model parameters (48M~vs.~130M) and inference speed (0.17s~vs.~2.02s per image). Code will be available at https://github.com/JingyeChen/MT-TransUNet.

연구 동기 및 목표

• 이전 방법들이 병변 세분화를 별도의 작업으로 간주하여 세분화의 이점을 忽시하는 한계를 해결하기 위해.
• 세분화와 분류 간 공유된 특징을 활용하는 공동 학습 프레임워크를 개발하여 성능 향상시키기 위해.
• 픽셀 수준의 레이블이 없는 이미지에 대해 이중 작업 및 주의 영역 일관성 손실을 통해 모델의 강건성을 향상시키기 위해.
• 이전의 다중 작업 모델들보다 높은 계산 효율성을 유지하면서 최신 기술 수준의 성능를 달성하기 위해.
• 두 작업 모두에 공통된 인코더를 공유하여 데이터 효율적이고 모델 효율적인 학습 및 추론을 가능하게 하기 위해.

제안 방법

• 트랜스포머 인코더-디코더 아키텍처 내에서 작업별 분류 및 세분화 토큰을 도입하여 다중 작업 특징 조율을 가능하게 한다.
• 이중 작업 일관성 손실을 적용하여 클래스 주의 맵을 세분화 맵과 일치시켜 약한 지도 학습 데이터에서 일반화 성능을 향상시킨다.
• 주의 영역 일관성 손실을 적용하여 분류 및 세분화 헤드 간 일치를 강제하여, 특히 픽셀 수준의 레이블이 없을 경우에 유용하다.
• 두 작업 간 공통 인코더를 사용하여 모델 파라미터 수를 감소시키고 추론 속도를 향상시킨다.
• 반감도 학습 설정에서 세분화 감독을 위한 서로서로 대체할 수 있는 레벨셋 함수를 활용하여 진정한 마스크가 없는 이미지 처리에 대비한다.
• 세분화와 분류가 공유된 표현과 일관성 제약 조건을 통해 상호 보완적으로 지도하는 다중 작업 학습 파이프라인을 설계한다.

실험 결과

연구 질문

• RQ1 병변 세분화와 분류의 공동 학습이 단일 작업 접근 방식에 비해 두 작업의 성능을 향상시킬 수 있는가?
• RQ2 이미지 수준의 레이블만 존재할 경우, 분류 및 세분화 예측 간 일관성을 어떻게 강제하여 강건성을 향상시킬 수 있는가?
• RQ3 공통된 트랜스포머 기반 아키텍처가 이전의 다중 작업 모델들보다 훨씬 작은 모델 크기와 더 빠른 추론 속도로 최신 기술 수준의 성능를 달성할 수 있는가?
• RQ4 제안된 방법이 미세조정 없이 다양한 데이터셋으로 일반화되는 정도는 어느 정도인가?
• RQ5 이중 작업 및 주의 영역 일관성 손실이 약한 지도 학습 데이터에서 성능에 어떤 영향을 미치는가?

주요 결과

• MT-TransUNet는 ISIC-2017 및 PH2 데이터셋에서 최신 기술 수준의 성능를 달성하였으며, 세분화의 Jaccard 점수는 88.5/92.2, Dice 점수는 93.6/95.9이며, 분류 정확도는 97.1%이다.
• 모델의 추론 시간은 1장당 0.17초로, MB-DCNN의 2.02초 대비 91.6% 향상되었다.
• 단지 4800만 개의 파라미터로, MB-DCNN의 13000만 개 파라미터 대비 63% 감소한 모델 크기를 확보하였다.
• ISIC-2017에서 사전 학습한 후 PH2에서 테스트한 결과, MB-DCNN 대비 인식 정확도 3.1% 향상 및 Jaccard 점수 2.8% 향상으로 일반화 성능이 향상되었다.
• 이중 작업 및 주의 영역 일관성 손실은 데이터 증강에 대한 강건성을 크게 향상시키고, 픽셀 수준의 레이블이 없는 이미지에서의 성능을 향상시켰다.
• 세분화와 분류의 공동 학습은 상호 성능 향상을 이끌어내어, 피부 내시경 영상 분석에서 공유 표현 학습의 이점을 입증하였다.