Why do we use a linear interpolation of fake and real data to penalize the gradient of discriminator in WGAN-GP

Question

I'm trying to better frame/summarize the formulations and motivations behind Wasserstein GAN with gradient penalty, based on my understanding.

For the basic GAN we are trying to optimize the following quantity:

$$\min_\theta \max_\phi \mathbb{E}_{x \sim p_{data}(x)}[D_\phi(x)] + \mathbb{E}_{z \sim p_G(z)}[1-D(G_\theta(z))]$$

The problem is that the dissimilarity measure between the two probabilities given by Jensen-Shannon divergence will not take into account any distance in a Euclidean sense. That's why we consider the Wasserstein distance defined as:

$$W(p_{data}, p_G) := \inf_\gamma \, \,\mathbb{E}_{(x,y) \sim \gamma(x,y)}\|x-y\|$$

that will account for a proper distance of our distributions. Computing it is very hard so we rely on Kantorovich-Rubinstein duality which states we can rewrite $W$ as:

$$W(p_{data},p_G) = \sup_{\|f\|_L \le 1}\mathbb{E}_{x \sim p_{data}(x)}[f_\phi(x)] - \mathbb{E}_{z \sim p_{G}(x)}[f_\phi(G_\theta(z))]$$

Now the crucial point, to enforce the constraint of $1$-Lipschitz continuity of the discriminator we add a penalty term to bound the norm of the gradient of $f$, so the final loss we consider is:

$$\mathcal{L} = \mathbb{E}_{x \sim p_{data}(x)}[f_\phi(x)] - \mathbb{E}_{z \sim p_{G}(x)}[f_\phi(G_\theta(z))] + \lambda \, \mathbb{E}_{\hat{x}}[(\|\nabla_{\hat{x}} f_\theta(\hat{x})\|-1)^2]$$

where

\begin{equation} \hat{x} = tx + (1-t)z \end{equation} $t \in [0,1]$.

Now, I've understood that we bound the slope of discriminator because we want toavoid the vanishing gradient problem and keep gradient signal in order to make the generator learn, but why do we actually penalize the gradient of discriminator with respect to a linear interpolation of real and fake data?

Answer 1

In [1], section 4, the authors mention the following:

Sampling distribution We implicitly define $p_{\hat{x}}$ sampling uniformly along straight lines between pairs of points sampled from the data distribution $p_{data}$ and the generator distribution $p_{G}$. This is motivated by the fact that the optimal critic contains straight lines with gradient norm 1 connecting coupled points from $p_{data}$ and $p_{G}$ (see Proposition 1). Given that enforcing the unit gradient norm constraint everywhere is intractable (1), enforcing it only along these straight lines seems sufficient and experimentally results in good performance (2).

So, there are 2 reasons for enforcing $\nabla f_{\phi} = 1$ over $\hat{x} = tx + (1 - t) G(z)$. Starting with point (1),

(1) Enforcing $\nabla f_{\phi} = 1$ everywhere would amount to sampling $x$ from the whole space $\mathbb{R}^{d}$ uniformly. By the curse of dimensionality, this is unfeasible.

Now, let us consider why does it suffice to enforce on samples from $p_{\hat{x}}$. Note that, due to displacement interpolation [2, Theorem 7.21], the interpolation of two probability measures under the Wasserstein metric corresponds to following geodesics in the ambient space.

In the case of GANs, let $x \sim p_{data}$ and $\tilde{x} \sim p_{model}$, where $p_{model} = G_{\sharp}p_{G}$. The interpolation $p_{t}$ of these two distributions is supported on points $\hat{x}_{t}=tx + (1-t)\tilde{x}$, that is, the distribution $p_{\hat{x}}$, for fixed $t$, is actually the interpolation in probability space of $p_{data}$ and $p_{model}$. As consequence, the authors propose enforcing Lipschitzness on a much more reduced region of $\mathbb{R}^{d}$, that is, the support of $p_{data}$, $p_{model}$ and interpolations $p_{t}$. This may be justified (although the authors do not give theoretical reasons) via the Manifold Hypothesis, that is, the data is supposed to lie on a non-linear low-dimensional region of $\mathbb{R}^{d}$ (so that we do not need to regularize over the whole space).

Remark: the book by Villani [2] is rather technical, and for a reader not familiarized with optimal transport it may be challenging to understand it. A somewhat easier text is the book by Santambrogio [3], which discusses displacement interpolation under the Wasserstein metric in chapter 5. All in all I consider that the authors had an interesting theoretical intuition, and that they validated their approach empirically, but the question on why does this work is not very clear.

References

[1] Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., & Courville, A. C. (2017). Improved training of wasserstein gans. Advances in neural information processing systems, 30.

[2] Villani, C. (2009). Optimal transport: old and new (Vol. 338, p. 23). Berlin: springer.

[3] Santambrogio, F. (2015). Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63), 94.