[논문 리뷰] A-Eval: A Benchmark for Cross-Dataset Evaluation of Abdominal Multi-Organ Segmentation

Ziyan Huang, Zhongying Deng외 12명|arXiv (Cornell University)|2023. 09. 07.

Colorectal Cancer Screening and Detection인용 수 6

한 줄 요약

이 논문은 FLARE22, AMOS, WORD, TotalSegmentator의 네 개의 대규모 데이터셋에서의 학습 데이터를 사용하여 복부 다장기 분할 모델의 교차 데이터셋 평가를 위한 A-Eval을 소개한다. 평가 대상은 이들의 검증 세트와 BTCV의 학습 세트이다. 주요 기여는 공동 학습, 의사 레이블링, 다중 모odal 학습이 일반화 능력을 크게 향상시킨다는 것을 입증한 것이며, 모델 크기 증가가 특정 지점 이후에는 수익 감소 현상에 들어간다는 점을 밝혀냈다.

ABSTRACT

Although deep learning have revolutionized abdominal multi-organ segmentation, models often struggle with generalization due to training on small, specific datasets. With the recent emergence of large-scale datasets, some important questions arise: extbf{Can models trained on these datasets generalize well on different ones? If yes/no, how to further improve their generalizability?} To address these questions, we introduce A-Eval, a benchmark for the cross-dataset Evaluation ('Eval') of Abdominal ('A') multi-organ segmentation. We employ training sets from four large-scale public datasets: FLARE22, AMOS, WORD, and TotalSegmentator, each providing extensive labels for abdominal multi-organ segmentation. For evaluation, we incorporate the validation sets from these datasets along with the training set from the BTCV dataset, forming a robust benchmark comprising five distinct datasets. We evaluate the generalizability of various models using the A-Eval benchmark, with a focus on diverse data usage scenarios: training on individual datasets independently, utilizing unlabeled data via pseudo-labeling, mixing different modalities, and joint training across all available datasets. Additionally, we explore the impact of model sizes on cross-dataset generalizability. Through these analyses, we underline the importance of effective data usage in enhancing models' generalization capabilities, offering valuable insights for assembling large-scale datasets and improving training strategies. The code and pre-trained models are available at \href{https://github.com/uni-medical/A-Eval}{https://github.com/uni-medical/A-Eval}.

연구 동기 및 목표

• 다른 복부 다장기 분할 데이터셋 간의 모델 일반화 능력 평가에 대한 표준화된 평가 부족 문제를 해결하기 위해.
• 공동 학습, 의사 레이블링, 다중 모달 학습과 같은 데이터 사용 전략이 새로운 데이터셋에서의 모델 성능에 미치는 영향을 조사하기 위해.
• 모델 크기가 복부 장기 분할에서 교차 데이터셋 일반화에 미치는 영향을 평가하기 위해.
• 단일 데이터셋 성능을 넘어서 공정하고 종합적인 평가가 가능한 통합 벤치마크를 제공하기 위해.

제안 방법

• 공식 학습 세트를 활용하여 FLARE22, AMOS, WORD, TotalSegmentator의 네 개의 대규모 공개 데이터셋에서 구성된 벤치마크를 구축한다.
• 이 네 데이터셋의 검증 세트와 BTCV의 학습 세트를 활용하여 다섯 개의 서로 다른 데이터셋으로 구성된 다양성 있는 평가 세트를 구성한다.
• 다양한 전략을 사용하여 모델을 학습하고 평가한다: 개별 데이터셋 학습, 미학습 데이터의 의사 레이블링, 모ality 혼합(CT/MR), 모든 데이터셋에 대한 공동 학습.
• 여덟 가지 공통 장기 카테고리에서 DSC(Dice Similarity Coefficient)와 NSD(Normalized Surface Dice)를 사용하여 모델 성능을 평가한다.
• 다양한 크기의 모델(STU-Net-S, -B, -L, -H)를 체계적으로 비교하여 모델 용량이 일반화에 미치는 영향을 분석한다.
• 모델 유형당 20회 교차 데이터셋 평가를 수행하고 평균 DSC와 NSD, 표준편차를 계산하여 안정성과 일반화 능력을 평가한다.

Figure 1 : Comparison of the original evaluation approach versus our proposed A-Eval benchmark for assessing model generalizability. (a) Original evaluation involves training and testing on the same dataset, providing good results but leaving uncertainty when applied to other datasets. (b) A-Eval, o

실험 결과

연구 질문

• RQ1 대규모 복부 분할 데이터셋에서 학습된 모델들이 다른 이미징 프로토콜과 애너테이션을 가진 다른 데이터셋으로 일반화될 수 있는가?
• RQ2 공동 학습, 의사 레이블링, 다중 모달 학습과 같은 다양한 데이터 사용 전략이 교차 데이터셋 일반화에 어떤 영향을 미치는가?
• RQ3 다양한 복부 데이터셋에서 모델 크기가 일반화 성능에 어떤 영향을 미치는가?
• RQ4 모델 용량을 늘리면 항상 일반화 성능이 향상되는가, 아니면 포화점이 존재하는가?

주요 결과

• FLARE22, AMOS, WORD, TotalSegmentator의 네 개의 대규모 데이터셋에서 공동 학습을 수행한 결과, BTCV 데이터셋에서 평균 DSC 93.81%와 NSD 93.81%를 기록하여 개별 데이터셋 학습보다 뛰어난 성능을 보였다.
• 의사 레이블링을 통해 평균 DSC가 FLARE22 전용 학습 시 89.28%에서 90.94%로 향상되어, 미학습 데이터 활용의 가치를 입증했다.
• AMOS의 CT 및 MR 데이터를 활용한 다중 모달 학습으로 평균 DSC가 91.65%에서 92.53%로 상승하여, 모달리티 다양성이 일반화 능력을 향상시킨다는 것을 보여주었다.
• 더 큰 모델(STU-Net-L)은 더 높은 평균 DSC(86.64%)와 NSD(90.14%)를 기록했지만, 가장 큰 모델(STU-Net-H)은 추가 향상 없이 포화 상태에 도달함을 보여주어 포화점이 존재함을 시사했다.
• STU-Net-H 모델은 평균 DSC 86.16%와 NSD 89.95%를 기록하여 STU-Net-L 모델보다 略적으로 낮은 성능을 보였으며, 이는 극단적 규모에서의 과적합 또는 비효율성을 시사한다.
• 벤치마크 분석 결과, 데이터 중심 전략—특히 공동 학습과 데이터 증강—이 모델 크기 증가만으로는 달성할 수 없는 일반화 능력 향상에 더 큰 기여를 한다는 점을 확인했다.

Figure 2 : Visualization of the performance of STU-Net-L models trained individually on different datasets (FLARE22 [ 24 ] , AMOS CT [ 16 ] , AMOS MR [ 16 ] , WORD [ 21 ] , TotalSegmentator [ 42 ] ) and validated on multiple datasets (FLARE22 [ 24 ] , AMOS CT [ 16 ] , AMOS MR [ 16 ] , WORD [ 21 ] ,