MedSteer: Counterfactual Endoscopic Synthesis via Training-Free Activation Steering

이 논문은 엔도스코프 이미지 생성 시 해부학적 구조를 유지하면서 병변 개념만 정밀하게 변경할 수 있는 훈련 없는 활성화 조종 프레임워크인 'MedSteer'를 제안하고, 이를 통해 기존 방법보다 우수한 반사실적 데이터 생성 및 하류 진단 성능 향상을 입증했습니다.

핵심

🍳 문제 상황: "재료만 바꾸고 요리 모양은 그대로!"

의료 AI 를 훈련시키려면 '병이 있는 사진'과 '건강한 사진'을 많이 보여줘야 합니다. 하지만 현실에서는 병이 있는 장기의 사진을 구하기 어렵고, 건강한 사람과 병든 사람의 장기는 모양이 완전히 달라서 비교하기 힘듭니다.

기존의 AI 기술은 **"병이 있는 장기를 그려줘"**라고 명령하면 병이 있는 사진을, **"건강한 장기를 그려줘"**라고 하면 건강한 사진을 만들어냅니다.
하지만 이 방법은 전체적인 그림을 처음부터 다시 그리는 것과 같습니다.

• 비유: "치킨 요리"를 "피자"로 바꾸라고 하면, AI 는 치킨을 없애고 피자를 새로 그립니다. 그런데 이 과정에서 접시 모양, 배경, 조명까지 모두 바뀌어버립니다.
• 문제점: AI 가 병을 찾는 훈련을 하려는데, "병"만 바뀐 게 아니라 "접시"까지 바뀌어버리면 AI 는 병을 찾는 게 아니라 "접시 모양"을 외우게 되어 엉뚱한 결론을 내립니다.

기존의 다른 방법들 (사진을 수정하는 기술) 은 원래 사진을 바탕으로 수정하려 하지만, 원래 사진의 흔적을 지우지 못해 병이 사라진 자리에 찌그러짐이나 왜곡이 생깁니다.

---

✨ 해결책: MedSteer (메드스티어) - "요리사의 마법 지팡이"

저자들은 MedSteer라는 새로운 기술을 개발했습니다. 이 기술은 아예 새로운 사진을 그리는 게 아니라, 같은 '씨앗 (노이즈)'에서 시작해서 병만 선택적으로 지우는 마법을 부릅니다.

1. 마법 지팡이 만들기 (패스톨로지 벡터)

• 상황: AI 가 "병이 있는 치킨"과 "건강한 치킨"을 그리는 과정을 관찰합니다.
• 행동: 두 그림의 차이점, 즉 '병'이라는 개념이 AI 의 뇌 (주의 메커니즘) 에 어떤 신호로 들어오는지를 분석합니다.
• 결과: 이 차이를 **"병을 지우는 마법 지팡이 (벡터)"**로 만듭니다. 이 지팡이는 "병"이라는 개념만 정확히 건드리고, "접시"나 "배경"에는 전혀 손대지 않습니다.

2. 마법 부리기 ( steering )

• 상황: 이제 AI 에게 "병이 있는 치킨"을 그리라고 명령합니다.
• 행동: 그림이 그려지는 중간중간, 마법 지팡이를 살짝 휘둘러서 '병'이라는 신호만 약하게 만들어버립니다.
• 결과: 놀랍게도 접시 모양, 배경, 조명, 장기의 전체적인 구조는 그대로 유지된 채, 오직 '병'만 사라진 건강한 치킨이 나옵니다.

---

🌟 왜 이 기술이 특별한가요?

1. 완벽한 쌍둥이 (Counterfactual):

   • 기존 방법은 병이 있는 사진과 건강한 사진이 서로 다른 그림이 될 수 있습니다.
   • MedSteer 는 같은 씨앗에서 시작해서 오직 병만 뺀 사진을 만듭니다. 마치 "병이 있는 내 사진"과 "병이 없는 내 사진"을 비교하는 것처럼, 나머지 모든 것이 100% 똑같습니다.

2. 학습이 필요 없음 (Training-Free):

   • 다른 방법들은 AI 를 다시 가르치거나 (Fine-tuning), 사진을 일일이 표시 (마스크) 해줘야 했지만, MedSteer 는 기존에 만들어진 AI 를 그대로 사용해도 됩니다. 약간의 지시만 주면 되죠.

