Learning a Sampling-Free Variational DNN Plugin from Tiny Training Sets to Refine OOD Segmentation With Uncertainty Estimation

기술 요약: 소규모 학습 세트로부터 샘플링이 필요 없는 변분 DNN 플러그인 학습을 통한 불확실성 추정 기반 OOD 세그멘테이션 정교화

문제 정의
심층 신경망(DNN)은 분포 외(out-of-distribution, OOD) 의료 영상에 적용될 때 성능이 크게 저하되는 문제를 자주 겪는다. 이 문제는 서로 다른 임상 현장 간의 영상 프로토콜, 스캐너 하드웨어 및 재구성 방식의 필연적인 변화로 인해 발생한다. 대상 도메인 데이터로 모델을 재학습시키는 것이 잠재적인 해결책이 될 수 있지만, 새로운 의료 데이터셋을 획득하고 주석(annotation)을 다는 데 따르는 높은 비용 때문에 이는 종종 비현실적이다. 또한, 기존의 세그멘테이션 정교화 방법들은 학습 중에 OOD 데이터에 접근해야 하거나 광범적인 사전 학습에 의존하는 경우가 많아, 초기 모델 개발 단계에서 OOD 데이터를 사용할 수 없는 시나리오에는 적합하지 않다. 과제는 OOD 학습 데이터에 접근할 수 없는 상태에서 OOD 이미지에 대한 저하된 세그멘테이션을 교정하는 동시에, 신뢰할 수 있는 불확실성 추정을 제공하고, 특히 소규모의 분포 내(in-distribution, ID) 학습 세트만 사용 가능한 상황에서도 계산 효율성을 유지하는 것이다.

방법론: VarDeepPCA
저자들은 기존 DNN이 생성한 OOD 데이터 상의 저하된 세그멘테이션 맵을 복원하거나 정교화하기 위해 설계된, 가벼운 변분 DNN 프레임워크인 VarDeepPCA를 제안한다. 핵심 혁신은 영상 획득 과정에서의 도메인 변화에도 불구하고 대체로 불변하는 관심 해부학적 구조의 고유한 기하학적 사전 정보(geometric priors)를 활용하는 데 있다.

1. 기하학적 사전 정보 학습: 원본 이미지로부터 학습하는 대신, VarDeepPCA는 소수의 ID 세그멘테이션 맵(100~200쌍)으로부터 직접 주요 변이 모드(principal modes of variation)를 학습한다. 이는 유효한 해부학적 기하 구조가 저차원 매니폴드(manifold) 상에 존재한다는 가설에 기반한다.
2. 심플렉스 기반 잠재 표현: 프레임워크는 인코더-디코더 구조를 채택한다. 인코더는 입력 세그멘테이션 맵을 $K$-차원 특징 벡터로 매핑하며, 이 벡터는 소프트맥스(softmax) 함수를 통해 $(K-1)$-차원 심플렉스(simplex) 상에 존재하는 확률 벡터로 변환된다. 이 벡터는 $K$개의 주요 변이 모드의 혼합을 나타낸다.
3. 샘플링이 필요 없는 변분 학습: 본 연구의 핵심적인 이론적 기여는 베이지안 프레임워크 내에서 소프트맥스 매핑을 재해석한 것이다. 저자들은 소프트맥스 출력이 디리클레 사전 확률(Dirichlet prior)로 매개변수화된 잠재 범주형 분포(latent categorical distribution)에 대해 명시적인 베이지안 주변화(Bayesian marginalization)를 암시적으로 수행함을 입증하였다. 이를 통해 샘플링이 필요 없는(sampling-free) 학습 및 추론이 가능하다. 값비싼 몬테카를로 샘플링을 요구하는 전통적인 변분 오토인코더(VAE)와 달리, VarDeepPCA는 출력 분포의 평균과 분산을 해석적으로 폐쇄형(closed form)으로 도출한다.
   • 디코더는 평균 세그멘테이션 맵($M$)과 픽셀별 분산 맵($V$)을 출력한다.
   • 분산은 디코더의 자코비안(Jacobian)과 잠재 범주형 분포의 분산을 사용하여 계산된다.
   • 불확실성 맵은 분산 맵의 제곱근으로 도출된다.
4. 2단계 정교화 알고리즘: 베이스라인 DNN이 생성한 OOD 세그멘테이션 $\tilde{Z}$를 교정하기 위해 다음 과정을 거친다:
   • 필터링(Filtering): 저하된 맵을 인코더-디코더에 통과시켜 비주요 변이 성분을 제거하여 필터링된 맵 $Z$를 생성한다.
   • 투영(Projection): 필터링된 맵을 학습된 유효 기하 구조의 매니폴드 위로 명시적으로 투영한다. 이는 필터링된 맵의 가능도(likelihood)를 학습된 가우시안 분포 하에서 최대화하도록 경사 상승법(gradient ascent)을 통해 잠재 특징 벡터를 최적화함으로써 달성되며, 이후 디코딩을 통해 복원된 세그멘테이션 $M^*$과 그에 따른 불확실성 $U^*$를 얻는다.

