Weight Space Correlation Analysis: Quantifying Feature Utilization in Deep Learning Models

기술 요약: 특징 활용량 정량화를 위한 가중치 공간 상관관계(Weight Space Correlation) 분석

문제 정의
의료 영상 분야의 딥러닝 모델은 모델이 실제 임상 신호가 아닌 혼란 변수 메타데이터(예: 스캐너 모델, 획득 프로토콜)에 의존하는 "지름길 학습(shortcut learning)"에 취약합니다. 선행 연구들은 이러한 혼란 정보가 이미지 임베딩 내에 흔히 *인코딩(encoded)*되어 있음을 입증했으나(Gichoya et al., 2022; Glocker et al., 2023), 중요한 격차가 존재합니다. 즉, 임베딩 공간에 특정 특징이 존재한다는 사실만으로는 모델이 최종 예측을 위해 해당 특징을 실제로 *활용(utilize)*하는지를 증명할 수 없습니다. 임베딩 분포를 분석하거나 대조적(counterfactual) 테스트 세트를 생성하는 기존 방식들은 분류 헤드(classification head)가 이러한 혼란 변수를 실제로 사용하는지 여부를 직접적으로 다루지 못합니다. 본 논문은 "모델의 결정 경계가 임상적으로 무관한 요인들과 능동적으로 정렬되는가?"라는 근본적인 질문을 다룹니다.

방법론: 가중치 공간 상관관계(WSC) 분석
저자들은 주 임상 작업(primary clinical task)의 분류 헤드와 보조 메타데이터 작업(auxiliary metadata tasks)의 헤드 간의 정렬도를 측정함으로써 특징 활용량을 정량화하는 해석 가능한 방법론인 WSC 분석을 제안합니다. 이 방법은 다음 세 단계로 진행됩니다:

1. 선형 표현(Linear Representation): 이 방법은 선형 분류 헤드의 가중치($W$)를 잠재 임베딩 공간에 작용하는 "어텐션 벡터(attention vectors)"로 취급합니다. $C$개의 클래스를 가진 작업의 경우, 가중치 행렬 $W \in \mathbb{R}^{C \times d}$는 각 클래스에 대한 결정 방향을 정의합니다.
2. 데이터 매니폴드로의 투영(Projection onto Data Manifold): 의미 있는 차원에서 비교가 이루어지도록 하기 위해, 분류 가중치를 주성분 분석(PCA)을 사용하여 고유한 데이터 매니폴드에 투영합니다. 저자들은 상위 $k$개의 주성분을 추출하여(최소 50개를 확보하고 분산의 99%를 포착) 투영 행렬 $P$를 형성합니다. 가중치는 $W' = W P^\top$로 변환됩니다. 이 단계는 직교하며 정보가 없는 차원을 필터링합니다.
3. 정렬도 정량화(Quantifying Alignment): 이 방법은 주 작업($W'_A$)과 보조 메타데이터 작업($W'_m$)의 투영된 가중치 벡터 사이의 코사인 유사도를 계산합니다. 임상 작업의 결정 경계와 메타데이터 작업(예: 스캐너 유형)의 결정 경계 사이에 높은 상관관계가 나타나는 것은 지름길 학습의 정량적 대리 지표가 되며, 이는 모델이 두 예측을 위해 동일한 잠재 특징에 의존하고 있음을 나타냅니다.

주요 기여

• 새로운 지표: 메타데이터의 인코딩과 추론 시의 능동적 활용을 구분하는 WSC 분석을 도입했습니다.
• 유도된 편향을 통한 검증: 저자들은 태아 평면(fetal plane) 클래스와 스캐너 모델 사이에 강한 상관관계를 인위적으로 도입한 "유도된 편향 데이터셋(Induced Bias Dataset)"을 생성하여 방법론을 검증했습니다. WSC 분석은 결과적으로 나타난 가중치 벡터의 높은 정렬을 성공적으로 탐지한 반면, 상관관계가 없는 "균형 잡힌 데이터셋(Balanced Dataset)"에서는 낮은 정렬을 보여, 이 방법이 지름길 학습에 대한 민감성을 가짐을 확인했습니다.
• 임상 모델에 적용: 이 방법론을 산전 조기 출산(sPTB) 예측을 위한 최첨단 모델인 SA-SonoNet에 적용했습니다.

