Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and Transport Phenomena

이 논문은 Wasserstein 거리가 데이터 분포 간 차이를 정량화할 수 있지만 그 원인을 설명하기 어렵다는 한계를 극복하기 위해, 설명 가능한 AI 기법을 활용하여 해당 거리를 데이터 하위 그룹, 입력 특성 또는 해석 가능한 부분 공간과 같은 구성 요소에 효율적으로 귀속시키는 새로운 방법을 제안합니다.

핵심

🌊 1. 문제 상황: "두 강물의 차이"를 숫자로만 알 수 있다?

가상 상황을 상상해 보세요.

• 강 A (Source): 아침에 흐르는 맑은 강물입니다.
• 강 B (Target): 오후에 흐르는 강물인데, 중간에 폭포가 있고, 물살이 세지며, 어떤 곳은 모래가 쌓여 있습니다.

우리는 이 두 강물의 상태를 비교할 때, 단순히 "두 강물의 거리는 50km 입니다"라고 숫자만 알려준다면 어떨까요?

• "어디가 50km 차이 나는 거지?"
• "폭포 때문일까? 모래 때문일까?"
• "물살이 세진 건 왜일까?"

기존의 워터스틴 거리는 두 강물 (데이터) 의 차이를 계산하는 아주 정교한 자입니다. 하지만 이 자는 **"차이의 총합"**만 알려줄 뿐, **"어떤 부분이 그 차이를 만든 주범인지"**는 알려주지 않습니다. 마치 "이 집의 전기세는 10 만 원이다"라고만 알려주고, "에어컨이 8 만 원, 냉장고가 2 만 원이다"라고 알려주지 않는 것과 같습니다.

💡 2. 해결책: "WaX (워터스틴 설명기)"의 등장

이 논문은 **WaX (Wasserstein Distances Made Explainable)**라는 새로운 방법을 제안합니다.
WaX 는 **"전기세 명세서"**처럼, 두 데이터 세트의 차이 (워터스틴 거리) 가 어떤 특징 (Feature) 이나 어떤 데이터 포인트 때문에 생겼는지 세세하게 설명해 줍니다.

🏗️ 비유: 건물의 구조를 뒤집어 보기

WaX 는 복잡한 수식을 마치 레고 블록으로 만든 기계처럼 다시 조립합니다.

1. 기계 분해 (Neuralization): 워터스틴 거리를 계산하는 복잡한 과정을, 신경망 (Neural Network) 이라는 레고 구조로 바꿉니다.
2. 역주행 (Propagation): 기계가 작동한 결과 (거리 값) 를 보고, 그 원인이 된 레고 블록 (데이터의 특징) 들로 거꾸로 추적합니다.
   • "아! 이 레고 블록 (예: '키'라는 특징) 이 30% 를 차지했구나!"
   • "이 블록 (예: '나이'라는 특징) 이 70% 를 차지했구나!"

이렇게 하면 **"두 데이터가 다른 이유는 '키' 차이 때문이지, '나이' 때문이 아니야!"**라고 명확하게 알 수 있게 됩니다.

---

🚀 3. WaX 가 실제로 어떤 일을 해내나? (세 가지 사례)

이 기술은 실제로 매우 유용하게 쓰입니다.

🛠️ 사례 1: "나쁜 친구"를 찾아내서 버리기 (도메인 적응)

• 상황: 한 병원에서 만든 AI 가 다른 병원에 적용될 때, 병원의 벽색이나 조명 같은 '불필요한 차이' 때문에 오작동할 수 있습니다.
• WaX 의 역할: "이 AI 가 병원의 벽색 (불필요한 특징) 때문에 판단을 잘못하고 있어! 이 부분을 잘라내자!"라고 알려줍니다.
• 결과: 불필요한 특징을 제거하면 AI 는 더 똑똑하고 튼튼해집니다.

🐌 사례 2: "오징어"가 어떻게 자라나? (이해하기 어려운 현상 분석)

• 상황: 작은 오징어 (데이터) 가 1 년 뒤 커진 오징어로 변하는 과정을 관찰합니다.
• WaX 의 역할: 단순히 "커졌다"가 아니라, **"작은 오징어들은 '무게'가 많이 늘었지만, 큰 오징어들은 '길이'가 더 많이 늘었다"**는 식으로, 그룹별로 어떻게 변했는지 세분화해서 보여줍니다.
• 결과: 복잡한 생물의 성장 과정을 '하드'하게 쪼개서 이해할 수 있게 됩니다.

📸 사례 3: "사진첩"의 차이 찾기 (데이터셋 비교)

• 상황: '셀럽 사진 (CelebA)'과 '야외 인물 사진 (LFW)' 두 가지 사진 데이터가 있습니다.
• WaX 의 역할: 두 사진첩의 차이를 분석합니다.
  • "셀럽 사진에는 여배우들이 많고, 야외 사진에는 정치인이나 남성이 많네."
  • "셀럽 사진에는 선글라스를 쓴 사람이 많고, 야외 사진에는 테니스를 치는 사람이 있네."