3. 실제 의료 효과:

   • 실험 결과, 이 방법으로 만든 사진으로 AI 를 훈련시키니 종양 (폴립) 을 찾는 정확도가 크게 향상되었습니다. AI 가 병의 특징만 제대로 학습했기 때문입니다.

4. 투명한 과정 (Interpretability):

   • 이 기술은 AI 가 어디서, 언제 병을 지웠는지 시각적으로 보여줍니다. 마치 "여기서 병을 지웠어요"라고 표시해주는 것처럼, AI 의 판단 근거를 인간이 이해할 수 있게 해줍니다.

📝 한 줄 요약

MedSteer는 AI 에게 "병만 지우고 나머지는 그대로 두세요"라고 정확히 지시하는 마법 지팡이입니다. 덕분에 의료 AI 는 병의 특징만 집중해서 학습할 수 있게 되었고, 더 정확한 진단을 내릴 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

내시경 이미지 분석에서 병변 감지기를 훈련할 때, 모델이 질병 특유의 특징 (pathology-specific features) 이 아닌 해부학적 구조나 배경과 같은 교란 변수 (confounders) 에 반응하는 것을 방지하는 것이 핵심 과제입니다. 이를 해결하기 위해 '질병'과 '정상' 이미지를 해부학적으로 일치시키는 쌍 (paired images) 을 생성하는 것이 이상적이지만, 기존 방법은 다음과 같은 한계가 있습니다.

• 텍스트 프롬프트 재사용 (Re-prompting): "폴립"에서 "정상"으로 프롬프트를 변경하면 생성 모델의 전체 생성 궤적 (trajectory) 이 다시 시작되어 해부학, 질감, 배경이 완전히 달라지는 새로운 이미지가 생성됩니다.
• 역변환 기반 편집 (Inversion-based Editing): DDIM 역변환 (inversion) 을 사용하여 원본 이미지를 기반으로 편집하는 방법 (PnP, h-Edit 등) 은 재구성 오차 (reconstruction error) 를 발생시킵니다. 이로 인해 원본의 구조가 왜곡되거나 (structural drift) 배경이 변형되는 문제가 발생합니다. 또한, RadEdit 와 같은 방법은 지역별 마스크 주석이 필수적이어서 적용이 어렵습니다.

2. 제안 방법론: MedSteer

저자들은 MedSteer라는 훈련이 필요 없는 (training-free) 활성화 제어 (activation steering) 프레임워크를 제안합니다. 이 방법은 원본 이미지, 마스크 주석, 미세 조정 (fine-tuning) 없이도 생성된 이미지 쌍에서 오직 병리학적 개념만 변경하고 나머지 구조는 완벽하게 보존합니다.

핵심 구성 요소

1. 백본 및 개입 지점 (Backbone & Intervention Point):

   • 고정된 (frozen) Diffusion Transformer (DiT, PixArt-α 기반) 를 사용합니다.
   • 텍스트 의미론이 이미지 토큰에 도달하는 유일한 통로인 크로스 어텐션 (Cross-Attention, CA) 레이어를 개입 지점으로 선정합니다.

2. 병리 벡터 추정 (Pathology Vector Estimation):

   • 대조적 프롬프트 쌍: "염색된 폴립" (Positive) 과 "폴립" (Negative) 과 같이 병리 개념 (P) 만이 다른 프롬프트 쌍을 준비합니다.
   • 벡터 계산: 여러 랜덤 시드 (Z 개) 로 생성된 이미지의 CA 특징을 평균화한 후, 두 개념 간의 차이를 계산하여 단위 벡터 ($v_{l,t}$) 를 추출합니다.
   • 이 벡터는 이미지 특정이 아닌 개념 특정 (concept-specific) 이므로 오프라인에서 한 번 계산하면 추론 시 재사용됩니다.

