데이터 사이의 거리: KL 발산부터 Wasserstein 거리까지

딥러닝 모델의 성능을 높이거나, 특히 연합학습(Federated Learning) 환경에서 기관 간의 데이터 불균형(Non-IID) 문제를 다룰 때, 우리가 가장 먼저 마주하는 질문은 이것입니다. "두 데이터 분포가 얼마나 다른가?"

 

이 질문에 답하기 위해 통계학과 정보이론에서는 다양한 '거리(Distance)'와 '발산(Divergence)' 개념을 사용합니다. 오늘은 가장 대표적인 4가지 지표를 물리학적 직관과 데이터 과학의 관점에서 비교해 보겠습니다.

1. KL Divergence (Kullback-Leibler Divergence)

"정보의 손실을 측정하다"

KL 발산은 정보이론의 핵심 개념으로, 실제 분포 $P$를 근사 분포 $Q$로 표현했을 때 발생하는 정보의 손실량(상대 엔트로피)을 측정합니다.

• 수식: $D_{KL}(P || Q) = \sum P(x) \log \frac{P(x)}{Q(x)}$
• 특징:
  • 비대칭성: $D_{KL}(P || Q) \neq D_{KL}(Q || P)$ 입니다. 즉, '거리'의 수학적 정의를 만족하지 못하는 '발산'입니다.
  • 정확도 측정: $Q$가 $P$를 얼마나 잘 설명하는지 측정하는 용도로 주로 쓰입니다.
• 한계: 두 분포의 서포트(Support, 확률이 0이 아닌 영역)가 다를 경우 값이 무한대로 발산하거나 정의되지 않는 문제가 있어 수치적 안정성이 떨어질 수 있습니다.

2. Jensen-Shannon Distance (JSD)

"대칭성을 갖춘 KL 발산의 진화"

JSD는 KL 발산의 비대칭성 문제를 해결하기 위해 고안되었습니다. 두 분포의 평균 분포 $M = \frac{1}{2}(P+Q)$를 상정하고, 각각과의 KL 발산을 합산합니다.

• 수식: $JSD(P || Q) = \sqrt{\frac{1}{2}D_{KL}(P || M) + \frac{1}{2}D_{KL}(Q || M)}$
• 특징:
  • 대칭성 및 유계성: 항상 $0 \leq JSD \leq 1$ (로그 밑이 2인 경우) 사이의 값을 가지며 대칭적입니다.
  • 안정성: 두 분포가 겹치지 않더라도 KL 발산처럼 무한대로 튀지 않습니다.
• 연합학습에서의 활용: 클라이언트 간의 라벨 분포 차이를 빠르게 비교할 때 유용합니다.

3. Earth Mover's Distance (EMD)

"물리학적 직관: 흙을 옮기는 비용"

EMD는 컴퓨터 비전 분야에서 주로 쓰이는 용어로, "한 분포의 진흙 더미를 다른 모양의 분포로 변형하기 위해 필요한 최소한의 일(Work)"을 의미합니다.

• 비유: 물리학에서 입자를 이동시킬 때 (질량) x (이동 거리)를 계산하는 것과 같습니다.
• 특징:
  • 단순히 특정 지점의 확률값 차이만 보는 것이 아니라, '어디서 어디로' 옮겨야 하는지 기하학적 거리를 고려합니다.
  • 데이터의 형상(Shape)이 비슷하더라도 위치가 멀리 떨어져 있다면 거리가 멀게 측정됩니다.

4. Wasserstein Distance (WD)

"수학적으로 완성된 최적 운송 이론"

수학적으로 EMD를 확률 공간으로 확장한 것이 바로 Wasserstein Distance입니다. 보통 1차 Wasserstein Distance($W_1$)를 의미하며, GAN(WGAN)의 성공 이후 AI 분야에서 가장 주목받는 지표가 되었습니다.

• 수식: $W(P, Q) = \inf_{\gamma \in \Pi(P, Q)} \mathbb{E}_{(x, y) \sim \gamma} [||x - y||]$
• 왜 Wasserstein인가? (장점):
  • 연속성: 두 분포가 전혀 겹치지 않아도 거리가 멀어짐에 따라 값이 선형적으로 증가합니다. 이는 경사 하강법(Gradient Descent)을 사용할 때 매우 중요한 속성입니다.
  • Non-IID 분석의 핵심: 연합학습에서 로컬 모델들의 가중치나 데이터 특징(Feature)의 이동을 추적할 때, 단순한 통계량 비교보다 훨씬 깊은 기하학적 통찰을 제공합니다.

한눈에 비교하기

구분	KL Divergence	JS Distance	Wasserstein (EMD)
수학적 성격	Divergence (비대칭)	Metric (대칭)	Metric (대칭)
측정 대상	정보량의 차이	분포의 겹침 정도	최적 운송 비용 (기하학)
값의 범위	$0 \sim \infty$	$0 \sim 1$	$0 \sim \infty$
주요 장점	계산이 매우 빠름	비교가 직관적이고 안정적	분포가 겹치지 않아도 유효함
한계점	비대칭성, 수치적 불안정	멀리 떨어진 분포에 둔감	계산 복잡도가 높음

맺음말: 데이터 딜레마를 해결하는 도구

연합학습에서 기관 간 데이터는 저마다의 '공간'에 존재합니다. 이 데이터들의 거리를 정확히 측정할 수 있다는 것은 곧, 모델의 편향을 정량화하고 보안 위협(데이터 오염 등)을 감지할 수 있는 기준선이 생긴다는 뜻입니다.

 

특히 Wasserstein Distance는 단순한 수치 비교를 넘어 데이터의 흐름과 기하학적 구조를 이해하게 해준다는 점에서, 물리학적 세계관을 데이터 과학에 투영하는 가장 아름다운 도구 중 하나가 아닐까 생각합니다.