[논문 리뷰] The Swiss Army Knife for Image-to-Image Translation: Multi-Task Diffusion Models

Julia Wolleb, Robin Sandkühler외 2명|arXiv (Cornell University)|2022. 04. 06.

Radiomics and Machine Learning in Medical Imaging인용 수 17

한 줄 요약

이 논문은 회귀 또는 분할 모델의 외부 기울기를 사용하여 기본 디퓨전 모델을 재학습하지 않고도 노이즈 제거 디퓨전 샘플링을 안내하는 다중 작업 확산 모델을 제안한다. 이 방법은 노화된 얼굴 이미지나 뇌 종양 성장 시뮬레이션과 같은 제어 가능한 이미지 간 변환을 가능하게 하며, 목표 외 특징을 유지하면서도, 단일 사전 훈련된 DDPM을 사용하여 얼굴 노화 및 MR 영상 종양 생성 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Recently, diffusion models were applied to a wide range of image analysis tasks. We build on a method for image-to-image translation using denoising diffusion implicit models and include a regression problem and a segmentation problem for guiding the image generation to the desired output. The main advantage of our approach is that the guidance during the denoising process is done by an external gradient. Consequently, the diffusion model does not need to be retrained for the different tasks on the same dataset. We apply our method to simulate the aging process on facial photos using a regression task, as well as on a brain magnetic resonance (MR) imaging dataset for the simulation of brain tumor growth. Furthermore, we use a segmentation model to inpaint tumors at the desired location in healthy slices of brain MR images. We achieve convincing results for all problems.

연구 동기 및 목표

• 사용자 지정 작업을 위한 이미지 간 변환을 가능하게 하기 위해 단일 사전 훈련된 확산 모델을 사용한다.
• 배경 및 구조적 세부 정보를 유지하면서 뿐만 아니라, 작업 관련 이미지 특징만 수정하는 과제를 해결한다.
• 분류 작업에서의 분류자 안내 기법을 분류 작업 외의 회귀 및 분할 작업으로 확장한다.
• 건강한 MR 스캔에 뇌 종양을 포함한 현실적인 합성 데이터를 생성한다.
• 추론 시기 기울기 주입을 통해 동일한 확산 모델을 다양한 작업에 재사용할 수 있음을 보여준다.

제안 방법

• 사전 훈련된 노이즈 제거 확산 확률 모델(DDPM)을 태스크 지정 없이 학습된 소스 데이터셋에서 사용하여 기본 생성자로 활용한다.
• 동일한 데이터셋에서 독립적으로 훈련된 외부 태스크 지정 모델(회귀 또는 분할)을 사용하여 안내 신호를 제공한다.
• DDIM 샘플링 중에 각 단계에서 노이즈가 있는 입력에 대한 태스크 지정 모델의 출력 기울기를 기반으로 노이즈 제거 과정을 안내한다.
• 가이드 강도는 스케일링 인자 $ s_t $ 로 제어되며, 이는 양수 또는 음수 값으로 설정되어 목표 속성(예: 연령 또는 종양 크기)을 증가 또는 감소시킬 수 있다.
• 회귀의 경우, 목표 값 $ i $ 를 사용하여 $ s_t = i - R(x_t, t) $ 를 계산한다. 여기서 $ R $ 은 단계 $ t $ 에서의 회귀 모델 예측이다.
• 분할의 경우, 레이블 마스크가 목표 종양 위치를 정의하며, 분할 모델의 손실 기울기가 샘플링 중에 종양 유사 특징을 이미지에 주입하는 데 사용된다.

실험 결과

연구 질문

• RQ1 재학습 없이도 단일 사전 훈련된 확산 모델이 여러 이미지 간 변환 태스크에 사용될 수 있는가?
• RQ2 외부 회귀 및 분할 모델의 기울기 안내가 샘플링 중 연령 또는 종양 크기와 같은 이미지 속성을 효과적으로 제어할 수 있는가?
• RQ3 목표 외 특징(예: 배경, 옷, 머리카락)이 수정된 이미지와 일치하는가?
• RQ4 이 방법은 데이터 증강을 위해 건강한 MR 슬라이스에 뇌 종양을 포함한 현실적인 합성 의료 영상을 생성할 수 있는가?
• RQ5 정밀도 및 속성 제어 측면에서 이 방법의 성능이 태스크 지정 확산 모델과 비교해 어떻게 되는가?

주요 결과

• 이 방법은 얼굴 이미지에 대한 연령 회귀를 성공적으로 수행하여, 40세 입력에서 원하는 연령(예: 10세, 20세, 60세, 80세)의 현실적인 이미지를 생성한다.
• 뇌 MR 영상의 경우, 회귀 손실 기울기를 조정하여 종양 성장 또는 위축을 정확하게 시뮬레이션하며, 목표 값 $ i $ 를 통해 종양 크기를 제어할 수 있다.
• 분할 기반 안내 방법은 사용자가 정의한 위치에 건강한 뇌 MR 슬라이스에 종양을 성공적으로 인painting 하였으며, 영역별 변화를 강조하는 차이 맵을 통해 확인되었다.
• 동일한 기본 DDPM이 모든 태스크—회귀, 종양 성장 시뮬레이션, 종양 인painting—에서 재사용되었으며, 이는 프레임워크의 유연성을 입증한다.
• 결과적으로 목표 속성과 관련된 특징만 수정되고, 배경, 얼굴 특징, 옷 등 다른 요소는 입력 이미지와 일치함을 보여주었다.
• 이 방법은 시각적 결과가 매우 매력적이며, 최소한의 잡음과 높은 구조적 정밀도를 유지하면서도 얼굴 노화 및 뇌 종양 시뮬레이션 태스크에서 뛰어난 성능을 보였다.