KLIP: localized distribution shift detection via KL-divergence with diffusion priors in Inverse Problems

KLIP에 대한 설명: 일상적인 언어와 비유를 사용한 해설

거대한 문제: 흐릿한 사진 속에서 "다른 하나" 찾아내기

당신이 퍼즐을 풀려고 하는데, 조각들이 빠져 있고 오직 몇 개의 흐릿한 단서(흩어진 퍼즐 조각이나 번진 사진 같은 것들)만 가지고 있다고 상상해 보세요. 의료 영상의 세계에서 이것은 **역문제(inverse problem)**라고 불립니다. 의사들은 제한적이고 노이즈가 섞인 데이터(예: 몇 개의 X선 각도)로부터 환자의 내부(예: 간)를 보여주는 선명한 이미지를 재구성해야 합니다.

보통 AI 모델은 이전에 보았던 것을 바탕으로 빈 조각을 채우도록 훈련됩니다. 만약 모델이 건강한 간을 학습했다면, 흐릿한 데이터를 보았을 때 모델은 건강한 간을 "추측"하여 그려내려 할 것입니다.

도전 과제: 만약 환자에게 종양이 있다면 어떻게 될까요?

• AI는 자신이 알고 있는 대로 여전히 건강한 간을 그리려고 시속할 수 있습니다.
• 종양은 전체 이미지가 아닌 특정 지점에만 나타나는 국소적인(localized) 문제입니다.
• 기존의 AI 도구들은 여기서 실패하는 경우가 많습니다. 그들은 "이 전체 이미지가 이상하다"라고 말하거나, 무엇을 찾아야 할지 알기 위해 종양의 예시를 미리 보여줘야 합니다. 하지만 현실 세계에서는 모든 가능한 희귀 질환을 AI에게 미리 보여줄 수 없습니다.

해결책: KLIP (더 "의심 측정기")

저자들은 KLIP이라는 새로운 도구를 제안합니다. 이것을 어떤 범죄가 일어나고 있다는 것을 미리 알지 못해도 범죄를 감지해내는 의심 측정기라고 생각해보세요.

KLIP은 다음 두 가지를 비교하여 작동합니다:

1. "기대치" (사전 확률, The Prior): AI가 자신의 학습을 바탕으로 이미지가 어떻게 보여야 한다고 생각하는 것 (예: "나는 건강한 간을 기대한다").
2. "현실" (사후 확률, The Posterior): 흐릿한 단서에 강제로 맞추었을 때 이미지가 실제로 어떻게 보이는가 (예: "데이터를 보니 여기에 이상한 덩어리가 있다").

만약 "현실"이 "기대치"와 매우 다르다면, 측정기가 울립니다. 두 차이가 클수록, 무언가 특이한 것(분포 외 데이터, OOD)이 일어나고 있을 가능성이 높습니다.

핵심 비결: 확대하기와 시간 여행하기

이 논문의 주요 돌파구는 KLIP이 단순히 전체 이미지를 한꺼번에 보는 것이 아니라는 점입니다. KLIP은 작고 숨겨진 문제를 찾기 위해 두 가지 영리한 기술을 사용합니다.

1. "그리드 탐색" (공간적 제한, Spatial Restriction)

당신이 아주 작은 동전을 찾기 위해 커다란 지도를 보고 있다고 상상해 보세요. 지도 전체를 그냥 보고 있다면 동전은 보이지 않을 것입니다.

• 기존 방식: 지도 전체를 보고 "대체로 정상이다"라고 말합니다.
• KLIP 방식: 지도를 작은 격자(grid) 형태의 작은 사각형들로 나눕니다. 그리고 각 사각형을 개별적으로 확인합니다. 만약 어떤 작은 사각형 하나에서 "우리가 기대하는 것"과 "데이터가 말하는 것" 사이의 차이가 크게 나타난다면, KLIP은 그 특정 사각형을 지목합니다.
• 결과: KLip은 나머지 이미지가 완벽하더라도 간 내부의 작은 종양이나 얼굴의 흉터를 정확히 짚어낼 수 있습니다.

2. "시간 여행" 기술 (타임스텝 제한, Timestep Restriction)