3. 공간 선택적 병리 제어 (Spatially Selective Pathology Steering, SSPS):

   • 동일 노이즈 시드: 두 이미지 (원본 개념과 반사실 개념) 를 동일한 노이즈 시드로 생성하여 궤적을 공유하게 합니다.
   • 코사인 유사도 게이트 (Cosine-Similarity Gate, CSG): 병리학적 변화는 공간적으로 국소적이므로, 모든 토큰을 균일하게 수정하면 안 됩니다. 각 토큰의 활성화가 병리 벡터와 얼마나 정렬되어 있는지 코사인 유사도 (내적) 로 점수 ($\sigma_{l,t}$) 를 매깁니다.
   • 활성화 업데이트: 정렬된 토큰의 활성화에서 병리 성분을 감산하여 수정합니다 ($h'_{l,t} = h_{l,t} - \alpha \sigma_{l,t} v_{l,t}$). 이 과정에서 해부학, 질감, 시점 등 수직 성분 (orthogonal components) 은 그대로 유지됩니다.

3. 주요 기여 (Key Contributions)

1. 훈련 없는 반사실 쌍 생성: 원본 이미지나 주석 없이도, 동일한 노이즈 궤적을 공유하는 두 이미지를 생성하여 구조적 일치를 보장합니다.
2. 공간적 해석 가능성 (Spatial Interpretability): 토큰별 코사인 유사도 게이트를 통해 모델이 어디에서 (공간적 위치) 언제 (확산 단계) 개념을 변경하는지 시각화할 수 있는 내장된 메커니즘을 제공합니다. 이는 역변환 기반 방법들이 lacking 한 기능입니다.
3. 엔트angled(얽힌) 의료 레이블 처리: 텍스트로 분리하기 어려운 의료 클래스 레이블 (예: 폴립 형태 + 염색) 을 효과적으로 분리하고 제어할 수 있음을 입증했습니다.

4. 실험 결과 (Results)

Kvasir v3 및 HyperKvasir 데이터셋을 사용하여 세 가지 실험을 수행했습니다.

• 하류 폴립 감지 (Downstream Polyp Detection):

  • 생성된 증강 데이터를 사용하여 폴립 감지기를 훈련했을 때, MedSteer 는 ViT 기반 모델에서 AUC 0.9755를 기록했습니다.
  • 이는 재프롬프팅 (0.9083) 이나 기존 편집 방법 (PnP: 0.9518, h-Edit: 0.9312) 보다 월등히 높은 성능으로, 반사실 구조가 모델의 일반화 능력을 향상시킴을 입증했습니다.

• 반사실 생성 품질 (Counterfactual Generation):

  • 개념 전환율 (Flip Rate): 폴립→정상 (0.800), 궤양성 대장염→정상 (0.925), 식도염→정상 (0.950) 으로, 기존 최첨단 방법들보다 높은 전환율을 보였습니다.
  • 배경 보존: 역변환 기반 방법들은 재구성 오차로 인해 배경이 변형되었으나, MedSteer 는 Bg-SSIM 및 Bg-PSNR 지표에서 가장 높은 점수를 기록하여 구조적 보존이 우수함을 입증했습니다.

• 염색 분리 (Dye Disentanglement):

  • "염색된 폴립"에서 "염색 제거"를 수행할 때, MedSteer 는 염색 제거율 (DDR) 에서 **75%**의 성공률을 보였습니다 (h-Edit: 10%, PnP: 20%). 이는 텍스트로 분리 불가능한 속성 (형태 vs 염색) 을 성공적으로 분리했음을 의미합니다.

5. 의의 및 결론 (Significance)

MedSteer 는 의료 영상 생성 분야에서 구조적 일치를 유지하면서 병리학적 개념만 정밀하게 조작할 수 있는 새로운 패러다임을 제시합니다.

• 데이터 증강: 교란 변수를 제거한 고품질의 반사실 데이터를 생성하여, 질병 감지 모델의 신뢰성과 일반화 성능을 극대화합니다.
• 해석 가능성: 생성 과정의 각 단계에서 모델이 어떤 토큰을 수정하는지 시각화함으로써, 의료 AI 의 투명성을 높입니다.
• 실용성: 추가적인 훈련이나 주석이 필요 없어, 레이블이 부족한 의료 데이터 환경에서도 즉시 적용 가능합니다.

이 연구는 생성형 AI 가 단순한 이미지 생성을 넘어, 인과적 (causal) 인 의료 데이터 증강 및 설명 가능한 AI(XAI) 도구로서 핵심적인 역할을 할 수 있음을 보여줍니다.