Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging

기술 요약: 모델 붕괴 없는 엔트로피 최소화

문제 정의
테스트 시간 적응(Test-Time Adaptation, TTA)은 센서, 환경 또는 데이터 획득 프로토콜의 변화로 인해 데이터 분포가 변하는 실제 환경에서 머신러닝 모델을 배포하는 데 필수적입니다. 의료 영상 분야에서 이러한 변화는 종종 미묘하고 이질적이지만(예: 서로 다른 스캐너 또는 염색 프로토콜), 성능 저하를 심각하게 초래할 수 있습니다. 엔트로피 최소화(Entropy Minimization, EM)는 레이블이 없는 테스트 데이터에 대한 예측 불확실성을 줄이는 것을 목표로 하며, TTA의 지배적인 목적 함수가 되었습니다. 그러나 EM은 모델이 입력과 관계없이 단일 클래스를 예측하도록 수렴하는 **모델 붕괴(model collapse)**라는 치명적인 실패 모드에 직면합니다.

저자들은 이러한 붕괴의 근본 원인이 단순히 적응 과정 자체 때문이 아니라, 그들이 **예측 편향(prediction bias)**이라 명명한 현상 때문임을 밝혀냈습니다. 분포 변화가 발생하면 표현 공간(representation space) 내의 서로 다른 클래스에 해당하는 특징 클러스터들이 이동하거나 부분적으로 병합될 수 있는 반면, 분류기의 결정 경계는 고정되어 있습니다. 이는 특정 클래스가 과하게 표현되고 다른 클래스는 억제되는 예측 클래스 분포의 왜곡을 유발합니다. 표준 EM은 이러한 기존의 편향을 더욱 강화하여 병합된 클러스터를 조여줌으로써, 모델을 사소한 단일 클래스 솔루션으로 몰아넣는 양의 피드백 루프를 생성합니다.

방법론: 분포 변화 편향 감소 (DSBR)
이러한 실패 모드를 완화하기 위해, 저자들은 소스 데이터나 레이블에 접근하지 않고 오직 테스트 시간에만 작동하도록 설계된 편향 교정 목적 함수인 **분포 변화 편향 감소(Distribution Shift Bias Reduction, DSBR)**를 제안합니다.

1. 핵심 메커니즘: DSBR은 각 예측 클래스의 엔트로피 최소화 손실 기여도를 균등하게 만듦으로써 예측 편향을 해결합니다. 지배적인 클러스터가 그래디언트 업데이트를 좌우하도록 두는 대신, DSBR은 손실을 재가중하여 어떤 단일 클래스도 최적화 신호를 압도하지 못하게 합니다.
2. 수학적 정식화: 이 방법은 각 예측 클러스터에 동일한 확률 질량($1/K$)을 할당하는 균형 잡힌 타겟 분포 $p^*_X$를 정의합니다. 중요도 재가중(importance reweighting)을 사용하여, 샘플 $x$에 대한 손실은 $1 / (K \cdot \hat{p}_{\hat{Y}}(\hat{y}(x)))$로 스케일링됩니다. 여기서 $\hat{p}_{\hat{Y}}$는 예측된 클래스의 추정 빈도입니다. 이는 과소 표현된 클러스터의 샘플이 더 많은 기여를 하고, 과하게 표현된 클서터의 샘플은 기여를 덜 하도록 보장하여, 병합된 클러스터를 능동적으로 분리하도록 유도합니다.
3. 구현: 테스트 시점에는 실제 예측 클래스 분포를 알 수 없으므로, DSBR은 배치 전체에 걸친 클러스터 빈도의 지수 이동 평균(Exponential Moving Average, EMA)을 사용하여 이를 근사합니다. 이는 작은 배치 크기가 퇴보적인 배치별 추정을 초래할 수 있는 고차 클래스 데이터셋(예: $K=1000$인 ImageNet)에서 특히 중요합니다.

주요 기여

• 예측 편향의 식별: 본 논문은 특징 클러스터의 병합으로 인해 적응이 시작되기 전부터 존재하는 예측 클래스 분포의 왜곡인 '예측 편향'이 EM의 모델 붕괴의 근본 원인임을 입증했습니다. 저자들은 EM이 이 편향을 생성하는 것이 아니라, 붕괴가 일어날 때까지 이를 증폭시킨다는 점을 보여줍니다.
• DSBR 알고리즘: 저자들은 클래스 기여도를 균형 있게 맞추기 위해 엔트로피 손실을 재가중함으로써 이 실패 모드를 구체적으로 겨냥하는 DSBR을 도입합니다. 다른 접근 방식과 달리, 이 방법은 샘플 필터링, 도메인 특화 증강 또는 복잡한 정규화에 의존하지 않습니다.
• 포괄적 평가: 연구는 네 가지 의료 영상 데이터셋(WILDSCamelyon, Histopantum, MammoBench, GliomaMRI)과 자연 영상 벤치마크인 ImageNet-C를 아우르는 도전적인 적응 시나리오를 설계했습니다.

결과

• 붕마 방지: 의료 영상 벤치마크에서 표준 EM 방법들(Tent 등)은 특정 소스-타겟 도메인 쌍에 대해 빈도 수준(chance-level)의 정확도로 자주 붕괴됩니다. DSBR은 이러한 붕괴를 지속적으로 방지하며, 모든 구성에서 최적에 가까운 정확도를 달est합니다.
• 성능 향상: DSBR은 최첨단 베이스라인(Tent, SAR, ROID, DeYO 포함)보다 우수한 성능을 보였습니다. 의료 데이터셋에서 DSBR은 비적응(non-adapted) 베이스라인 대비 ResNet 기준 평균 +4.2%, ViT 기준 +2.6%의 성능 향상을 달성하여 차순위 방법들을 크게 앞질렀습니다.
• 강건성: ImageNet-C에서 DSBR은 심각한 오염, 혼합 오염 및 극단적인 레이블 변화 하에서도 탁월한 안정성을 입증했습니다. EMA 감쇠율을 적절히 조정한다면 배치 크기가 1인 경우에도 효과적으로 작동한 반면, 다른 방법들은 종종 실패하거나 붕괴되었습니다.
• 메커니즘 검증: 특징 공간 임베딩 시각화를 통해 EM은 클러스터를 병합하고 붕괴시키는 반면, DSBR은 적응 과정 내내 잘 분리된 클러스터를 유지함을 확인했습니다.

의의
본 논문은 엔트로피 최소화를 안정화하려면 왜 실패하는지에 대한 원리적이고 기계적인 이해가 필요하다고 주장합니다. 예측 편향을 모델 붕괴의 주요 동력으로 격리함으로써, 저자들은 이 단일 수치를 교정하는 것만으로도 TTA를 안정화하기에 충분하다는 것을 보여줍니다. 이 연구의 의의는 의료 영상과 같은 중요한 도메인에서의 견고한 테스트 시간 적응이 복잡한 휴리스틱이나 데이터 필터링이 아니라, 모델 붕괴를 유도하는 최적화 역학을 직접적으로 교정하는 데 있다는 것을 입증한 데 있습니다. DSBR은 하이퍼파라미터가 적고 단순한 솔루션을 제공하며, 완전히 테스트 시간에만 작동하므로 신뢰성이 가장 중요한 임상 배포에 매우 적합합니다.