We introduce a new method for training GANs by applying the Wasserstein-2 metric proximal on the generators. This approach is based on the gradient operator induced by optimal transport theory, which connects the geometry of the sample space and the parameter space in implicit deep generative models. From this theory, we obtain an easy-to-implement regularizer for the parameter updates. Our experiments demonstrate that this method improves the speed and stability in training GANs in terms of wallclock time and Fréchet Inception Distance (FID) learning curves.
[summary]
This paper considers natural gradient learning in GAN learning, where the Riemannian structure induced by the Wasserstein-2 distance is employed. More concretely, the constrained Wasserstein-2 metric $d_W$, the geodesic distance on the parameter space induced by the Wasserstein-2 distance in the ambient space, is introduced (Theorem 1). The natural gradient on the parameter space with respect to the constrained Wasserstein-2 metric is then derived (Theorem 2). Since direct evaluation of $G^{-1}$ poses difficulty, the authors go on to considering a backward scheme using the proximal operator (3), yielding:
(i) The Semi-Backward Euler method is proposed via a second-order Taylor approximation of the proximal operator $d_W^2$ (Proposition 3).
(ii) From an alternative formulation for $d_W$ (Proposition 4), the authors propose dropping the gradient constraint to define a relaxed Wasserstein metric $d$, yielding a simple proximal operator given by the expected squared Euclidean distance in the sample space used as a regularizer (equation (4)). The resulting algorithm is termed the Relaxed Wasserstein Proximal (RWP) algorithm.
[pros]
The proposal provides an easy-to-implement drop-in regularizer framework, so that it can straightforwardly be combined with various generator update schemes.
[cons]
Despite all the theoretical arguments given to justify the proposal, the resulting proposal may simply be viewed as a naive application of the proximal operator.
[Quality]
See [Detailed comments] section below.
[Clarity]
This paper is basically clearly written.
[Originality]
Providing justification to the proximal operator approach in GAN learning via natural gradient with respect to the Riemannian structure seems original.
[Significance]
See [Detailed comments] section below.
[Detailed comments]
To the parameter space $\Theta\subset\mathbb{R}^d$, one can consider introducing several different Riemannian structures, including the conventional Euclidean structure and that induced by the Wasserstein-2 metric. Which Riemannian structure among all these possibilities would be natural and efficient in GAN training would not be evident, and this paper discusses this issue only in the very special single instance in Section 2.3. A more thorough argument supporting superiority of the Riemannian structure induced by the Wasserstein-2 metric would thus be needed in order to justify the proposed approach.
In relation to this, the result of comparison between WGAN-GP with and without SBE shown in Figure 5 is embarrassing to me, since it might suggest that the proposed framework aiming at performing Wasserstein natural gradient is not so efficient if combined with WGAN-GP. The natural gradient is expected to be efficient when the underlying coordinate system is non-orthonormal (Amari, 1998). Starting with the gradient descent iteration derived from the backward Euler method in (3), which is computationally hard, the argument in this paper goes on to propose two methods: the Semi-Backward Euler method via a second-order Taylor approximation to the backward Euler scheme (Proposition 3), and RWP in (4) via approximation (dropping of the gradient constraint and finite-difference approximation in the integral with respect to $t$) of an alternative simpler formulation for the Wasserstein metric (Proposition 4). These two methods involve different approximations to the Semi-Backward Euler, and one would like to know why the approximations in the latter method is better in performance than those in the former. Discussion on this point is however missing in this paper.
In Section 3, it would have been better if the performance be compared not only in terms of FID but also the loss considered (i.e., Wasserstein-1), since the latter is exactly what the algorithms are trying to optimize.
Minor points:
Page 4: The line just after equation (4) should be moved to the position following the equation giving $d(\theta_0,\theta_1)^2$.
In the reference list, the NIPS paper by Gulrajani et al. appears twice.