Med-SegLens: Latent-Level Model Diffing for Interpretable Medical Image Segmentation

Motivation

• Known: CNN 및 transformer 기반 segmentation 모델은 강력한 벤치마크 성능을 보이지만 내부 표현이 불투명하며, Grad-CAM, SHAP 등 post-hoc saliency 기반 해석법은 예측과 상관된 영역만 보여줄 뿐 내부 표현이 실패를 어떻게 유발하는지에 대한 인과적 통찰은 제공하지 못한다.
• Gap: 기존 mechanistic interpretability 연구는 주로 language model이나 ViT/CLIP 같은 classification 백본에 집중되어 있고, dense prediction 과제인 medical image segmentation에 대한 mechanistic 분석 및 model diffing을 통한 dataset shift 연구는 거의 탐구되지 않았다.
• Why: 임상 현장에서 이질적인 인구집단에 배포되는 segmentation 모델의 실패를 재훈련 없이 원인 규명하고 교정할 수 있다면, 고위험 의료 AI의 신뢰성과 적응성을 크게 높일 수 있다.
• Approach: SegFormer와 U-Net을 healthy(IXI), adult/pediatric/SSA glioma(BraTS) 데이터셋별로 독립적으로 학습한 뒤, 각 모델의 중간 activation에 BatchTopK sparse autoencoder를 학습시켜 latent를 추출하고, Hungarian algorithm 기반 cosine similarity 매칭으로 shared 및 dataset-specific latent를 식별한다.

Achievement

Figure 5. Mean Dice versus steering strength (α) for adult-specific (left), universally shared (middle), and random (rig

1. Med-SegLens 프레임워크 제안: segmentation 모델을 위한 최초의 mechanistic model-diffing 프레임워크로, 데이터셋 간 표현을 해석 가능한 형태로 원칙적으로 비교 가능케 함.
2. 자동 latent 해석 기법: geometry 및 spatial 정보에 기반한 자동 해석 파이프라인을 제안하여 latent에 해부학적·병리학적 의미를 부여함.
3. shared/population-specific latent 규명: cross-dataset model diffing을 통해 안정적인 공유 backbone 표현과 dataset shift를 유발하는 population-specific latent를 구분함.
4. 인과적 개입을 통한 실패 교정: 실패 사례의 70%에서 latent-level intervention으로 성능을 복구하고, 가장 영향받은 클래스의 Dice score를 39.4%에서 74.2%로 향상시킴.
5. 재훈련 없는 cross-dataset adaptation 개선: class-specific bias를 완화하여 재훈련 없이도 도메인 적응 성능을 향상시킴.

How

Figure 4. Feature swapping between Adult and Pediatric SAEs.

• IXI(healthy), BraTS 2023(adult/pediatric/SSA glioma) T1-weighted MRI 데이터셋을 사용하여 동일 아키텍처(SegFormer-B4, U-Net)를 각 데이터셋별로 독립 학습하여 4개의 dataset-specific 모델 생성.
• 각 모델의 중간 activation에서 BatchTopK SAE를 학습하여 encoder/decoder 가중치를 얻고, sparse latent representation z, h(z), 재구성 activation z̃ 도출.
• 데이터셋 쌍마다 encoder/decoder 가중치에 대한 cosine similarity matrix를 계산하고 Hungarian algorithm으로 매칭을 수행하여, πenc(i)=πdec(i) 및 유사도 임계값 τ를 만족하는 latent를 shared latent로, 그렇지 않은 latent를 dataset-specific latent로 분류.
• 식별된 dataset-specific latent에 대해 ablation 및 steering(강도 α 조절) 방식의 targeted intervention을 수행하여 segmentation 실패에 대한 인과성을 검증하고 성능 개선 효과를 측정.
• feature swapping 실험(Adult↔Pediatric SAE)을 통해 latent의 population-specific 역할을 추가 검증.

Originality

• segmentation과 같은 dense prediction task에 mechanistic interpretability 및 model diffing 개념을 최초로 적용.
• SAE 기반 latent 분해와 Hungarian algorithm 기반 cross-architecture/cross-dataset alignment를 결합하여 shared/population-specific latent를 정량적으로 구분하는 절차 제안.
• geometry 및 spatial grounding 기반 automated interpretation 파이프라인으로 latent에 해부학적·병리학적 semantic을 자동 부여.
• latent-level ablation/steering을 통한 재훈련 없는 causal intervention으로 segmentation failure를 직접 교정한 실증적 증명.

Limitation & Further Study

• 실험이 뇌 MRI(brain glioma, subcortical) 도메인 및 두 아키텍처(SegFormer, U-Net)에 국한되어 있어, 다른 장기·모달리티·아키텍처로의 일반화 가능성이 검증되지 않음.
• 70%의 실패 사례에서만 복구가 이루어졌다는 점에서 나머지 실패 유형에 대한 latent 개입의 한계 및 실패 원인 분석이 부족함.
• SAE의 sparsity 수준(k), 매칭 임계값(τ) 등 하이퍼파라미터에 대한 민감도 분석이 제한적으로 제시된 것으로 보이며, 이들 선택이 결과에 미치는 영향에 대한 심층 검토가 필요함.
• 후속 연구로는 다양한 영상 모달리티 및 장기에 대한 확장, 완전 자동화된 latent 기반 실시간 모니터링 시스템 개발, intervention의 임상적 검증 등이 필요함.

Evaluation

총평: medical image segmentation에 mechanistic interpretability와 model diffing을 접목한 신선하고 실용적인 시도로, latent-level intervention을 통한 재훈련 없는 실패 교정이 인상적이나 다른 도메인·아키텍처로의 일반화 검증이 추가로 필요하다.