기술 요약: Tracing the Oracle (TrO)

문제 정의

사전 학습된 확산 모델(Pretrained diffusion models)은 희소 뷰(sparse-view) 및 제한된 각도(limited-angle) 재구성(reconstruction)과 같이 매우 불량하게 설정된(ill-posed) 3D 컴퓨터 단층 촬영(CT) 역문제(inverse problems)를 해결하는 데 상당한 잠재력을 제공한다. 그러나 이들의 실질적인 적용은 두 가지 주요 제한 사항으로 인해 어려움을 겪고 있다:

1. 계산 오버헤드(Computational Overhead): 추론 과정은 높은 충실도(fidelity)를 달야하기 위해 일반적으로 수천 번의 연속적인 신경 함수 평가(NFE)를 필요로 하며, 이는 상당한 계산적 부담을 지운다.
2. 비효격적인 타임스텝 스케줄링(Inefficient Timestep Scheduling): 기존 방법들은 대개 균등(Uniform-t), 코사인(Cosine), EDM과 같은 균등하거나 휴리스틱한 타임스ップ 스케줄에 의존한다. 이러한 스케줄은 역방향 조건부 확산 확률 미분 방정식(reverse conditional diffusion SDE)의 비균등한 진화 과정을 포착하지 못한다. 결과적으로, 엄격한 계산 예산(예: 10단계 미만) 하에서 작동할 때 상당한 절단 오차(truncation errors)를 유발한다. 또한, 무조건부 생성(unconditional generation)을 위해 설계된 표준 최적화 전략(GITS와 같은)은 조건부 역문제에서는 효과적이지 않은데, 이는 측정 제약 조건이 무조건부 생성 매니폴드를 방해하고, 조건부 솔버의 확률적 특성이 결정론적 궤적 가정을 무효화하기 때문이다.

방법론

저자들은 3D CT 재구성을 위한 타임스텝 스케줄링을 최적화하기 위해 설계된 플러그 앤 플레이(plug-and-play) 프레임워크인 **Tracing the Oracle (TrO)**를 제안한다. 핵심 철학은 밀도 있게 샘플링된 수치 적분 궤적을 고충실도 "오라클(oracle)"로 취급하고, 이 오라클에 대한 누적 오차를 최소화하는 희소하고 최적화된 타임스텝 부분 집합을 추출하는 것이다.

이 방법론은 세 가지 핵심 구성 요소로 이루어진다:

1. 오라클 궤적 구축 (Oracle Trajectory Construction)

참조 궤적($T_{oracle}$)은 데이터의 일부에 대해 많은 횟수의 평가($N=200$)를 사용하는 분해 확산 샘플링(Decomposed Diffusion Sampling, DDS) 프레임워크를 사용하여 생성된다. 이 조밀한 경로는 특정 측정 조건 하에서의 역방향 확산 과정에 대한 그라운드 트루스(ground truth) 역할을 한다.

2. 노이즈 재사용 메커니즘 (Noise Reuse Mechanism)

확률적 설정에서 희소 근사치와 오라클 사이의 오차를 정량화할 때 직면하는 결정적인 과제는 독립적인 노이즈 샘플링에 의해 도입되는 분산이다. 만약 근사치와 오라클이 노이즈를 독립적으로 샘플링한다면, 둘 사이의 기하학적 거리는 변동하여 진정한 결정론적 절단 오차를 가릴 것이다.
이를 해결하기 위해, TrO는 노이즈 재사용(noise reuse) 기술을 채택한다. 소스 타임스텝 $t_i$에서 타겟 타임스텝 $t_j$로의 전이 비용을 평가할 때, 이 방법은 오라클이 $t_{j-1}$에서 $t_j$로 전이하는 동안 원래 주입되었던 정확한 확률적 노이즈 벡터($\epsilon_{j-1}$)를 재사용한다. 이를 통해 결정론적 절단 오차를 확률적 분산으로부터 분리하여, 정확한 오차 정량화를 가능하게 한다.

3. 이산 최단 경로 최적화 (Discrete Shortest-Path Optimization)

스케줄링 작업은 이산 최단 경로 문제로 정식화된다. 목표는 누적 전이 비용을 최소화하는 $L$개의 타임스텝($L \ll N$) 부분 집합을 선택하는 것이다.

• 비용 함수 (Cost Function): 타임스텝 $t_i$와 $t_j$ 사이의 비용 $C_{i,j}$는 기하학적 거리(절단 오차)와 정규화 항의 합으로 정의된다:
  $$C_{i,j} = \|x_{t_j|t_i} - x_{t_j}\|_2 + \kappa \left( (t_i - t_j) - \frac{T}{L} \right)^2$$
  정규화 항은 네트워크가 훈련된 조건부 매니폴드 밖에서 작동하도록 강제하는 것을 방지하기 위해 균등한 시간 간격으로부터의 큰 편차에 페널티를 부여한다.
