TCAM-Diff: 트라이플레인 표현으로 3D 의료영상 확산 모델의 부담을 줄이다

도입

3D 의료영상 생성은 본질적으로 메모리와 계산 비용이 큰 문제다. 영상이 단순한 2D 이미지가 아니라 체적 데이터이기 때문에, 해상도가 높아질수록 처리해야 할 복셀 수가 빠르게 늘어난다. arXiv 논문 “TCAM-Diff: Triplane-Aware Cross-Attention Medical Diffusion Model”은 이 병목을 줄이기 위해 조밀한 3D 볼륨을 직접 모델링하는 대신, 트라이플레인 표현을 학습하고 그 위에서 확산 모델을 적용하는 TCAM-Diff를 제안한다.

핵심 내용

• 트라이플레인 표현으로 3D 정보 압축: TCAM-Diff는 dense volume에서 triplane representation을 학습한다. 이는 3D 구조 정보를 여러 개의 2D 특징 평면으로 나누어 담는 방식으로, 전체 체적 데이터를 그대로 다루는 것보다 메모리 부담을 줄일 수 있다.

• decoder-only autoencoder 사용: 논문은 이 모델이 decoder-only autoencoder 방식으로 트라이플레인 표현을 학습한다고 설명한다. 또한 과적합을 막기 위해 generalization operations를 활용한다.

• 트라이플레인 인식 cross-attention 확산 모델: 생성 단계에서는 triplane-aware cross-attention diffusion model이 사용된다. 이는 트라이플레인 특징을 단순한 잠재 벡터처럼 취급하지 않고, 각 평면 사이의 정보를 학습하고 통합하도록 설계된 점이 특징이다.

• 사전 학습 디코더로 3D 볼륨 복원: 확산 모델이 만든 특징은 사전 학습된 디코더 모듈을 통해 3D 볼륨으로 변환된다. 결과적으로 무거운 3D 생성 문제를 더 간결한 특징 공간의 생성 문제로 바꾸는 구조다.

• 다양한 해상도 데이터셋 실험: 실험은 BrainTumour 128×128×128, Pancreas 256×256×256, Colon 512×512×512 데이터셋에서 수행됐다. 재구성 품질은 MSE와 SSIM으로, 생성 품질은 W-GAN critic으로 평가했다.

의미와 영향

TCAM-Diff의 의의는 고해상도 3D 의료영상 생성에서 표현 방식 자체를 다시 설계했다는 점에 있다. 전체 볼륨을 직접 생성하는 접근보다, 디코딩 가능한 트라이플레인 특징을 생성하는 접근은 제한된 GPU 메모리 환경에서도 더 큰 3D 데이터를 다룰 가능성을 열어준다. 이는 합성 의료영상 생성, 데이터 증강, 3D 의료영상 표현 학습 연구에 활용될 수 있다.

다만 논문의 결과는 제시된 데이터셋과 지표 안에서 해석해야 한다. 저자들은 유사한 크기의 latent space를 가진 기존 encoder-decoder 방법보다 더 나은 재구성 및 생성 결과를 보고했지만, 임상적 활용을 위해서는 병변의 사실성, 해부학적 일관성, 개인정보 보호, 편향, 다운스트림 작업 성능 등 추가 검증이 필요하다. 그럼에도 TCAM-Diff는 3D 의료영상 확산 모델을 더 효율적으로 만들기 위한 중요한 방향을 보여준다.

출처: arXiv