• 결과: 데이터 세트가 어떤 편향을 가지고 있는지, 어떤 주제가 빠졌는지 한눈에 파악할 수 있습니다.

---

🌟 4. 왜 이것이 중요한가요?

과거에는 "두 데이터가 다르다"는 사실만 알 수 있었습니다. 하지만 WaX는 "왜, 어디서, 어떻게" 다른지 이해할 수 있게 (Explainable) 만들어줍니다.

• 투명성: AI 가 왜 그런 판단을 내렸는지, 데이터가 왜 다른지 그 이유를 투명하게 보여줍니다.
• 신뢰성: 의료나 과학 같은 중요한 분야에서, 데이터의 차이를 정확히 이해해야만 신뢰할 수 있는 결정을 내릴 수 있습니다.
• 효율성: 불필요한 데이터를 제거하고, 진짜 중요한 부분에만 집중하게 도와줍니다.

📝 한 줄 요약

"WaX 는 두 데이터 세트의 '거리'를 계산하는 자를, 그 거리가 '왜' 생겼는지 설명해주는 '해설가'로 만들어주는 기술입니다."

이 기술 덕분에 우리는 이제 데이터의 차이를 단순히 숫자로만 보는 것이 아니라, 그 이면에 숨겨진 이야기와 원인을 깊이 있게 이해할 수 있게 되었습니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: Wasserstein 거리 (또는 Optimal Transport, OT) 는 두 데이터 분포 간의 차이를 측정하는 강력한 도구로, 시간에 따른 프로세스 분석이나 데이터 내의 불균일성 탐지에 널리 사용됩니다.
• 문제점:
  • 기존에는 단순히 Wasserstein 거리의 값이나 이를 계산하는 과정에서 얻어지는 수송 계획 (Transport Plan, Coupling, $\gamma^\star$) 만을 분석했습니다.
  • 그러나 수송 계획만으로는 어떤 데이터 하위 그룹, 입력 특징 (Feature), 또는 해석 가능한 부분 공간 (Subspace) 이 높은 (또는 낮은) Wasserstein 거리에 기여하는지 구체적으로 파악하기 어렵습니다.
  • 특히, 수송 계획은 특정 모델 사양 (예: $p, q$ 파라미터) 에 따라 민감하게 변할 수 있는 '병목 현상'이나 '이상치'의 영향을 제대로 반영하지 못해, 데이터 시프트의 근본 원인을 설명하는 데 한계가 있었습니다.
• 목표: Wasserstein 거리를 구성하는 다양한 데이터 요소 (개별 데이터 포인트, 입력 특징, 개념적 하위 공간) 에 대한 기여도 (Attribution) 를 정량화하여, 거리 값이 왜 발생했는지를 설명 가능한 AI (XAI) 기법을 통해 해석하는 방법론을 제시하는 것.

2. 제안 방법론: WaX (Methodology)

저자들은 WaX (Wasserstein distances made explainable) 라는 새로운 프레임워크를 제안합니다. 이는 Layer-wise Relevance Propagation (LRP) 과 신경망화 (Neuralization) 기법을 Wasserstein 거리에 적용한 것입니다.

핵심 절차

1. 신경망화 (Neuralization):
   • Wasserstein 거리 계산 ($W_p$) 을 최적 수송 문제의 쌍대 형식이 아닌 원형 (Primal) 형식으로 재정의합니다.
   • 최적 커플링 행렬 $\gamma^\star$ 를 고정하고, 입력 분포와 $W_p$ 사이의 함수적 관계를 단순화하여 2 계층 신경망 구조로 변환합니다.
     • Layer 1: 두 분포의 인스턴스 쌍 $(x_k, y_l)$ 간의 거리 계산 ($z_{kl} = \|x_k - y_l\|_q$).
     • Layer 2: 커플링 가중치 $\gamma^\star$ 를 적용하여 거리들을 가중 합산하고 $p$-norm 을 취하여 최종 $W_p$를 계산.
2. 전파 (Propagation - LRP):
   • 변환된 신경망 구조를 통해 역전파 (Backward Pass) 를 수행하여 출력값 ($W_p$) 을 입력 특징과 인스턴스 쌍으로 분해합니다.
   • LRP 규칙 정의:
     • 인스턴스 쌍 기여도 ($R_{kl}$): $R_{kl} = \frac{\gamma^\star_{kl} \cdot z_{kl}^\alpha}{\sum \gamma^\star_{kl} \cdot z_{kl}^\alpha} W_p$
     • 입력 특징 기여도 ($R_i$): $R_i = \sum_{kl} \frac{|x_{ki} - y_{li}|^\beta}{\sum_i |x_{ki} - y_{li}|^\beta} R_{kl}$
   • 여기서 $\alpha$ 와 $\beta$ 는 하이퍼파라미터로, 설명의 집중도 (어떤 샘플이나 특징에 더 많은 관련성을 할당할지) 를 조절합니다. 저자들은 $\alpha = p$, $\beta = \min(p+2, q)$ 라는 휴리스틱을 제안하여 비선형성이 강한 모델에서도 안정적인 설명을 제공합니다.

