[논문 리뷰] GenerateCT: Text-Conditional Generation of 3D Chest CT Volumes

İbrahim Ethem Hamamcı, Sezgin Er외 14명|arXiv (Cornell University)|2023. 05. 25.

Radiomics and Machine Learning in Medical Imaging인용 수 4

한 줄 요약

GenerateCT는 자유형 영상의학 보고서에서 고해상도, 고정밀도 3D 흉부 CT 볼륨을 생성하기 위해 순차적 시각 트랜스포머를 사용하여 3D CT 볼륨을 인코딩하고, 텍스트-이미지 트랜스포머를 통해 다중 모odal 간 정렬을 수행하며, 텍스트 조건부 초해상도 확산 모델을 활용하는 새로운 텍스트 조건부 3D 흉부 CT 생성 프레임워크이다. 데이터 증강에서 최신 기술 성능을 달성하여, 100,000개의 합성 볼륨을 사용해 훈련할 경우 다중 이상징후 분류 성능을 11% 향상시키고, 실재 데이터만을 사용하는 모델보다 8% 높은 성능을 기록하였다.

ABSTRACT

GenerateCT, the first approach to generating 3D medical imaging conditioned on free-form medical text prompts, incorporates a text encoder and three key components: a novel causal vision transformer for encoding 3D CT volumes, a text-image transformer for aligning CT and text tokens, and a text-conditional super-resolution diffusion model. Without directly comparable methods in 3D medical imaging, we benchmarked GenerateCT against cutting-edge methods, demonstrating its superiority across all key metrics. Importantly, we evaluated GenerateCT's clinical applications in a multi-abnormality classification task. First, we established a baseline by training a multi-abnormality classifier on our real dataset. To further assess the model's generalization to external data and performance with unseen prompts in a zero-shot scenario, we employed an external set to train the classifier, setting an additional benchmark. We conducted two experiments in which we doubled the training datasets by synthesizing an equal number of volumes for each set using GenerateCT. The first experiment demonstrated an 11% improvement in the AP score when training the classifier jointly on real and generated volumes. The second experiment showed a 7% improvement when training on both real and generated volumes based on unseen prompts. Moreover, GenerateCT enables the scaling of synthetic training datasets to arbitrary sizes. As an example, we generated 100,000 3D CTs, fivefold the number in our real set, and trained the classifier exclusively on these synthetic CTs. Impressively, this classifier surpassed the performance of the one trained on all available real data by a margin of 8%. Last, domain experts evaluated the generated volumes, confirming a high degree of alignment with the text prompt. Access our code, model weights, training data, and generated data at https://github.com/ibrahimethemhamamci/GenerateCT

연구 동기 및 목표

• 3D 의료 영상에서의 데이터 부족 및 개인정보 제약 문제를 해결하기 위해 고정밀도 3D 흉부 CT 볼륨의 텍스트 조건부 생성을 가능하게 하는 것.
• 의료 영상 분야에서 사전 학습된 3D 시각-언어 모델의 부족 문제를 해결하기 위해 새로운 계단식 프레임워크를 도입하여 3D CT 합성에 기여하는 것.
• 합성 CT 볼륨의 임상적 유용성을 데이터 증강 및 다중 이상징후 분류에 대한 제로샷 일반화 능력 평가를 통해 검증하는 것.
• 쌍방향 훈련 데이터가 필요 없이 자유형 의료 텍스트 프롬프트를 사용하여 확장 가능하고 개인정보 보호에 적합한 데이터 생성을 가능하게 하는 것.
• 코드와 사전 학습 모델을 오픈소스화하여 향후 연구의 기초를 마련하는 것.

제안 방법

