[논문 리뷰] Spatiotemporal Feature Learning Based on Two-Step LSTM and Transformer for CT Scans

Chih–Chung Hsu, Chi-Han Tsai외 3명|arXiv (Cornell University)|2022. 07. 04.

COVID-19 diagnosis using AI인용 수 6

한 줄 요약

이 논문은 CT 영상에서 코로나19 분류를 위한 이중 단계의 시공간 특징 학습 프레임워크를 제안한다. 2D CNN을 사용해 슬라이스 수준의 임베딩을 추출한 후, LSTM과 Swin Transformer 모델을 활용해 시간적 모델링을 수행한다. 이 방법은 상태 기반 성능을 달성하며, LSTM 기반의 낮은 가짜 음성 비율과 Swin 기반의 낮은 가짜 양성 비율을 보이며, 진단 신뢰도에서 상호 보완적인 강점을 입증한다.

ABSTRACT

Computed tomography (CT) imaging could be very practical for diagnosing various diseases. However, the nature of the CT images is even more diverse since the resolution and number of the slices of a CT scan are determined by the machine and its settings. Conventional deep learning models are hard to tickle such diverse data since the essential requirement of the deep neural network is the consistent shape of the input data. In this paper, we propose a novel, effective, two-step-wise approach to tickle this issue for COVID-19 symptom classification thoroughly. First, the semantic feature embedding of each slice for a CT scan is extracted by conventional backbone networks. Then, we proposed a long short-term memory (LSTM) and Transformer-based sub-network to deal with temporal feature learning, leading to spatiotemporal feature representation learning. In this fashion, the proposed two-step LSTM model could prevent overfitting, as well as increase performance. Comprehensive experiments reveal that the proposed two-step method not only shows excellent performance but also could be compensated for each other. More specifically, the two-step LSTM model has a lower false-negative rate, while the 2-step Swin model has a lower false-positive rate. In summary, it is suggested that the model ensemble could be adopted for more stable and promising performance in real-world applications.

연구 동기 및 목표

• 환자 간 CT 스캔 해상도와 슬라이스 수의 일관성 부족으로 인한 딥러닝 일반화의 어려움을 해결한다.
• 다양한 길이의 CT 스캔에서 시공간 특징을 안정적으로 추출하여 신뢰할 수 있는 코로나19 분류를 위한 강건한 방법을 개발한다.
• 이중 단계 파이프라인을 통해 특징 추출과 시간 모델링을 분리하여 과적합을 줄이고 성능을 향상시킨다.
• LSTM 기반 및 트랜스포머 기반 모델의 진단 신뢰도를 비교하며, 가짜 음성 및 가짜 양성 비율에 중점을 둔다.

제안 방법

• 개별 CT 슬라이스에서 의미적 임베딩을 추출하기 위해 224차원의 임베딩을 생성하기 위해 EfficientNet-B3를 2D CNN 기반으로 사용한다.
• 이중 단계 접근법을 적용: 먼저 슬라이스 수준의 특징을 추출한 후, 양방향 LSTM(4층, 128개의 은닉 유닛)을 사용해 슬라이스 간 시간적 의존성을 모델링한다.
• 트랜스포머 기반 브랜치의 경우, 스캔 당 224개의 슬라이스를 무작위로 샘플링(224개 미만일 경우 복원 추출)하여 Swin Transformer의 입력 길이를 표준화한다.
• 정확도 향상과 분산 감소를 위해 10~50개의 무작위 슬라이스 샘플링 세트에 대한 예측 평균을 취해 앙상블 추론을 수행한다.
• 학습에 Adam 옵timizer를 사용하며, 학습률은 0.0001, 가중치 감쇠는 0.0005로 설정하고, 최종 예측에 0.5의 분류 임계값을 적용한다.

실험 결과

연구 질문

• RQ1 일관되지 않은 해상도와 슬라이스 수를 가진 다양한 길이의 CT 스캔을 효과적으로 다룰 수 있는 이중 단계 딥러닝 프레임워크는 가능한가?
• RQ2 특징 추출과 시간 모델링을 분리함으로써 의료 영상 데이터에서 과적합을 줄이고 분류 성능을 향상시킬 수 있는가?
• RQ3 코로나19 CT 분류에서 LSTM 기반 및 트랜스포머 기반 모델은 가짜 음성 및 가짜 양성 비율 측면에서 어떻게 비교되는가?
• RQ4 LSTM과 Swin 트랜스포머의 상호보완적인 아키텍처를 앙상블하면 더 안정적이고 신뢰할 수 있는 진단 결과를 도출할 수 있는가?

주요 결과

• 이중 단계 LSTM 모델은 민감도(SE) 0.900을 기록하여 다른 모델보다 낮은 가짜 음성 비율을 보였다.
• 이중 단계 Swin 트랜스포머 모델은 특이도(SP) 0.981을 기록하여 다른 모델보다 낮은 가짜 양성 비율을 보였다.
• 2D CNN 베이스라인은 매크로-F1 스코어 0.930을 기록하였으며, 무작위 슬라이스 샘플링 세트를 50번 평균 처리한 후 0.937로 향상되었다.
• 이중 단계 LSTM 모델은 2D CNN 및 ADLeaST보다 민감도에서 뛰어나 참 양성 케이스를 더 잘 탐지하는 것으로 나타났다.
• 이중 단계 Swin 모델은 최고의 특이도(0.981)를 기록하여 비코로나19 케이스를 정확히 식별하는 능력이 뛰어나다는 것을 입증했다.
• LSTM 및 Swin 브랜치의 모델 앙상블은 실제 임상 진단의 안정성을 확보하기 위해 낮은 가짜 음성 및 낮은 가짜 양성 성능을 결합하는 데 유용하다고 제안된다.