실험 결과
연구에는 두 개의 비공개 임상 초음파 데이터셋(태아 해부학적 분류 데이터셋 및 sPTB 예측을 위한 자궁경부 데이터셋)이 사용되었습니다.

• 인코딩 가능성 vs. 활용성: 태아 데이터셋에서 베이스라인 모델들은 메타데이터(스캐너, 병원 ID)가 이미지 임베딩 내에 매우 높게 인코딩되어 있음(프로빙 시 높은 AUROC)을 확인했습니다. 그러나 WSC 분석 결과, 주 평면 분류 작업의 가중치 벡터는 스캐너 분류 작업과 대체로 직교(낮은 상관관계)하는 것으로 나타났습니다. 이는 메타데이터를 인코딩하도록 명시적으로 훈련된 멀티태스크 모델에서도 동일하게 나타났으며, 이는 주 작업이 스캐너 지름길에 의존하지 않음을 시사합니다.
• 유도된 지름길 탐지: 특정 평면이 특정 스캐너와 상관관계를 갖도록 데이터셋에 편향을 주었을 때, 평면과 스캐너 가중치 사이의 WSC 값이 유의미하게 증가했습니다. 이는 WSC가 모델이 지름길 전략을 채택하도록 강제될 때 이를 탐지할 수 있음을 확인시켜 주었습니다.
• SA-SonoNet 분석: sPTB 예측 모델에 적용한 결과는 다음과 같습니다:
  • 임상적 관련성: sPTB 분류 가중치는 자궁경부 길이 및 픽셀 간격과 같은 임상적으로 관련 있는 요인들과 중간 정도의 상관관계를 보였습니다.
  • 혼란 변수로부터의 분리: 가중치는 임상적으로 무관한 획득 요인, 특히 스캐너 유형으로부터 분리되어 있었습니다.
  • 주의 사항: 출생 체중이 sPTB와 높은 WSC를 보였으나, 저자들은 프로빙 헤드가 출생 체중을 효과적으로 예측하지 못했기 때문에 이는 기능적으로 무의미하다고 언급했습니다. 이는 해당 정보가 견고하게 인코딩되지 않았음을 나타냅니다.

의의 및 주장
본 논문은 WSC 분석이 의료 영상의 모델 신뢰성을 검증하기 위한 필수적인 도구를 제공한다고 주장합니다. 분류기가 실제 임상 신호와 무관한 특징을 활용하는지 조사함으로써, 연구자들은 단순한 성능 지표를 넘어 내부의 결정 로직을 이해할 수 있습니다. 저자들은 메타데이터가 의료 영상 임베딩에 흔히 존재하지만, SA-SonoNet과 같은 모델은 결정 경계가 임상적으로 유의미한 신호와 정렬되고 스캐너 모델과 같은 가짜 상관관계로부터 분리되도록 설계되거나 훈련될 수 있다고 결론지었습니다. 이는 모델이 데이터셋의 아티팩트를 이용하는 것이 아니라 의도된 병리학을 학습하고 있음을 확인할 수 있는 방법을 제시합니다.

한계점
저자들은 이 방법이 임베딩 공간과 클래스 예측 사이의 선형 관계를 가정하고 있다는 점을 인정합니다. 이는 표준 딥러닝 관례와 일치하지만 복잡한 비선형 상호작용은 포착하지 못할 수 있습니다. 또한, 현재 방법은 개별 메타데이터 속성을 분석하는데, 실제 세계의 지름길은 여러 혼란 변수의 비선형적 조합을 포함할 수 있으며, 향후 연구에서 이를 조사하는 것을 목표로 합니다. 마지막으로, 낮은 WSC가 이미지 품질 변화나 특징 변화와 같은 다른 요인에 의한 도메인 변화에 대한 강건성을 보장하지는 않지만, 그러한 실패의 원인이 지름길 학습 때문은 아님을 배제해 준다는 점을 명시했습니다.