RAFM: Retrieval-Augmented Flow Matching for Unpaired CBCT-to-CT Translation

이 논문은 소규모 의료 데이터셋에서 비쌍합 CBCT-to-CT 변환의 불안정성을 해결하기 위해 DINOv3 인코더와 글로벌 CT 메모리 뱅크를 활용한 검색 기반 가짜 쌍 (retrieval-guided pseudo pairs) 구성 전략을 도입한 '검색 증강 흐름 매칭 (RAFM)'을 제안하고, SynthRAD2023 벤치마크에서 기존 방법들을 능가하는 성능을 입증했습니다.

핵심

1. 왜 이 기술이 필요한가요? (문제 상황)

암 치료 (방사선 치료) 에서는 두 가지 사진이 필요합니다.

• CT 사진: 마치 고화질 지도처럼 뼈와 장기의 위치가 정확하고, 방사선량을 계산하는 데 필수적입니다.
• CBCT 사진: 환자가 치료대에 누워 치료받는 동안, 매번 찍는 사진입니다. 하지만 이 사진은 화질이 흐리고 (아티팩트), 뼈나 장기의 정확한 위치 (수치) 를 알기 어렵습니다.

문제는 무엇일까요?
치료 계획은 정밀한 CT 로 세우는데, 실제 치료는 흐린 CBCT 로 진행됩니다. 그래서 "흐린 CBCT 를 고화질 CT 로 바꿔주는 인공지능"이 필요합니다.

하지만 여기서 큰 걸림돌이 있습니다.

• 완벽한 짝 (Pair) 이 없습니다. 같은 환자를 같은 시간에 CT 와 CBCT 로 동시에 찍는 경우가 거의 없습니다. 시간이 지남에 따라 환자의 몸이 조금씩 변하고, 찍는 각도도 달라지기 때문입니다.
• 기존 기술들은 이 '완벽하지 않은 짝'을 강제로 맞추려다 보니, 뼈가 뒤틀리거나 장기가 사라지는 엉뚱한 결과가 나오기도 했습니다.

---

2. RAFM 이 해결책입니다 (해결 방법)

이 연구팀은 **RAFM(검색을 활용한 흐름 매칭)**이라는 새로운 기술을 개발했습니다. 이를 '유령 지도 찾기' 게임에 비유해 볼까요?

🧩 비유: 유령 지도 찾기 게임

• 상황: 당신은 흐릿한 지도 (CBCT) 를 가지고 있습니다. 하지만 정확한 지도 (CT) 가 어디에 있는지 모릅니다.
• 기존 방법 (랜덤 짝짓기): 주변에 있는 지도를 아무거나 하나 집어서 "이게 내 지도의 정확한 버전이야!"라고 가정하고 수정합니다. 하지만 그 지도가 완전히 다른 동네라면? 지도는 엉망이 됩니다.
• RAFM 의 방법 (검색 활용):
  1. 지식 저장고 (기억 은행): 수많은 정밀 CT 지도들을 한곳에 모아둡니다.
  2. 스마트 검색 (DINOv3): 흐릿한 지도 (CBCT) 를 보면, AI 가 "이 지도의 특징은 '산이 있고 강이 흐르는' 형태구나"라고 분석합니다.
  3. 가장 비슷한 짝 찾기: 저장고에서 "산과 강이 있는" 정밀 CT 지도를 찾아냅니다.
  4. 자연스러운 변환: 이제 흐릿한 지도를 그 정밀한 지도로 부드럽게 이어주는 길을 그립니다. (이걸 'Flow Matching'이라고 합니다.)

핵심 아이디어:
무작위로 짝을 지어주는 게 아니라, 내용 (해부학적 구조) 이 가장 비슷한 CT 사진을 찾아서 짝을 맞춰줍니다. 그래서 환자의 뼈나 장기 위치가 흐트러지지 않고 자연스럽게 고화질로 바뀝니다.

---

3. 이 기술이 얼마나 뛰어난가요? (결과)

연구팀은 실제 데이터 (SynthRAD2023) 로 실험을 해보았습니다.