주요 기여

• 새로운 프레임워크: 소규모 ID 데이터셋만을 사용하여 OOD 세그멘테이션을 정교화하는 플러그인으로서, OOD 데이터나 광범한 사전 학습의 필요성을 제거하는 VarDeepPCA를 도입하였다.
• 샘플링이 필요 없는 변분 추론: 소프트맥스 함수에 대한 이론적 재해석을 통해 정확한 분포 모델링과 계산 효율적인 샘플링 프리 학습 및 추론을 가능하게 하며, 폐쇄형 불확실성 추정을 제공한다.
• 경량화된 아키텍처: 매우 작은 규모의 학습 세트(100~200개 이미지)로부터 학습할 수 있도록 설계되었으며, 파라미터 수가 적어(1.02M ~ 2.72M) 데이터가 제한적인 임상 환경에 적합하다.
• 포괄적인 검증: 4가지 서로 다른 임상 응용 분야(심근, 망막 신경절 경계, 전립선, 태아 머리)에 대해 14개의 공개 데이터셋을 사용하여 검증하였으며, 트랜스포머, 확산 모델(diffusion models), 파운데이션 모델을 포함한 15가지 기존 방법들과 비교하였다.

결과
본 논문은 VarDeepPCA가 기존 방법들을 일관되게 능가함을 보여주는 광범위한 정량적 및 정성적 결과를 보고한다:

• 세그멘테이션 정확도: 4가지 응용 분야와 15가지 베이스라인 전체에 걸쳐, 기존 DNN에 VarDeepPCA를 결로하여 증강함으로써 OOD 데이터에서의 경계 오차(HD95 및 ASD로 측정)를 유의미하게 감소시켰다. 예를 들어, 심근 세그멘테이션의 경우 OOD 데이터셋에서의 평균 HD95가 평균 13.213.6 픽셀에서 7.28.0 픽셀로 감소하여 임상적으로 수용 가능한 수준의 성능을 확보하였다.
• 해부학적 타당성: 이 방법은 OOD 실패 사례에서 흔히 나타나는 해부학적으로 타당하지 않은 기하 구조(예: 끊어진 덩어리 형태)를 학습된 매니폴드로 투영함으로써 효과적으로 복원하였다.
• 불확실성 교정: VarDeepPCA는 ProbUNet 및 PHISeg과 같은 확률적 베이스라인 모델의 불확실성 추정 교정 성능(NCC, US, TACE로 측정)을 크게 향상시켰다. 생성된 불확실성 맵은 베이스라인에 비해 더 국소적이며 해부학적 경계와 잘 일치하였다.
• 효율성: 이 플러그인은 학습 및 추론 시간에 최소한의 오버헤드만을 추가하며, 베이스라인 방법에서 사용하는 것 이상의 추가 학습 데이터를 요구하지 않는다.
• 강건성: 민감도 분석 결과, 본 방법은 잠재 차원($K$) 및 학습 세트 크기($S$)의 변화에 강건하며, 단 100개의 학습 샘 샘플만으로도 우수한 성능을 보였다.

의의 및 주장
저자들은 VarDeepPCA가 DNN의 임상 배치를 가로막는 핵심 장벽인, 비용이 많이 드는 재학습 없이 새로운 영상 사이트로 일반화하지 못하는 문제를 해결한다고 주장한다. 소규모 ID 데이터셋으로부터 유효한 해부학적 기하 구조의 분포를 명시적으로 모델링함으로써, 이 프레임워크는 OOD 세그멘테이션을 정교화하기 위한 원칙적이고 계산 효율적인 솔루션을 제공한다.

본 논문은 기존 방식들이 학습을 위해 OOD 데이터가 필요하거나 높은 계산 비용(예: 확산 모델, 대규모 트랜스포머)을 초래하는 것과 달리, VarDeepPCA가 세그멘테이션 정확도와 불확실성 추정 모두를 개선하는 가볍고 샘플링이 필요 없는 대안을 제공한다는 점을 강조한다. 저자들은 이 방법이 기하학적 일관성에 의존하며, 매우 무정형적인 병변(예: 종양)이나 초기 오류가 너무 심각하여 올바른 기하 구조가 학습된 매니폴드 밖에 있는 경우 어려움을 겪을 수 있다는 제약 사항을 겸허히 인정하였다. 그러나 이러한 경우에도 프레임워크는 출력을 유효한 해부학적 기하 구조로 성공적으로 투영한다. 이 연구는 VarDeepPCA를 기존의 세그멘테이션 모델을 강화할 수 있는 다재다능한 플러그인으로 위치시키며, 이는 비 DNN 방식까지 확장될 수 있다.