• GenerateCT는 순차적 시각 트랜스포머(CT-ViT)를 활용하여 3D CT 볼륨을 토큰으로 순차적으로 인코딩함으로써 다양한 축방향 시야 및 고해상도 슬라이스 재구성 기능을 제공한다.
• 이중방향 텍스트-이미지 트랜스포머를 통해 자유형 영상의학 보고서에서 생성된 텍스트 임베딩과 인코딩된 CT 토큰 간의 정렬을 수행하며, 교차 어텐션 감독을 위해 마스크된 CT 토큰 예측을 활용한다.
• 계단식 텍스트 조건부 초해상도 확산 모델은 입력 프롬프트의 의미적 일관성을 유지하면서도 저해상도 축 슬라이스의 내부 해상도를 향상시킨다.
• 프레임워크는 CT-ViT의 재구성 손실과 초해상도의 확산 손실을 조합하여 엔드 투 엔드로 훈련되며, 텍스트 프롬프트가 조건 신호로 사용된다.
• 동일한 텍스트 프롬프트를 반복적으로 재사용하여 합성 데이터셋의 크기를 임의로 확장할 수 있도록 유연한 데이터 생성 기능을 제공한다.
• 모델 아키텍처는 3D 볼륨 전반에 걸쳐 공간적 일관성과 해부학적 현실감을 유지하면서도 복잡한 자유형 영상의학 보고서와의 정확한 일치를 보장하도록 설계되었다.

Figure 2 : The GenerateCT architecture consists of three main components. (1) The CT-ViT encoder architecture processes the embeddings of CT patches from raw slices S through a spatial transformer followed by a causal transformer (auto-regressive in-depth), generating CT tokens. (2) The vision-langu

실험 결과

연구 질문

• RQ1 자유형 영상의학 보고서에서 텍스트 조건부 생성 모델이 고정밀도이고 현실적인 3D 흉부 CT 볼륨을 생성할 수 있는가?
• RQ2 GenerateCT가 생성한 합성 CT 볼륨은 하류의 다중 이상징후 분류 모델의 성능을 어느 정도 향상시킬 수 있는가?
• RQ3 GenerateCT는 제로샷 설정에서 외부 데이터셋과 예측되지 않은 텍스트 프롬프트에 대해 얼마나 잘 일반화되는가?
• RQ4 GenerateCT가 생성한 합성 데이터는 분류 작업에서 실재 데이터만을 사용하는 모델의 성능을 뛰어넘을 수 있는가?
• RQ5 데이터 규모와 프롬프트 다양성이 생성된 CT 볼륨의 품질과 일반화 능력에 미치는 영향은 어떠한가?

주요 결과

• 20,000개의 실재 CT 볼륨과 20,000개의 GenerateCT가 생성한 합성 CT 볼륨을 혼합한 데이터셋으로 다중 이상징후 분류기를 훈련한 결과, 실재 데이터만으로 훈련한 경우에 비해 평균 정밀도(AP)가 11% 향상되었다.
• 실재 데이터셋 크기의 5배에 해당하는 100,000개의 합성 CT 볼륨으로만 훈련된 분류기는 실재 데이터만 사용하는 모델보다 평균 AP에서 8% 향상되고, 평균 AUROC에서 4% 향상되었다.
• RadChestCT 데이터셋에서 제로샷 평가를 수행한 결과, 예측되지 않은 프롬프트에서 생성된 합성 데이터로 훈련된 분류기는 평균 AP 0.146을 기록하여 실재 데이터 기준선(0.177)과 유사한 성능을 보였다.
• 제로샷 설정에서 실재 데이터와 합성 데이터를 조합한 결과, 평균 AUROC 0.623과 평균 AP 0.190을 기록하여 실재 및 합성 데이터 단독 기반 모델보다 뛰어난 성능을 보였다.
• 분야 전문가들은 생성된 CT 볼륨이 입력 텍스트 프롬프트와 높은 일치도를 보이며, 강력한 의미적 일관성과 해부학적 정확성을 확보하고 있음을 확인하였다.
• 반복적인 프롬프트 재사용을 통해 100,000개의 합성 볼륨을 생성함으로써 모델의 확장성과 대규모 개인정보 보호 기반 데이터 증강의 유용성을 입증하였다.

Figure 3 : Examples of generated chest CT volumes within the practical HU scale limits of $[-1000~{}\text{HU},+1000~{}\text{HU}]$ . Each example includes three sequential slices from the same volume, highlighting GenerateCT’s ability to maintain consistency across multiple slices.