• 화질: 흐릿한 CBCT 가 CT 처럼 선명해졌습니다. (오차가 가장 적음)
• 구조: 환자의 뼈와 장기 모양이 원래대로 유지되었습니다. (다른 방법들은 장기가 뭉개지거나 사라지는 경우가 많았음)
• 속도: 복잡한 과정을 거치는 기존 기술들보다 훨씬 빠르고 안정적입니다.

한 줄 요약:

"RAFM 은 흐릿한 치료 중 사진을, 내용 (해부학적 구조) 이 가장 비슷한 정밀 CT 지도를 찾아내어, 자연스럽게 고화질로 변환해주는 똑똑한 기술입니다."

이 기술 덕분에 환자는 더 정확한 방사선 치료를 받을 수 있게 되었고, 의사들은 치료 계획을 더 신뢰할 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem Definition)

• 배경: 방사선 치료 계획 수립에는 Hounsfield Unit (HU) 값이 정확한 전자 밀도 정보를 제공하는 CT 가 표준으로 사용됩니다. 반면, 치료 중 적응형 방사선 치료 (Adaptive Radiotherapy) 에 유용한 CBCT 는 촬영이 용이하지만 심한 아티팩트와 부정확한 HU 값으로 인해 직접적인 선량 계산에 사용하기 어렵습니다.
• 과제: CBCT 를 CT 와 유사한 이미지로 변환 (Synthetic CT, sCT 생성) 하는 작업이 필요합니다.
• 핵심 난제:
  • 데이터 불일치: CBCT 와 CT 는 시간적 간격, 해부학적 변화, 등록 (registration) 오차 등으로 인해 쌍을 이루지 않은 (Unpaired) 데이터가 대부분입니다.
  • 기존 방법의 한계: GAN 기반 방법은 훈련 불안정성과 구조 왜곡 문제가 있고, 확산 (Diffusion) 모델은 복잡한 훈련 파이프라인이 필요합니다.
  • Flow Matching 의 한계: Rectified Flow (RF) 는 이론적으로 쌍을 이루지 않은 학습이 가능하지만, 의료 영상은 데이터 양이 적고 배치 크기가 작아 무작위 또는 배치 내 (batch-local) 페어링 시 의미론적으로 일치하지 않는 엔드포인트 쌍이 생성되어 훈련이 불안정해지고 해부학적 구조가 손상될 수 있습니다.

2. 제안 방법: RAFM (Methodology)

저자들은 **Retrieval-Augmented Flow Matching (RAFM)**을 제안하여 쌍을 이루지 않은 CBCT-to-CT 변환 문제를 해결합니다.

2.1 핵심 아이디어

• 검색 기반 페이징 (Retrieval-guided Coupling): 무작위 또는 배치 내 페어링 대신, 전역적인 CT 메모리 뱅크 (Memory Bank) 에서 CBCT 슬라이스와 특징 (Feature) 이 가장 유사한 CT 슬라이스를 검색하여 가상의 쌍 (Pseudo Pair) 을 구성합니다.
• 비적대적 (Non-adversarial) 접근: GAN 과 달리 적대적 훈련 없이 결정론적 운송 (Deterministic Transport) 을 통해 분포 간 변환을 수행합니다.

2.2 세부 기술

1. Rectified Flow (RF) 기반:
   • 소스 분포 (CBCT) 에서 타겟 분포 (CT) 로의 직선 경로를 따라 결정론적 미분방정식 (ODE) 을 학습합니다.
   • 목표는 $x_0$ (CBCT) 와 $x_1$ (CT) 사이의 속도장 $v_\theta(x, t)$를 학습하여 $t=0$에서 $t=1$로 이동시키는 것입니다.