• 솔버 (Solver): 동적 계획법(Dynamic Programming, DP)이 이 이산 경로 문제의 전역 최적해를 찾는 데 사용된다. 탐욕적(greedy) 접근 방식과 달리, DP는 추출된 스케줄이 누적 오차의 전역 최솟값임을 보장한다.

캘리브레이션 세트에서 오프라인으로 계산되면, 최적화된 스케줄은 3D 볼륨 재구성을 위한 추론 단계에서 추가적인 계산 오버헤드 없이 적용된다.

주요 기여

1. 플러그 앤 플레이 프레임워크: 저자들은 3D 역문제에서 누적 궤적 오차를 최소화하기 위해 타임스텝 스케줄링을 이산 최단 경로 문제로 정식화하는 TrO를 도입한다.
2. 노이즈 재사용 메커니즘: SDE 기반 샘플링 궤적 내에서 결정론적 절단 오차를 확률적 분산으로부터 격리하는 새로운 기술을 도입하여, 표준 결정론적 지표가 실패하는 조건부 설정에서도 정확한 오차 정량화를 가능하게 한다.
3. 성능 향상: 광범에한 평가를 통해 TrO가 사전 정의된 휴리스틱 스케줄에 비해, 특히 엄격한 계산 예산(예: 8~15 NFE) 하에서 재구성 충실도(PSNR 및 SSIM)를 일관되게 개선함을 입증하였다.

실험 결과

프레임워크는 최첨단 DDS 솔버와 결가된 AAPM 데이터셋의 희소 뷰 CT(SV-CT) 및 제한된 각도 CT(LA-CT) 작업에 대해 평가되었다.

• 정량적 성능: TrO는 다양한 NFE 예산(8, 10, 15 단계)에 대해 다양한 베이스라인(Uniform-t, Uniform-$\lambda$, Quadratic, EDM, Cosine)보다 우수한 성능을 보였다.
  • 8-view SV-CT 작업의 8 NFE 환경에서, TrO는 Axial에서 38.53, Coronal에서 39.78의 PSNR을 달성하여 다음으로 우수한 방법(Quadratic)을 상당한 차이로 앞질렀다.
  • TrO와 베이스라인 간의 성능 격차는 NFE 예산이 줄어들수록 커졌으며, 이는 적은 단계의 레짐(regime)에서의 효율성을 강조한다.
  • TrO는 또한 경쟁력 있는 결과를 얻기 위해 일반적으로 수천 번의 NFE를 필요로 하는 다른 확산 기반 솔버들(DiffMBIR, Score-Med, MCG)보다 우수한 성능을 보여주었다.
• 정성적 분석: 시각적 비교를 통해, 특히 재구성 오류가 발생하기 쉬운 영역에서 향상된 아티팩트 억제 및 세부 사항 보존 능력을 보여주었다.
• 적응형 동작: 최적화된 스케줄 분석 결과, TrO가 특정 측정 연산자에 적응함을 확인했다. 저주파 정보가 신뢰할 수 없는 LA-CT의 경우, TrO는 누락된 구조를 합성하기 위해 역과정의 초기 단계에 더 많은 단계를 할당한 반면, SV-CT의 경우 고주파 정밀도를 위해 후기 단계에 더 집중하였다.

의의 및 주장

본 논문은 TrO가 3D 의료 영상 적용에 있어 확산 모델의 핵심 병목 현상인 계산 비용과 재구성 품질 사이의 트레이드오프를 해결한다고 주장한다. 고충실도 오라클에 기반한 타임스텝 스케줄링을 최적화 문제로 다루고 노이즈 재사용을 통해 확률적 오차 정량화 문제를 해결함으로써, TrO는 다음을 가능하게 한다:

• 고충실도 적은 단계 재구성: 계산 비용의 극히 일부(예: <15 단계)만 사용하여 고NFE 솔버와 대등한 재구성 품질을 달성한다.
• 작업별 적응: 역문제의 저하 특성(예: 희소 뷰 vs 제한된 각도)에 따라 중요한 진화 단계에 계산 자원을 동적으로 할당하는 능력.
• 제로 오버헤드: 최적화는 오프라인으로 수행되므로, 추론 단계에서 표준 샘플링 과정 이상의 추가적인 계산 부담을 주지 않는다.

저자들은 TrO가 적은 단계의 샘플링과 계산 비용이 많이 드는 고NFE 솔버 사이의 성능 격차를 좁힘으로써, 3D 의료 영상 분야에서 확산 모델의 효율적이고 견고한 배포를 촉진한다고 결론짓는다.