확장: U-WaX (Subspace-based Explanations)

• U-WaX는 WaX 를 부분 공간 (Subspace) 기반 설명으로 확장한 것입니다.
• 입력 공간을 직교 행렬 $U$를 통해 여러 개념 (Concept) 이나 하위 시프트 (Sub-shift) 로 분해합니다.
• Wasserstein 거리를 각 하위 공간별 기여도로 분해하여, 복잡한 수송 현상을 여러 개의 독립적인 요인으로 해체 (Disentanglement) 할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

1. 최초의 Wasserstein 거리 설명 기법: 데이터 분포 간 거리를 개별 특징이나 인스턴스 수준에서 설명하는 체계적인 XAI 방법론을 최초로 제시했습니다.
2. 이론적 및 계산적 효율성:
   • 보존성 (Conservation): 할당된 관련성 점수의 합이 원래 Wasserstein 거리와 정확히 일치합니다 ($\sum R_i = W_p$).
   • 경사도 연결: 특정 파라미터 설정 시 LRP 규칙이 경사도 (Gradient) 계산과 수학적으로 동치임을 증명했습니다.
   • 효율성: 특징 제거 (Occlusion) 기반 방법보다 계산 비용이 훨씬 낮으며, 대규모 데이터셋에도 적용 가능합니다.
3. 유연한 해석 가능성:
   • 데이터 포인트, 입력 특징, 그리고 추상적인 개념 (Subspace) 단위로 설명을 제공할 수 있습니다.
   • 다양한 Wasserstein 모델 ($W_1, W_2$, Sinkhorn 등) 과 거리 메트릭 ($L_1, L_2, L_\infty$) 에 적용 가능합니다.

4. 실험 결과 (Results)

저자들은 다양한 데이터셋과 시나리오에서 WaX 의 성능을 검증했습니다.

• 설명 충실도 (Faithfulness):
  • SRG (Symmetric Relevance Gain) 지표를 사용하여 평가했습니다. WaX 는 MeanShift, Occlusion, Coupling 기반 베이스라인보다 일관되게 높은 점수를 기록했습니다.
  • 특히 $p, q$ 값이 커져 비선형성이 강해지거나 이상치에 민감한 경우, 기존 방법들은 실패했으나 WaX 는 모델의 민감도를 정확히 반영했습니다.
• 수송 현상 특성화:
  • 시계열 데이터 (Air Quality, Electricity 등) 에서 실제 이동 (Transport) 이 발생한 차원을 정확히 식별했습니다. 기존 분류기 기반 방법들은 분포가 겹칠 때 성능이 떨어졌으나, WaX 는 우수한 성능을 보였습니다.
• 실제 활용 사례 (Use Cases):
  1. 도메인 적응 (Domain Adaptation): WaX 를 통해 도메인별 편향 (Batch effects) 을 유발하는 특징을 식별하고 제거하여 분류기 robustness 를 향상시켰습니다.
  2. 노화 현상 분석 (Aging Phenomenon): 가재 (Abalone) 데이터셋을 이용해 1 년 간격의 변화를 분석했습니다. U-WaX 를 통해 전체적인 노화뿐만 아니라, 크기별/무게별 하위 그룹이 서로 다른 방식으로 성장하는 복잡한 패턴을 분리해냈습니다.
  3. 데이터셋 차이 탐지 (Dataset Differences): CelebA 와 LFW 얼굴 데이터셋을 비교하여, 인종/성별 편향, 안경 착용 여부, 스포츠 관련 배경 등 구체적인 시각적/의미적 차이를 발견하고 텍스트 단어로 설명했습니다.

5. 의의 및 결론 (Significance)

• 블랙박스 해석의 확장: Wasserstein 거리는 단순히 "차이가 있다"는 수치만 제공했으나, WaX 를 통해 "어떤 특징이, 어떤 방식으로, 얼마나 큰 영향을 미쳐 차이가 발생했는지" 를 정량적으로 설명할 수 있게 되었습니다.
• 실용적 가치:
  • 모델 검증: 사용자가 설정한 Wasserstein 모델이 의도한 대로 (예: 이상치에 민감하게, 혹은 특정 특징에 집중하여) 작동하는지 검증하는 도구로 활용 가능합니다.
  • 데이터 품질 관리: 대규모 데이터셋 간의 미세한 편향을 발견하여 데이터 정제 (Data Curation) 및 모델 학습 전략 수립에 기여합니다.
  • 과학적 통찰: 복잡한 물리적 현상이나 생물학적 과정 (예: 세포 발달, 노화) 을 분포 수준에서 분석할 때, 그 메커니즘을 해석 가능한 형태로 도출해냅니다.

이 연구는 Optimal Transport 와 Explainable AI 의 교차점을 개척하여, 데이터 과학자들이 분포 기반 분석을 더 깊이 있고 신뢰할 수 있게 수행할 수 있는 새로운 도구를 제공했습니다.