2. 검색 강화 페이징 전략:
   • DINOv3 인코더 사용: 고정된 (Frozen) DINOv3 인코더를 사용하여 CT 슬라이스들의 특징 벡터를 추출합니다.
   • CT 메모리 뱅크 (Memory Bank): 전체 CT 데이터셋을 직접 검색하는 대신, 미니배치 단위로 업데이트되는 '롤링 (Rolling)' 형태의 CT 특징 메모리 뱅크를 유지합니다.
   • 페어링 과정: 현재 CBCT 미니배치의 각 슬라이스에 대해 메모리 뱅크에서 코사인 유사도가 가장 높은 CT 슬라이스를 검색하여 $(x_0, x_1)$ 쌍을 생성합니다. 이는 환자 ID 나 정렬된 쌍 정보를 사용하지 않으므로 엄격한 Unpaired 설정을 유지합니다.
3. 훈련 및 추론:
   • 생성된 가상의 쌍을 사용하여 Flow Matching 손실 함수를 최적화합니다.
   • 추론 시에는 학습된 ODE 를 10 단계의 오일러 적분 (Euler integration) 으로 풀어서 CBCT 를 CT 로 변환합니다.

3. 주요 기여 (Key Contributions)

1. 의료 영상용 RF 최적화: 소규모 데이터셋과 작은 배치 크기 환경에서 Flow Matching 의 실용적 효과를 입증하고, 무작위 페어링의 불안정성을 해결하기 위해 검색 기반 페이징을 도입했습니다.
2. RAFM 프레임워크 제안: DINOv3 특징 공간과 전역 메모리 뱅크를 활용한 새로운 Unpaired 변환 아키텍처를 제안하여, 해부학적 구조 보존과 이미지 품질을 동시에 향상시켰습니다.
3. 엄격한 평가 프로토콜: SynthRAD2023 데이터셋에서 Subject-level True-Unpaired (훈련 시 CBCT 와 CT 가 완전히 다른 환자) 설정을 적용하여 기존 방법들과 공정한 비교를 수행했습니다.

4. 실험 결과 (Results)

SynthRAD2023 (골반 영역) 데이터셋에서 GAN 기반 (CycleGAN, GcGAN, CUT) 및 Diffusion/SB 기반 (SynDiff, UNSB) 방법들과 비교했습니다.

• 정량적 성능 (Quantitative Metrics):
  • FID (Fréchet Inception Distance): 53.29 (기존 최강인 UNSB 의 62.91 대비 개선) → 분포의 현실성이 가장 뛰어남.
  • MAE (Mean Absolute Error): 101.2 HU (SynDiff 의 104.2 대비 개선) → HU 값 정확도 향상.
  • SSIM/PSNR: 각각 80.96%, 25.15 dB 로 모든 방법 중 최고 성능 기록.
  • SegScore (해부학적 일관성): 75.77% 로, TotalSegmentator 를 이용한 장기 분할 정확도가 가장 높음.
• 정성적 성능 (Qualitative):
  • 아티팩트 제거가 더 깔끔하고 해부학적 구조가 안정적으로 유지됨을 시각적으로 확인했습니다.
• 효율성:
  • GAN 의 순환 일관성 (Cycle-consistency) 이나 Diffusion 의 다단계 샘플링에 비해 훈련 및 추론 속도가 효율적입니다.

5. 의의 및 결론 (Significance & Conclusion)

• 해부학적 보존의 중요성 입증: 쌍을 이루지 않은 학습에서도 의미론적으로 유사한 페어링 (Semantic Coupling) 을 구성하는 것이 해부학적 구조를 보존하고 훈련을 안정화하는 데 결정적임을 증명했습니다.
• 실용적 가치: 쌍을 이루지 않은 의료 데이터를 활용하여 고품질의 Synthetic CT 를 생성할 수 있는 새로운 패러다임을 제시하며, 적응형 방사선 치료의 실용성을 높일 수 있습니다.
• 코드 공개: 제안된 방법론의 코드는 GitHub 에서 공개되어 재현성과 후속 연구를 장려합니다.

요약하자면, RAFM은 소규모 의료 데이터 환경에서 Flow Matching 의 이론적 장점을 실용적으로 구현하기 위해 검색 기반의 페이징 전략을 도입함으로써, 쌍을 이루지 않은 CBCT-to-CT 변환에서 기존 최첨단 방법들을 능가하는 성능을 달성한 획기적인 연구입니다.