Gradient Penalties and Convergence

Dive into Deep Learning · §16.4

Gradient penalties and convergence
the Dirac-GAN · circles and spirals · zero-centered penalties · the R3GAN recipe

Equilibrium Values Do Not Ensure Convergence

Two questions remain after changing the objective:

  • The pairing objective removes mode-dropping basins but still saturates on disjoint supports, where its generator gradient vanishes.
  • Best-response analysis characterizes the equilibrium, but training uses coupled gradient steps whose convergence requires a separate analysis.

The Dirac-GAN makes both issues exact for two point masses.

The Dirac-GAN

p = \delta_0, q_\theta = \delta_\theta, linear critic D_\psi(x) = \psi x:

V(\theta, \psi) = \ell(0) + \ell(-\psi\theta)

Simultaneous descent–ascent flow:

\dot\theta = \psi\,\ell'(-\psi\theta), \qquad \dot\psi = -\theta\,\ell'(-\psi\theta)

  • The field is orthogonal to the position: \tfrac{d}{dt}(\theta^2 + \psi^2) = 0.
  • Exact circles; eigenvalues \pm i\,\ell'(0) — rotation, no attraction.

Discretization Turns Circles into Outward Spirals

One simultaneous gradient step, by Pythagoras (field \perp position):

\|(\theta, \psi) + \eta v\|^2 = \|(\theta, \psi)\|^2 + \eta^2 \|v\|^2

  • Every step increases the distance to the solution — for every step size.
  • Update-map eigenvalues 1 \pm i\eta\,\ell'(0): spectral radius > 1 always.
  • Non-saturating weighting and the pairing objective leave the field at the equilibrium unchanged — the failure is in the dynamics, not the objective.

Zero-Centered Penalties Add Damping

R_1 = \tfrac{\gamma}{2} E_p\big[\|\nabla_x D\|^2\big], \quad R_2 = \tfrac{\gamma}{2} E_q\big[\|\nabla_x D\|^2\big], \quad R_1 + R_2 = \gamma\, E_m\big[\|\nabla_x D\|^2\big]

On the Dirac-GAN: \nabla_x D_\psi = \psi, so either penalty is \tfrac{\gamma}{2}\psi^2 — the two coincide here, and using both doubles the damping (\gamma \to 2\gamma) — giving

\lambda_{1,2} = -\frac{\gamma}{2} \pm \sqrt{\frac{\gamma^2}{4} - \ell'(0)^2}

  • Negative real part for every \gamma > 0; critically damped at \gamma = 2|\ell'(0)|.
  • (R3GAN’s Eq. 12 prints this without the square — a typo; the Jacobian’s determinant is \ell'(0)^2.)

The Penalized Game Measures Linearized W2

Near equilibrium every payoff linearizes to a\,\langle p - q, D\rangle, and

\sup_D \Big\{ a \langle p - q, D\rangle - \gamma \int m \|\nabla_x D\|^2 \Big\} = \frac{a^2}{4\gamma}\, \|p - q\|^2_{\dot H^{-1}(m)}

This is the squared linearized W_2 distance. It depends on the difference p - q rather than the density ratio and remains locally sensitive to support displacement.

penalty at q = p geometry Dirac test
one-centered (WGAN-GP) rewards unit slope W_1 fails
zero-centered (R_1, R_2) rewards flat critic linearized W_2 converges

When One Penalty Is Not Enough

  • Near equilibrium, the supports overlap, so either penalty can control both measures; the local theorem requires only one.
  • Far from equilibrium, the penalty acts like implicit smoothing (Roth et al.). Smoothing p alone leaves critic gradients on generated samples uncontrolled.
  • In the R3GAN StackedMNIST ablation, R_1 alone diverges for both objectives over \gamma \in [0.1, 100]. StyleGAN2 nevertheless trains successfully on FFHQ with R_1 alone.

The Full Loss in Code

The implemented loss combines the pairing objective, a non-saturating generator update, and both penalties. The penalty function returns unscaled per-sample values of \|\nabla_x D\|^2; the caller applies the mean and \gamma/2 at every critic step.

def rpgan_loss_D(critic, real, fake):
    """Relativistic pairing loss for the critic: -E[log sigma(D(x) - D(y))]."""
    return jax.nn.softplus(critic(fake) - critic(real)).mean()


def rpgan_loss_G(critic, real, fake):
    """Non-saturating pairing loss for the generator."""
    return jax.nn.softplus(critic(real) - critic(fake)).mean()


def r1_r2_penalty(critic, real, fake):
    """Per-sample squared critic input gradients on real (R1) and fake (R2),
    before the gamma/2 scale."""
    def sq_grad_norm(x):
        x = jax.lax.stop_gradient(x)
        grad_fn = jax.grad(lambda xi: critic(xi[None, ...]).squeeze())
        grad = jax.vmap(grad_fn)(x)
        return (grad.reshape(x.shape[0], -1) ** 2).sum(axis=1)
    return sq_grad_norm(real), sq_grad_norm(fake)

Phase Portraits Confirm the Eigenvalues

All panels start from (\theta, \psi) = (1, 1). The flow retraces a circle, whereas the discrete iterates cross outward on every revolution. With the penalty, \gamma = 0.3 produces an inward spiral and \gamma = 1 produces a critically damped approach.

Mode Coverage on 25 Gaussians

GAN: 20/25 modes, reverse KL 0.75, on-mode fraction 0.05
    off-mode mass 0.95, nearest-center distance mean 0.77 / median 0.78, 26-bin KL 4.09
GAN + $R_1 + R_2$: 25/25 modes, reverse KL 0.19, on-mode fraction 0.03
    off-mode mass 0.97, nearest-center distance mean 0.69 / median 0.71, 26-bin KL 4.24
RpGAN + $R_1 + R_2$: 25/25 modes, reverse KL 0.22, on-mode fraction 0.04
    off-mode mass 0.96, nearest-center distance mean 0.67 / median 0.69, 26-bin KL 4.18
  • Plain GAN: stable training, but it never reaches all 25 modes and uses them unevenly (markedly higher reverse KL).
  • Either penalized configuration covers all 25 modes in the displayed runs and uses them more evenly. The two penalized losses are indistinguishable at this scale.
  • Every configuration places only a small fraction of its mass within 3\sigma of the centers. The experiment measures improved reach and balance, not a close fit to the mixture.

What the Toy Cannot Show

The RpGAN-vs-GAN coverage gap needs many modes under capacity pressure — StackedMNIST (cited):

loss modes / 1000 reverse KL
RpGAN + R_1 + R_2 1000 0.078
GAN + R_1 + R_2 693 0.927
either + R_1 only diverged —

In the toy experiment, the penalties improve reach and balance. The additional coverage advantage of pairing is supported by the cited StackedMNIST result. The appropriate \gamma is dataset-dependent (0.05–150 across R3GAN’s benchmarks).

Recap

  • A valid objective need not yield convergent training: Dirac-GAN flow follows circles, and simultaneous discrete updates spiral outward.
  • Zero-centered penalties give local convergence for every \gamma > 0, with critical damping at \gamma = 2|\ell'(0)|.
  • Near equilibrium, the penalized game is a scaled linearized W_2^2 distance, based on a difference of measures rather than a density ratio.
  • Far from equilibrium, symmetric regularization smooths both distributions.
  • In the experiments, the penalties govern convergence and balance; the StackedMNIST comparison attributes additional mode coverage to pairing.