JS Divergence & Wasserstein Distance

Comparing Probability Distance Metrics

Given two probability distributions $P$ and $Q$:

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│  1. KL DIVERGENCE        │  2. JS DIVERGENCE        │ 3. WASSERSTEIN DISTANCE  │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ D_KL(P || Q)             │ D_JS(P || Q)             │ W(P, Q)                  │
│ Asymmetric: D(P||Q) ≠    │ Symmetric & Bounded      │ Earth Mover Distance     │
│ D(Q||P). Explodes if     │ Range: [0, ln(2)].       │ Smooth gradients even    │
│ distributions don't overlap.| Constant for disjoint P, Q.| for disjoint support!   │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

1. KL vs JS Divergence

KL Divergence:

$$D_{KL}(P \parallel Q) = \int P(x) \ln \frac{P(x)}{Q(x)} dx$$

• Unbounded ($0 \to \infty$).
• Asymmetric ($D_{KL}(P \parallel Q) \neq D_{KL}(Q \parallel P)$).

Jensen-Shannon Divergence (Symmetrized KL):

$$M = \frac{1}{2}(P + Q)$$

$$D_{JS}(P \parallel Q) = \frac{1}{2} D_{KL}(P \parallel M) + \frac{1}{2} D_{KL}(Q \parallel M)$$

• Symmetric ($D_{JS}(P \parallel Q) = D_{JS}(Q \parallel P)$).
• Bounded between $0$ and $\ln(2) \approx 0.693$.

The Disjoint Support Problem in Classic GANs

Consider two parallel lines in 2D space: Real distribution $P$ at $x = 0$, Generated distribution $Q_\theta$ at $x = \theta$.

  Real P (x = 0)                   Generated Q_θ (x = θ)
       │                                  │
       │                                  │
       ▼                                  ▼
  Disjoint Support (Zero Overlap for θ ≠ 0)

• $D_{KL}(P \parallel Q_\theta) = \infty$ (Explodes!).
• $D_{JS}(P \parallel Q_\theta) = \ln(2) = \text{Constant}$.

Because JS Divergence is constant $\ln(2)$ for any $\theta \neq 0$, the derivative $\frac{\partial D_{JS}}{\partial \theta} = 0$.

The Generator receives ZERO GRADIENT to move toward $x = 0$!

2. Wasserstein Distance (Earth Mover's Distance)

Wasserstein-1 distance $W(P, Q)$ measures continuous transportation cost:

$$W(P, Q_\theta) = |\theta|$$

Notice that derivative $\frac{\partial W}{\partial \theta} = 1$ everywhere!

Wasserstein distance provides a smooth, linear gradient even when distributions have zero overlap!

  Distance Metric vs Parameter θ
  Metric
   ln(2) ┼─────────────────────── JS Divergence (Constant -> ZERO Gradient!)
         │                     /
         │                    /
       0 ┴───────────────────/──► Parameter θ
                                  Wasserstein Distance W(P, Q) = |θ| (SMOOTH GRADIENTS!)

WGAN & 1-Lipschitz Constraint

By Kantorovich-Rubinstein Duality:

$$W(P_r, P_g) = \sup_{|f|L \le 1} \mathbb{E}{x \sim P_r}[f(x)] - \mathbb{E}_{y \sim P_g}[f(y)]$$

Critic network $f_w(x)$ must be 1-Lipschitz Continuous:

$$|\nabla_x f_w(x)| \le 1 \quad \forall x$$

Enforced in WGAN-GP by adding a Gradient Penalty loss:

$$\mathcal{L}{\text{GP}} = \mathbb{E}{\hat{x}} \left[ \left( |\nabla_{\hat{x}} f_w(\hat{x})|_2 - 1 \right)^2 \right]$$

Say this out loud

KL divergence is asymmetric and explodes for disjoint distributions. JS divergence symmetrizes KL but saturates to constant ln(2) when distributions do not overlap, causing vanishing gradients in classic GANs. Wasserstein distance measures Earth Mover transport cost, providing smooth non-zero gradients even for disjoint distributions when the Critic network satisfies 1-Lipschitz continuity.

Follow-ups to expect

• What is Mode Collapse in GANs? Occurs when generator G learns to produce only a single high-probability image mode (e.g. outputting only 1 digit in MNIST) to fool discriminator D. WGAN eliminates mode collapse by providing continuous Wasserstein gradient feedback.
• What is Sliced Wasserstein Distance (SWD)? Projects high-dimensional probability distributions onto 1D random linear lines, computing 1D Wasserstein distance in closed form via simple sorting.