디퓨전 모델(여기서 사용된 AI)은 조각가가 작품을 만드는 것처럼 이미지를 구축합니다. 처음에는 거칠고 흐릿한 덩어리에서 시작하여 시간이 지남에 따라 세부적인 디테일을 더해갑니다.

• 초기 단계: AI는 큰 형태(얼굴의 윤곽이나 간의 일반적인 모양)를 파악하는 중입니다.
• 후기 단계: AI는 미세한 디테일(눈, 피부 질감, 종양의 구체적인 모양)을 추가합니다.

KLIP은 서로 다른 유형의 "이상함"이 서로 다른 시점에 나타난다는 점을 깨달았습니다.

• 크고 전역적인 변화는 초기에 명확히 드러날 수 있습니다.
• 아주 작고 구체적인 디테일(흉터나 작은 종양 등)은 AI가 과정의 마지막 단계에서 미세한 디테일을 추가할 때 비로소 명확해질 수 있습니다.
• KLIP 방식: KLIP은 전체 과정 동안 특정 순간마다 "의심 측정기"를 확인합니다. 이는 놓치기 쉬운 미세한 디테일을 잡아내는 데 도움을 줍니다.

테스트 대상

저자들은 이 "의심 측정기"를 두 가지 주요 시나리오에서 테스트했습니다.

1. 의료 스캔 (CT 스캔):

   • 학습: AI는 건강한 간만을 보았습니다.
   • 테스트: 그들에게 인공 종양이 추가된 흐릿한 간 스캔을 주었습니다.
   • 결과: KLIP은 단순히 "이 스캔은 이상하다"라고 말하는 데 그치지 않았습니다. AI가 종양을 본 적이 없음에도 불구하고, 종양의 정확한 위치를 강조하는 히트맵(heat map)을 그려냈습니다.

2. 얼굴 (디블러링, De-blurring):

   • 학습: AI는 유명인들의 사진(건강한 얼굴)을 학습했습니다.
   • 테스트: 흉터가 추가된 얼굴의 흐릿한 사진이나, 보철물(예: 금속 눈 이식물)을 가진 SF 캐릭터의 얼굴 사진을 주었습니다.
   • 결과: KLIP은 인간의 얼굴에 있는 흉터나 캐릭터의 얼굴에 있는 금속 이식물을 성공적으로 지목하여, 이를 일반적인 특징과 구분해 냈습니다.

왜 이것이 중요한가 (논문에 따르면)

• "나쁜 데이터 학습"이 필요 없음: 대부분의 AI는 "나쁜 것"(예: 종양)을 식별하는 법을 배우기 위해 그 예시를 보여줘야 합니다. KLIP은 이것이 필요하지 않습니다. 오직 "좋은 것"이 무엇인지만 알면 됩니다.
• 흐릿한 데이터에서도 작동: KLIP은 최종적인 선명한 이미지가 아니라, 간접적인 측정값(흐릿한 단서들) 위에서 직접 작동합니다.
• 건초더미 속에서 바늘 찾기: KLIP은 나머지 이미지에 방해받지 않고 작고 국소적인 문제를 찾는 데 탁월합니다.

언급된 한계점

저자들은 몇 가지 사항에 대해 솔직하게 밝히고 있습니다:

• 튜닝(Tuning): "그리드 사각형의 크기"와 "언제 확인할 것인지"를 주의 깊게 조정해야 합니다. 설정을 잘못 선택하면 성능이 떨어질 수 있습니다.
• 시뮬레이션: 이 연구는 컴퓨터 시뮬레이션(완벽한 이미지로부터 가상의 흐릿한 이미지를 생성하는 방식)을 사용하여 테스트되었습니다. 아직 실제 병원의 복잡하고 지저도한 데이터에 적용해 보지는 않았지만, 저자들은 실제 환경에서도 작동할 것이라고 믿고 있습니다.

요약하자면: KLIP은 AI에게 "이 흐릿한 단서가 당신이 기대하는 것과 일치합니까?"라고 묻는 똑똑한 방법입니다. 만약 대답이 "아니오"라면, KLIP은 그 불일치가 발생하는 정확한 위치를 찾기 위해 줌인(zoom-in)합니다. 이 과정에서 해당 불일치를 이전에 본 적이 없어도 말이죠.