다중 의료 영상 복원을 위한 작업 적응형 트랜스포머

초록

본 논문은 PET 합성, CT 저노이즈, MRI 초해상도 등 서로 다른 의료 영상 복원 과제를 하나의 모델로 동시에 수행할 수 있는 All‑in‑One 프레임워크인 TAT(Task‑Adaptive Transformer)를 제안한다. 작업 간 그래디언트 충돌을 방지하기 위해 작업별 가중치를 동적으로 생성하고, 작업 난이도에 따라 손실 가중치를 샘플 수준에서 자동 조정하는 두 가지 핵심 메커니즘을 도입한다. 실험 결과, TAT는 개별 작업에서도, 통합 All‑in‑One 설정에서도 최첨단 성능을 달성한다.

상세 분석

TAT는 기존 All‑in‑One 모델이 공유 파라미터에 고정된 가중치를 사용해 작업 간 간섭(task interference)과 불균형(task imbalance) 문제에 취약하다는 점을 정확히 짚어낸다. 이를 해결하기 위해 두 단계의 적응형 설계를 도입한다. 첫 번째는 Task‑Adaptive Weight Generation이다. 입력 이미지의 저해상도 특징을 Transformer Encoder를 통해 잠재 표현 I_LF 로 변환한 뒤, Gradient‑Detached 복사본을 Task Representation Extraction Network(TREN)으로 전달한다. TREN은 간단한 3×3 Conv 블록으로 구성되어, I_LF 로부터 작업 고유의 임베딩 Z∈ℝ^d 를 추출한다. 이 임베딩은 이후 MLP를 통해 depth‑wise convolution의 가중치 W_G 로 변환되며, 기존 공유 가중치 W_S 와 λ라는 학습 가능한 스칼라를 통해 선형 결합한다( W = W_S + λ·W_G ). depth‑wise convolution을 선택한 이유는 파라미터 수가 O(C) 로 선형 확장되어 메모리와 연산 효율성을 크게 개선하면서도 지역 정보를 보존할 수 있기 때문이다. 이렇게 생성된 작업‑특화 가중치는 Decoder의 각 Weight‑Adaptive Transformer Block(WATB)에 주입되어, 동일 파라미터가 작업마다 다른 업데이트 방향을 갖게 함으로써 그래디언트 충돌을 원천 차단한다. 두 번째는 Task‑Adaptive Loss Balancing이다. 기존 멀티태스크 학습에서 흔히 쓰이는 σ_t 기반 가중치 조정은 작업 수준에서만 동작한다. TAT는 샘플‑레벨 손실 동향을 반영하기 위해 L1 거리 세 가지(원본‑HQ, 원본‑예측, 예측‑HQ)를 concatenate하고, 이를 또다시 Gradient‑Detached 상태에서 MLP에 입력해 σ 를 예측한다. 최종 손실은 ½·σ²·L₁(예측, HQ)+logσ 형태로, σ가 클수록 해당 샘플의 손실이 억제되고, logσ 항이 과도한 스케일링을 방지한다. 이 메커니즘은 각 샘플의 학습 난이도에 따라 자동으로 가중치를 재조정하므로, 데이터 불균형이나 특정 작업의 과다 학습을 효과적으로 완화한다. 전체 아키텍처는 U‑shaped Encoder‑Decoder 구조에 3‑stage Transformer Encoder와 4‑stage Decoder를 배치하고, Encoder‑Decoder 사이에 TREN과 WATB를 삽입한다. 최종 출력은 Residual 형태로 복원된 고품질 이미지와 원본 저품질 이미지를 합산해 얻는다. 실험에서는 PET synthesis, CT denoising, MRI super‑resolution 세 가지 공개 데이터셋을 사용해 개별 및 All‑in‑One 설정 모두에서 PSNR, SSIM, RMSE 지표에서 기존 최고 성능을 상회한다. 특히 All‑in‑One 환경에서 작업 간 간섭을 최소화한 덕분에 개별 모델 대비 평균 PSNR가 0.5~1.2 dB 향상되었다. 코드와 모델 가중치가 공개돼 재현 가능성도 확보하였다.