\section{\label{sec:level1}Introduction}
In the past decade, deep generative models have proven themselves to be an impressive tool for sampling from enormously complex distributions. In particular, generative adversarial networks (GANs) have been used to produce realistic data, and represent a powerful framework for training generative models \cite{gan_review, style_based_gan, scene_generation, face_generation}. Consequently, understanding and improving the training of GANs is of considerable interest.

GANs comprise two neural networks: one called the \textit{generator}, $G_\theta$, and the other called the \textit{discriminator}, $D_\phi$ (parameterized by $\theta$ and $\phi$, respectively).
\begin{align}
    \text{Generator }& G_\theta: \mathbb{R}^n \rightarrow \mathbb{R}^d \\
    \text{Discriminator }& D_\phi: \mathbb{R}^d \rightarrow \mathbb{R}
\end{align}
The generator is a function which maps randomly selected points in the latent space to points in data-space.
The discriminator assigns scores to these simulated data-points, as well as to genuine samples from the data-set. During training, the discriminator's goal is to distinguish real data from simulated data (through high and low scores, respectively), while the generator's goal is to increase the score assigned to its outputs by the discriminator \cite{gan_paper, deep_learning_book, wgan, mmd_gan}.

Although GANs are both powerful and popular, they are notoriously hard to train.
The adversarial nature of the dynamics distinguishes a GAN's `loss', $\mathcal{L}(\phi, \theta)$, from a standard loss function -- one that is bounded from below and which the training algorithm seeks to minimize. Rather than living at the minimum, the ideal parameter settings here are at the saddle points of the loss landscape \cite{gan_paper}:
\begin{equation}
    \theta^* = \argmin_\theta \max_\phi \mathcal{L}(\phi, \theta).
\end{equation}
Convergence to such a fixed-point is difficult to attain, as it requires a careful balancing of the two competing networks during training.

One important form of nonconvergence commonly encountered during GAN training is known as \textit{mode collapse} \cite{veegan, mode_reg}. Mode collapse occurs when samples from the generator fail to capture the full diversity of modes present in the data-set. Instead, the generator's output `collapses', and it produces samples from relatively few of the available modes in the data distribution.

When mode collapse occurs, during training the generator will focus its distribution on a small subset of the overall data-distribution. Eventually, the discriminator learns to identify the concentrated output of the generator, at which point the generator will switch from its current specialization to another \cite{veegan, mode_reg}. The generator's output switching from mode to mode, rather than converging on the distribution as a whole, is a key symptom of mode collapse.

Many practically useful training techniques for avoiding mode collapse have been proposed, often involving modified objective functions and novel regularizers \cite{veegan, wgan, gan_paper, gan_review}.
Here, rather than constructing empirical methods for reducing mode collapse,
we seek to understand this phenomenon from the perspective of dynamical systems, determine the physical meaning of competing factors, and derive principles to guide the training of GANs.

\iffalse
We motivate an abstracted model for GANs in which the generator is replaced by $N$ particles in $\mathbb{R}^d$, echoing the simplification applied within Sobolev Descent \cite{sobolev}. Here, however, we additionally incorporate a static neural tangent kernel (NTK) -- causing the particles to correspond to the outputs of an infinite width generator neural network.
\fi

We map GANs to an effective model in which the output of the generator network is replaced by $N$ particles in $\mathbb{R}^d$, descending the loss landscape set by both the discriminator's score function and the collective state of $N$ particles. We additionally incorporate a static neural tangent kernel (NTK) -- a \iffalse key \fi feature of realistic GANs using an infinite-width generator. Within our effective model, the NTK induces a dependence of the velocity of any particle on the current state of all other particles.

As a result of the sampling procedure of generators, and the form of common NTKs, we show the presence of a type of universality -- in which many different types of infinite-width generator neural networks map to the same particle dynamics.

We argue that this effective model provides a simplified and interpretable context in which to understand mode collapse.
In particular, applying this model to a low-dimensional toy-problem, we show a transition between the regimes of mode collapse and convergence as a function of the NTK and the relative training time. We provide a physical interpretation in which we explain this transition in terms of the properties of the discriminator.

Finally, we also use this model to study GAN regularization --  the modification of the training objective in order to promote convergence. We indeed find that when a regularizer is introduced, it results in the reduction of mode collapse in our model GAN.
Additionally, by sweeping over regularization strengths, we are able to observe under-regularized, over-regularized, and critically-regularized regimes.

\iffalse
This further encourages the use of our model as a realistic proxy for true GANs.
\fi

\begin{figure}[H]
    \centering
    \includegraphics[width=.45\textwidth]{dynamics_diagram.pdf}
    \caption{
    An illustration of how input $z$ vectors (sampled from a high-dimensional sphere) map to particles in data-space.
    During GAN training, generator parameters, $\theta_t$ (upper left plot), evolve over time. Therefore, given fixed inputs $\{ z \}$ (bottom left), the set of data-space outputs also evolves in time: $\{X_t\}$. It is the dynamics of these points in data-space which we deal with directly here.}
    \label{fig:dynamics}
\end{figure}

\section{Training and Failure}

In GANs, the generator is a neural network, $G_{\theta}$, which is fed random inputs, $z \in \mathbb{R}^n$, selected from some noise distribution $q(z)$. Samples $G_{\theta}(z) \in \mathbb{R}^d$ are therefore implicitly sampled from the generator distribution, $p_{\theta}$. Different GAN implementations exist, many with distinct loss functions \cite{general_gan}. Here we consider loss functions of the form \cite{orig_mmd_gan, mmd_gan, wgan, general_gan}:
\begin{equation}
    \mathcal{L} \equiv \langle D_\phi(x) \rangle_{x\sim p(x)} - \langle D_\phi(G_\theta(z)) \rangle_{z \sim q(z)} - \lambda R(D_\phi),
\end{equation}
where $p(x)$ is the probability distribution of samples in the data-set, $q(z)$ is the probability distribution from which seeds in the latent space are sampled, and $R(D_\phi)$ represents a regularizer on the discriminator, limiting its magnitude under a norm of interest \footnote{Note, for example, that if $R(f_\phi) = ||f||^2_{\mathcal{H}}$ this becomes an MMD GAN, while if $R(f_\phi) = ||f||^2_{K}$, we have a Wasserstein GAN \cite{wgan, demystifying}} with $\lambda \geq 0$ the magnitude of the regularizer.

{\color{red} The way the loss function $\mathcal{L}$ is written suggests that it is Wasserstein GAN we are using, the original GAN uses KL divergence, which leads to $<log(D)>$ terms in the loss function. But, the footnote seems to suggest otherwise, i.e., it is WGAN only when the regularizer $R$ has certain form.  Please clarify.}
{\color{blue}
Wasserstein GANs use the fact (the Kantorovich-Rubinstein duality) that
$$
\sup_{|f|_{L} \leq 1} \langle f(x) \rangle_p - \langle f(x) \rangle_q
$$
is equal to the wasserstein (or earthmover's) distance between $p$ and $q$. This term: $|f|_{L}$, indicates that the sup is taken over all 1-lipshitz functions -- basically limiting their gradients. In a w-gan, the wasserstein distance is estimated by learning an $f$ (like a discriminator) which then provides us with an estimate of the wasserstein distance. We practically do this by limiting the discriminator's gradients clipping weights or something, so they dont get too large. Then, by backpropagating through f to minimize this value, we are incrementally reducing the estimate of the wasserstein distance.

However, if the sup is taken over a different class of functions (say the function norm is restricted to be below some value), and we maximize this
$$
\sup_{||f|| \leq 1} \langle f(x) \rangle_p - \langle f(x) \rangle_q
$$
we end up with a different distance function between $p$ and $q$ (some other distance besides wasserstein distance). The choice of what we're supp-ing over informs what that $R(D_\phi)$ term is, and so R determines the kind of distance we're minimizing.
} {\color{red} The paragraph of the text in question in particular would be a useful thing for Youssef to look over -- just to make sure I haven't accidentally said something wrong or imprecise}

Conceptually, the discriminator parameters, $\phi$, evolve to maximize this quantity, while the generator parameters evolve to minimize it.
\begin{equation}
    \dot{\phi} = \alpha_D \frac{d \mathcal{L}}{d \phi}, \quad
    \dot{\theta} = -\alpha_G \frac{d \mathcal{L}}{d \theta}.
\end{equation}
The discriminator and generator evolution each occur at individual learning rates $\alpha_D$ and $\alpha_G$.

Practically, in neural networks a loss function is defined at each training step using mini-batches of $N$ samples of real data and generated data, both of which are re-sampled for each training step:
\begin{equation}
    \mathcal{L}_N \equiv \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i)) - \lambda R(D_\phi).
\end{equation}
Training is performed in iterations. First, for $n_{disc.}$ steps, the discriminator is updated using stochastic gradient ascent:
\begin{equation}
    \phi \xleftarrow[]{} \phi + \alpha_D \nabla_\phi \mathcal{L}_N.
\end{equation}
Then, for a single step, the generator is updated analogously:
\begin{equation}
    \theta \xleftarrow[]{} \theta - \alpha_G \nabla_\theta \mathcal{L}_N.
    \label{eq:generator_update}
\end{equation}
The resulting stochastic gradient ascent and descent are repeated until convergence, or until training has been halted after a large number of iterations.

\iffalse
\begin{algorithm}
\caption{The GAN training algorithm}\label{alg:gan_training}
\begin{algorithmic}
\For{iteration number}
    \For{$n_{disc.}$}
        \State $\bullet$ Sample $N$ data-points, $\{ x_i \}$
        \State $\bullet$ Sample $N$ seed points, $\{ z_i \}$ from the latent space according to $q(z)$

        \State$\bullet$  Compute
        $$
        \mathcal{L}_N = \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i)) - \lambda R(D_\phi)
        $$
        \State and update discriminator parameters by ascending its stochastic gradient
        $$
        \phi \gets \phi + \alpha_D \nabla_\phi \mathcal{L}_N
        $$
    \EndFor
    \State $\bullet$ Sample $N$ seed points, $\{ z_i \}$ from the latent space according to $q(z)$

    \State $\bullet$ Compute
    $$
    \mathcal{L}^{gen.}_N =  - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i))
    $$

    \State and update generator parameters by descending its stochastic gradient
    $$
    \theta \gets \theta - \alpha_G \nabla_\theta \mathcal{L}^{gen.}_N
    $$
\EndFor
\end{algorithmic}
\end{algorithm}
\fi

\subsection{Mode Collapse}
Mode collapse occurs when the generator's outputs focus on a few of the available modes, rather than replicating the full data-distribution. During training, once the discriminator learns that the generator is focused at a particular mode, it assigns low scores to the data-points coming from this mode. The response of the generator is then to shift its output distribution to another mode.
Mode collapse is therefore characterized by the generator's distribution switching from mode to mode throughout training.

\section{Generator Particles and Universality}
During training, the parameters of the generator, $\theta$, evolve by following their individual stochastic gradients (Eq.~\ref{eq:generator_update}).
However, rather than concerning ourselves with the dynamics of these parameters, here we study the dynamics of generator outputs treated as particles in data-space. While generator parameters follow their own local (stochastic) gradients, the dynamics of generator outputs can be viewed as being explicitly correlated.

Taking the parameters to be time dependent, $\theta_t$, we can denote a point in data-space which is mapped to by the seed $z$ at time $t$ by
$$
X_t = G_{\theta_t}(z).
$$
We then see that updates in parameter space, $d\theta_t$, are related to updates in data-space, $d X_t$, by
\begin{equation}\label{eq:basic_data_dynamics}
    d X_t = \frac{d G_{\theta} (z)}{d \theta}^T|_{\theta=\theta_t} \frac{d \theta_t}{d t} dt.
\end{equation}
Under gradient dynamics, the generator parameters evolve according to
\begin{equation}\label{eq:sim_gan_gen}
     \dot{\theta}_t = -\alpha_G \ \frac{d \mathcal{L}}{d \theta_t}=\alpha_G \ \frac{d}{d \theta_t}\langle D_\phi(G_\theta(z))\rangle_{z\sim q(z)},
\end{equation}
and so
\begin{align}
    \frac{d \theta_t}{d t} &= \frac{d}{d \theta} \left(\int dz' q(z')  D(G_{\theta} (z'))\right)|_{\theta=\theta_t} \\
    &= \int dz' q(z')  D'(G_{\theta_t} (z')) \frac{d G_{\theta} (z')}{d \theta}|_{\theta=\theta_t}.
\end{align}
To see the corresponding data-space dynamics, we plug this in to (\ref{eq:basic_data_dynamics}) and write
\begin{equation}
    dX_t^i = \alpha_G \ dt  \int dz' \ \Gamma^{i, j}_{\theta_t}(z, z') \frac{d D(G_{\theta_t} (z'))}{d x_j} q(z'), \label{eq:data_space_dynamics_1}
\end{equation}
where $i$ and $j$ represent the components of the vector $X_t$, and repeated indices are summed over. We have also introduced the \textit{neural tangent kernel} (NTK) \cite{jacot_ntk}, $\Gamma^{i, j}_{\theta}(z, z')$, defined by
\begin{align}
    &\Gamma_\theta: \mathbb{R}^{n} \times \mathbb{R}^{n} \rightarrow \mathbb{R}^{d \times d},\\
    &\Gamma_\theta^{i, j}(z, z') = \sum_k \frac{d G^i (z)}{d \theta_k}\frac{d G^j (z')}{d \theta_k},
\end{align}
where $n$ and $d$ are the dimensions of the inputs and the data-space, respectively, and $\theta_k$ are the parameters of the generator. {\color{blue} In general, avoid using ``Above", use ``where" and add a ',' at the end of the preceding equation. I have changed some, you can do the rest.}

By examining Equation \ref{eq:data_space_dynamics_1}, we see that here the NTK, $\Gamma_{\theta}$, tells us the degree to which the gradients of the \iffalse loss \fi discriminator at other points $X'=G_\theta(z')$ influence the dynamics of the the point $X=G_\theta(z)$.

In general, NTKs evolve during training. However, as we increase the width of the network, the degree to which the NTK changes throughout training drops \cite{jacot_ntk, deep_learning_theory_book}. In the full infinite-width limit, the NTK becomes asymptotically constant during training; its value fixed at initialization: $\Gamma_{\theta_t} = \Gamma_{\theta_0}$.

The infinite-width regime is of particular interest, as the performance of neural networks has been observed to improve as their width is increased. Additionally, in this limit it becomes possible to derive analytical results, as certain theoretical aspects of neural networks simplify \cite{ deep_learning_theory_book, Hanin_Finite, Halverson_2021, jacot_ntk, wide_networks_are_linear}.

For particular architectures, the exact form of the infinite-width NTK can be worked out analytically. Such an exact NTK can be found, for instance, in the case of an infinite width network with a ReLU or Erf activation  \cite{wide_networks_are_linear}{\color{blue} references where this is done?}.

\subsection{Mapping to model GANs}

In generative adversarial networks, random seeds are provided to the generator by sampling from a so-called noise distribution, $q(z)$, at each iteration. Usually, this is taken to be a high-dimensional Gaussian. Noting that points from $\mathcal{N}^d(0, 1)$ are approximately on an $d$-dimensional sphere of radius $\sqrt{d}$ \footnote{
Note that a d-dimensional vector with elements selected from $\mathcal{N}^d(0, \sigma^2)$ will have an average squared length of $d \sigma^2$, and the relative standard deviation of this estimate will drop as $\sqrt{\frac{2}{d}}$.
},
we will take the noise distribution as being a uniform selection from a d-sphere.

For certain activations (most prominently, ReLU), if input seeds have a fixed magnitude, then the infinite-width NTK will be a function only of the angle between inputs: $\Gamma(z, z') = \Gamma(\varphi_{z, z'})$ \cite{wide_networks_are_linear, kernel_methods_in_deep_learning}.
Additionally, these samples selected from high-dimensional spheres will, with high probability, be nearly orthogonal \footnote{Two elements, $z$ and $z'$ uniformly selected from an $d$-dimensional sphere of radius $\sqrt{d}$ will have a dot product obeying
$$
 \cos(\varphi_{z, z'}) \sim \mathcal{N}(0, \sigma^2 = \frac{1}{d}).
 $$
}. Therefore, given such an NTK and a sample of inputs taken uniformly from a high-dimensional sphere, it becomes possible to estimate the distribution of NTK values within a mini-batch.

\begin{figure}[h]
\includegraphics[width=.5\textwidth]{ntk_histogram.pdf}

\caption[hist_caption]{
{\color{blue} This (``Above") is not the way figure caption should be written.} The distribution of sampled NTK values using two distinct network architectures. 50 inputs, $\{ z_i \}$, are sampled from a unit sphere in 100 dimensions. We then compute the NTK values of pairs of inputs.
In red is shown a distribution obtained using a single hidden layer network with an an Erf activation, while in blue we use a ReLU activation.
 {\color{blue} What is the network architecture?} The ReLU and Erf networks have respective weights sampled from $\mathcal{N}(\mu=0, \sigma^2 \approx 0.42 \text{ and } 1.18)$, and respective biases sampled from $\mathcal{N}(\mu=0, \sigma^2 \approx 1.17 \text{ and } 11.67)$.

 Despite the differences between the two networks, the NTK values are approximately characterized by the same two numbers; corresponding to $\Gamma(\varphi_{z, z'} \approx \pi/2)$ for distinct inputs, and $\Gamma(\varphi_{z, z} = 0)$ for pairs of the same input. This type of universality allows our framework to describe many infinite-width networks simultaneously. {\color{red} Do you just use the initial weights? What happens when you evolve the weights and look at the NTK distribution again?} {\color{blue} This note speaks to a more global restructuring I need to do -- to emphasize the fact that infinite-width networks have an asymptotically static ntk.}
}
\label{fig:NTK_distribution}
\end{figure}

Using these observations, we propose a simplification of the GAN training protocol. Within our simplified model, we take our generator to be of large width. The noise distribution, $q(z)$, is taken to be a uniform distribution over a high-dimensional sphere.
{\color{red} Steve's comment: here, should we add a bit on exactly what regime this applies to? What does `large $d$' mean here? I believe that as a function of $d$ and $\Gamma$, we want the width of the $g_2$ values to be small: $\Gamma'(\pi/2) \sqrt{d}$. This could be something Youssef might be able to provide insight into.}
Finally, (as in wide ReLU networks with inputs of fixed-magnitude), we take the NTK to be a function of the dot-product of inputs only.

The effect of the first of these assumptions is to fix the NTK at initialization \cite{wide_networks_are_linear, jacot_ntk}. The final two assumptions have the result of concentrating the pairwise NTK values obtained using one  sample of intputs, $\{z \}$, to two characteristic numbers which we denote $g_1$ and $g_2$. The first of these, $g_1 \equiv \Gamma(\varphi_{z, z} = 0)$, corresponds to evaluations involving the same point, and the second, $g_2 \equiv \Gamma(\varphi_{z, z'} = \pi/2)$, corresponds to pairs of distinct points.

The fact that an entire generator neural network, with its activation functions and individual weight and bias distributions, can be to an extent characterized by just two numbers, implies a sort of universality. Many different generator neural network architectures may be mapped on to the same system, parameterized only by ($g_1$, $g_2$). This universality can be observed in Figure \ref{fig:NTK_distribution}, in which two distinct networks (one with ReLU activations, the other using Erf) are observed to have very similar pairwise NTK values across a sample from a 100-dimensional sphere.

Additionally, we consider a restricted version of GAN training.
Rather than performing re-sampling from the noise distribution, taking new mini-batches from $q(z)$ at each training iteration, we instead train by effectively using one fized set of $N$ generator inputs, $\{z\}$, from $q(z)$. This same sample of inputs is used at each iteration of the training algorithm.

Based on these simplifications, we propose a coarse-grained NTK of the form
\begin{equation}\label{eq:general_ntk}
    \Gamma^{i, j}(z, z') = \delta_{i, j}\left(g_1 \delta_{z, z'} + g_2 (1-\delta_{z, z'})\right).
\end{equation}
This NTK is static throughout training, corresponding to an infinite-width generator.
Its two constant values, for diagonal ($g_1$) and off-diagonal ($g_2$) entries, characterize the NTK values for pairs of identical and distinct points, respectively.

\iffalse [We emphasize that within true GANs, training is not performed using a only single batch.
Rather, we take our cloud of particles to be an estimate of the full generator output distribution, with each point, for instance, representing the outputs from a region of the latent space (for example, all points $z'$ satisfying $\varphi_{z, z'}<\phi^*$).
We leave the mapping of generator networks to our model GAN setup for a future work, and here assume that such a mapping may be performed to obtain such a coarse-grained NTK.] \fi

\iffalse
correspond to the NTK value for $\varphi_{z, z}= 0$ and for $\varphi_{z, z'}\approx \pi/2$, respectively.
\fi

The form of this NTK allows us to further simplify -- ignoring the latent space entirely, and instead explicitly correlating particles in data-space:
\begin{equation}\label{eq:general_ntk_2}
    \Gamma^{i, j}_{a, b} = \delta_{i, j}\left(g_1 \delta_{a, b} + g_2 (1-\delta_{a, b})\right).
\end{equation}
Here, $a$ and $b$ index particles, while $i$ and $j$ index components in data-space. Out front, $\delta_{i, j}$ can be understood as implying a lack of correlation between the gradients of output degrees of freedom of an infinite width neural network.
Note that distinct particles are assigned a fixed NTK value -- regardless of their spatial position within data-space.

Using such an NTK, we can model the dynamics of points in output-space by writing
\begin{equation}\label{eq:data_space_dynamics_2}
    \frac{d X^i_a}{dt} \propto \frac{1}{N} \sum^{d, N}_{j, b} \Gamma^{i, j}_{a, b} \nabla_{x_j} D(X_b).
\end{equation}
[The dynamics are somewhat reminiscent of flocking behavior -- in which the velocities are found through a spatial average. Here, however, the average is not over velocities, but discriminator gradients. Additionally, the average is taken over all particles, rather than over a local region.]

{
\color{red}
Not sure how to include this:
We're working in the limit of all pairwise points having the same value. The width of the NTK values should be proportional to

$$
 \Gamma'(\pi / 2) / \sqrt{d}
$$

If we judge the width of the pairwise NTK values by the largest deviation from the mea, then we use $\phi^{-1}$, is the inverse of the normal cumulative distribution. This gives us
$$
\phi^{-1}(1 - \frac{1}{N}) \Gamma'(\pi / 2) / \sqrt{d}
$$

Is this the quantity we'd like to say is small here.
}

\subsection{Multi-Modal Target}

As a case-study, we consider a canonical two-dimensional problem of training a GAN on a distribution of 8 Gaussians arranged in a circle of radius 2, each Gaussian having a standard deviation 0.02. Since each Gaussian can be naturally taken to represent a distinct mode, this data distribution is used throughout GAN literature as a toy data-set for observing mode collapse \cite{veegan, mode_reg, youssef_mmd_paper}. Mode collapse in this context would correspond to a generator whose outputs are focused on one, or a subset, of the eight Gaussians. During training, mode collapse would cause these outputs to oscillate between these Gaussian modes, without splitting to cover all eight.

The generator particles are taken to be 2000 parameterized points in the plane, initialized as a Gaussian distribution with $\sigma=0.5$, while the discriminator is a ReLU network with 4 hidden layers of width 512 \footnote{weights are initialized using a Glorot uniform distribution \cite{glorot}, and biases are initialized at zero}.
The discriminator is updated during training according to gradient ascent, while the generator points, $X_i$, are updated using gradient descent. The full training routine is described in Algorithm \ref{alg:gan_training_2}.

\iffalse
we now write its distribution as
\begin{equation}\label{eq:new_dist}
p_{X}(x) = \frac{1}{N}\sum^N_i \delta(x-X_i)
\end{equation}
 Such a `cloud of points' represents a proxy for the full distribution in data-space of a generator neural network\footnote{Meaning the pushforward of the distribution $q(z)$ through a generator $G$.}.
 Note that by working with the cloud itself, and applying the dynamics given by equation (\ref{eq:data_space_dynamics_1}), we can ignore the details of the generator neural network while still learning about its properties\footnote{
 A similar point-cloud model was previously applied in \cite{sobolev}, where it was studied using a diagonal NTK.}. *\textbf{is this redundant?}*
 \fi

 \iffalse
 Here, instead of writing the NTK as a function of imagined inputs from a latent space, $z$, we will now express $\Gamma$ as a matrix with indices corresponding to points $X_i$ in data-space. Following the notation of (\ref{eq:general_ntk}), this takes the form
\begin{equation}\label{eq:updated_ntk_form}
    \Gamma^{k, l}(X_i, X_j) = \delta_{k, l}(g_1 \delta_{i, j} + g_2(1-\delta_{i, j}))
\end{equation}
By plugging equations (\ref{eq:new_dist}), (\ref{eq:updated_ntk_form}) into the generator dynamics given by equation (\ref{eq:data_space_dynamics_1}), we find that $X_i$ evolve according to
\begin{equation}\label{eq:data_space_dynamics_2}
    \frac{d X_i}{dt} \propto \frac{1}{N} \sum^N_j \Gamma_{i, j} \nabla_x D(X_j)
\end{equation}
with the constant of proportionality dependent on the learning rate of the generator.
\fi

In Figures \ref{fig:no_ntk_2d}, \ref{fig:with_ntk_2d} {\color{blue} Figure numbers should be arranged in ascending order as they are first cited in the text}, we show time slices of the training progress. The generator particles are shown in white on a heat-map of the discriminator values.
We begin by running an experiment using a diagonal NTK ($g_2 / g_1 = 0$ \footnote{
Noting the dynamics described in equation (\ref{eq:data_space_dynamics_1}), we normalize the generator's dynamics by the particle number, setting $g_1=2000$, $g_2=0$ so that $\frac{1}{N}\sum^N_{j, k}\Gamma(X_j, X_k) D'(X_k) = D'(X_k)$.
}).

Because this NTK is diagonal, the generator particles each independently ascend the local gradient of the discriminator: $\dot{X}^i_a \propto \nabla_{x_i} D(X_a).$ Visually (Fig.~\ref{fig:no_ntk_2d}), this corresponds to each particle (in white) drifting up the color gradient (taking steps towards lighter regions). Meanwhile, the discriminator modifies its parameters to increase the difference between the expectation on the real data and the generator particles -- assigning higher values (brighter colors) to the eight points in black, and lower values (darker colors) to the particles in white.

We observe the result of this dynamic in Figure \ref{fig:no_ntk_2d}. Initially, the discriminator assigns low values to the cluster of particles. However, the initial cloud of particles rapidly splits apart, and the adversarial dynamic results in informative gradients being passed to the generator particles, which quickly converge to the full multi-modal distribution.

In a second experiment, we begin with the same initialization, but instead use an NTK satisfying $g_2 / g_1 = 1/5$. Due to the non-trivial off-diagonal of the NTK, the velocities of the particles are correlated. As is shown in Figure \ref{fig:with_ntk_2d}, particles (in white) no longer split apart. Instead, generator particles stay together, the entire cluster shifting from mode to mode indefinitely.
The discriminator repeatedly attempts to assign low values (darker colors) to the cluster's spatial region.

\iffalse
Whereas before (when $g_2/g_1=0$), this would result in particles splitting apart and ascending their local gradients, the particles' velocities are now overwhelmingly dominated by the mean discriminator gradient experienced by the ensemble. Therefore, the entire cloud of particles shifts away from the minimum, following the average experienced discriminator gradient. The result is that the cluster shifts from mode to mode, while the discriminator tries to `catch up', and place its minimum at the position of the cluster.
\fi

This behavior is a key signature of mode collapse, and suggests an understanding of the phenomenon of mode collapse through the lens of our model.
For the remainder of this work, we will identify the observed failure of convergence and switching between modes as mode collapse.
By modifying $g_2 / g_1$, we can probe the onset of this mode collapse, and investigate which training algorithms and discriminators lead to improved performance.

\begin{figure}[h]
    \centering
    \includegraphics[width=.5\textwidth]{convergence_version_2.pdf}
    \caption
    {The dynamics of the model GAN system, using a diagonal NTK, shown at different time-slices during training. The generator particles are shown in white, while the discriminator values are shown using a color-plot, with higher values shown in brighter colors. Because the NTK is diagonal, the particles' velocities are not explicitly correlated, and they ascend their local gradients (from darker to brighter colors), $\dot{X} \propto \nabla_x D(X)$. Meanwhile, the discriminator attempts to maximize the difference between its expectation on the data (black points) and generator particles (in white). Initially, the discriminator assigns a low value to the cluster's location. Consequently, the cluster rapidly splits apart (step 100), each point following the local discriminator gradient.
     The combined adversarial dynamics are seen to result in convergence to the eight modes.}
    \label{fig:no_ntk_2d}
\end{figure}

\begin{figure}[H]
    \centering
    \includegraphics[width=.5\textwidth]{mode_collapse.pdf}
    \caption{
    Model GAN dynamics depicted over time (as in Fig.~\ref{fig:no_ntk_2d}) in which the NTK satisfies  $g_2/g_1 = 1/5$.
    The discriminator seeks to assign higher values (brighter colors) to real data-points (the eight points in black), and low values (darker colors) to generated points (shown in white).
    As had occurred in Fig.~\ref{fig:no_ntk_2d}, the discriminator initially places its minimum at the position of the generator particle cluster.
    Now, however, due to the large explicit correlation in particle velocity, particles do not split apart. Instead, we see that the cluster shifts away from discriminator minima before individual particles can split.
    As a result, the cluster of generator particles switches from mode to mode, as the discriminator attempts to `catch up' -- a behavior indicative of mode collapse. }
    \label{fig:with_ntk_2d}
\end{figure}

\begin{algorithm}
\caption{The coarse-grained GAN training algorithm}\label{alg:gan_training_2}
\begin{algorithmic}
\For{iteration number}
    \For{$n_{disc.}$}
        \State $\bullet$ Sample $N$ data-points, $\{ x_i \}$, from the 8-Gaussian distribution.

        \State$\bullet$  Compute
        $$
        \mathcal{L}_N = \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) - \frac{1}{N} \sum_{i=1}^N D_\phi(X_i) - \frac{\lambda}{2}\sum_j \phi_j^2
        $$
        \State and update discriminator parameters by ascending its stochastic gradient
        $$
        \phi \gets \phi + \alpha_D \nabla_\phi \mathcal{L}_N
        $$
    \EndFor
    \State $\bullet$ update $X_l$ according to equation (\ref{eq:data_space_dynamics_1})
    $$
    X_l \gets X_l + \alpha_G \ \frac{1}{N} \sum^N_l \Gamma_{k, l} \nabla_x D(X_l)
    $$
\EndFor
\end{algorithmic}
\end{algorithm}

\section{Model GAN Experiments}
\label{sec:mode_experiments}

We have observed that the ratio $g_2/g_1$ within Equation \ref{eq:general_ntk_2} may be increased to induce mode collapse. Apart from the architecture of the discriminator network, the remaining adjustable parameters of Algorithm \ref{alg:gan_training_2} concern the relative learning rate of the discriminator. The parameters $\alpha_D$ and $\alpha_G$ control the step-size of the discriminator and generator, respectively, while $n_{disc.}$ tunes the number of discriminator steps taken for each generator step.

We will therefore vary these parameters to examine the relationship between $g_2 / g_1$ and the learning rate of the discriminator.
The learning rate of the discriminator can be varied in two ways: by modifying the rate $\alpha_D$, or by modifying the value of $n_{disc.}$ used in algorithm. Here, we show the result of modifying $n_{disc.}$, leaving $\alpha_D$ experiments (which produce similar results) to the appendix.

\iffalse
As we vary $g_2/g_1$, we must take care to control for the overall effect that the NTK has on total velocity.
For example, if all points were initialized at the same location, $X$, then their velocity would obey
$$
\frac{d X_i}{dt} = \alpha_G \frac{1}{N}\sum_j \Gamma(X_i, X_j) D'(X_j) = \alpha_G \frac{D'(X)}{N} (g_1 + g_2 (N-1))
$$
To control for this effect, we vary $g_2/g_1$ while insisting that $g_1 + g_2 (N-1) = \text{const.}$ We take $g_1 + g_2 (N-1)=N$, so that $g_2 / g_1 = 0 \implies g_2=0, \ g_1 = N$, and $g_2 / g_1 = 1 \implies g_2= g_1 = 1$.
\fi

To characterize whether, at a given time, generator particles have converged or have collapsed to a single mode, we define a metric based on the entropy of the distribution. Letting $P_i$ be the fraction of particles for which the $i^{th}$ mode is the nearest,
 we define the following:
\begin{equation}\label{eq:mode_collapse_metric}
\text{Mode Collapse Metric} = \log(8) + \sum_i P_i \log(P_i).
\end{equation}
Note that $P_i = 1/8$, $i=1,2,\cdots8$ would give total mode coverage, and have a value of $0$. On the other hand, $P_1 = 1$, $P_{i>1}=0$ would correspond to all generator points being nearest to a single mode (giving a mode collapse metric value of $\log 8$).

\subsection{GAN Setup}

To maximize the interpretability of our results, we employ a simpler discriminator: one with a wide single hidden layer (2048 units):
\begin{equation}
    D(x) = a_i \sigma(w^j_i x^j + b_i).
\end{equation}
We will experiment with both Tanh and ReLU activations (Details can be found in the appendix).

The target data distribution is again taken to be the 8 Gaussians used previously\iffalse with modes set a distance 2 from the center, standard deviation = 0.02\fi. We take the number of generator particles to be 200, which are initialized at $(\frac{1}{\sqrt{2}}, \frac{1}{\sqrt{2}})$ with a standard deviation of $.1$.

The loss function is then defined as
\begin{multline}
  \mathcal{L} = \langle D(x) \rangle_\text{target} - \langle D(x) \rangle_{\text{gen.}} \\
  - \frac{1}{\text{width}} \left(\sum_{i, j} (w_i^j)^2 + \sum_i a_i^2 + \sum_i b_i^2 \right),
\end{multline}
where $\langle D(x) \rangle_{\text{gen.}}$ is the expectation of the discriminator on the generator distribution ($\frac{1}{N}\sum_i D(X_i)$), while $ \langle D(x) \rangle_\text{target}$ is its expectation on the data distribution (the eight Gaussians). The remaining terms represent an $L_2$ regularizer on the discriminator weights, placing an overall restriction on the discriminator.

Following the form of (\ref{eq:data_space_dynamics_2}), the particle velocities are given by:
\begin{align}
    \frac{d X_i}{dt} &= \frac{\alpha_G}{N} \sum_J^N \Gamma_{i, j} D'(X_{j}) \\
    &= \alpha_G\left( \frac{g_1-g_2}{N}D'(X_{i}) + g_2 \langle D'(X)\rangle \right).
    \label{eq:competition_eq}
\end{align}
And the entire system is trained using Algorithm \ref{alg:gan_training_2}.

\subsection{Results and Interpretation}
\label{interpretation}

We run the model GAN training algorithm for each ($g_2 / g_1$, $n_{disc}$) pair considered.
After training for a fixed number of iterations {\color{blue} what do you mean by this? do you mean "for a fixed large number of iterations"?}, and computing the mode collapse metric (Eq.~\ref{eq:mode_collapse_metric}) for all pairs, we can observe the behavior of the transition from convergence to mode collapse.

As can be seen from Equation \ref{eq:competition_eq}, each particle experiences competition between the mean discriminator gradient over the ensemble, $\langle D'(X)\rangle$, and the contribution from its local gradient, $\left( \frac{g_1 - g_2}{N}\right) D'(X_i)$ {\color{red} This inverse $N$ dependence makes the argument rather weak because this ``splitting term" is arbitrarily small at a large $N$.} {\color{blue} This is true, and needs to be addressed in a careful way. The experiments run here in effect use a single batch of $N$ inputs, and then train on these. I need to think about how precisely to convey this -- what it means practically, and how it relates to batch size. This might be another matter Youssef might be able to help us think about}.

The effect of this competition on the presence of mode collapse can understood by considering a cluster of generator particles localized to a region.

During training, the discriminator seeks to maximize the difference between its expectation on the training data and on the generator distribution.
Suppose to this end, the discriminator has formed its minimum within the region of generator support{\color{blue} This is not completely true as the discriminator is maximizing the difference between its expected values for the generated data and the training sample data)} {\color{green} This is a great point, and an important correction to make}. If the term involving local gradients dominates $g_2 \langle D'(X) \rangle$, then the particles in this cluster `split apart', each following its own individual gradient, regardless of the location of the discriminator minimum within the region.

However, for sufficiently large $g_2$, the term, $g_2 \langle D'(X) \rangle$, may be dominant. Now the location of the discriminator's minimum within the region becomes important, as it may determine both the magnitude and direction of $\langle D'(X) \rangle$. As is depicted in Figure \ref{fig:slip_away}, if the discriminator obtains a minimum far from the center of a cluster, $|\langle D'(X) \rangle|$ may become non-negligible, and cause the entire group of generator particles `slip away' -- even those on the opposing slope. Meanwhile, for a minimum closer to the cluster's center, the mean gradient experienced becomes sufficiently small to allow for the particles to `split apart'.

\begin{figure}[h]
\includegraphics[width=.5\textwidth]{slip_away_diagram.pdf}

\caption{
A schematic depicting how increased discriminator precision may result in the splitting up of generator particles. The cluster of generator particles is depicted as a uniform shaded region with distribution $p_\theta(x)$, while the discriminator is shown as a simple function of the form $D_\phi(x) \equiv a |x-b|$. Arrows indicate the local velocities of particles. In (a), The imprecision of the discriminator leads to a sub-minimal expectation, $\langle D(x) \rangle$, and significant $\langle D'(X) \rangle$ value, causing all particles to slip to the left due to all-to-all coupling through the NTK. On the other hand, in (b), the discriminator has found the optimal minimum directly at the center of the distribution, killing off average gradients and allowing particles to ascend their local slopes, splitting apart.
}
\label{fig:slip_away}
\end{figure}

 In this way, we conclude that the discriminator's precision in minimizing its value over a cluster of particles is related to its ability to split apart the clusters. More spatially precise discriminators may yield smaller values of $|\langle D'(X) \rangle|$, allowing local gradients to dominate particle dynamics.

Using the ReLU discriminator, the mode collapse data shows a power-law boundary separating mode collapse from convergence.
This is shown in Figure \ref{fig:relu_data} on a log-log scale.
Using a Tanh discriminator, we see a boundary which appears exponential, as is depicted on a log-linear plot in Figure \ref{fig:tanh_data}. Examples of generator particle distributions sampled across the transition can be found in Appendix \ref{app:sample_dists}.

\begin{figure}[t]
\includegraphics[width=.5\textwidth]{relu_window_3_scatter_5000_time.png}

\caption{
Experimental results taken after 5000 training iterations using a ReLU discriminator. The x-axis gives the value of $g_2 / g_1$ for a given experiment, while the y-axis indicates the number of training steps the discriminator is given at each iteration. Data is time-averaged, representing the mean of results taken at $5000 \pm n\times 20$, with $n = 0, 1, 2, 3$. The color of each data-point reflects the magnitude of mode collapse for the experiment. Brighter points indicate mode collapse while darker points correspond to convergence. A line is fit to $(1/2) \log 8$, indicating a power-law boundary. We see that as $g_2 / g_1$ is increased, the discriminator requires more `time' (a greater $n_{disc}$ value) in order to shift the training from mode collapse (bright points) to convergence (darker points).}
\label{fig:relu_data}
\end{figure}

\begin{figure}[t]
\includegraphics[width=.5\textwidth]{tanh_window_3_scatter_2500_time.png}

\caption{
Experimental results taken after 2500 training iterations using a Tanh discriminator.
The ratio, $g_2 / g_1$ is given on the x-axis for each experiment, while the y-axis gives the number of discriminator steps allowed at each iteration.
As for the ReLU results, data is time-averaged, showing the mean of results taken at $2500 \pm n\times 20$, with $n = 0, 1, 2, 3$. The magnitude of the mode collapse metric is shown in for each experiment, with bright/dark points indicating mode collapse/convergence. A line is fit to the values $(1/2) \log 8$, depicting a generally exponential transition. As in Fig.~\ref{fig:relu_data}, for greater $g_2 / g_1$, more training time (a greater $n_{disc}$) is required in order for training to be shifted from mode collapse (bight markers) to convergence (darker markers).
}
\label{fig:tanh_data}
\end{figure}

These two behaviors, power-law and exponential, are highly suggestive, and match another feature of wide ReLU and Tanh networks
referred to as the frequency principle \cite{frequency_vs_training, convergence_vs_freq, linear_f_principle}. As networks learn, they tend to first learn lower-frequency functions, before including higher-frequency contributions. This has the practical effect of producing a rate at which a network can learn a feature of spatial frequency $k$: $\gamma(k)$.
 For wide ReLU networks, $\gamma(k)$
is expected to be power law, while within Tanh networks,
an exponential $\gamma(k)$ is predicted \cite{linear_f_principle}.
{\color{red} What is "a feature of spatial frequency $k$"? Do you mean a spatial feature of size $1/k$? or the feature needs to repeat itself at a spatial frequency $k$? If it is the former, I would rather use feature scale (size) than feature frequency.}

If we identify spatial features of size $1/k$ as having a dominant spatial frequency of $k$, then simple arguments suggest (Appx.~\ref{section:precision_scaling_argument}) {\color{blue} This reference shows up as ``B" in the pdf file.} that in order to split a cluster, the maximum allowable spatial error grows as $(g_1 - g_2)/g_2$ {\color{blue} ``what is "the scale of the spatial precision"? do you mean ``the maximum spatial error"?}{\color{blue}I should change the way i talk about this to really make it clear what I mean by the spatial scale... This is related to the previous restructuring I noted I should do}.
This implies that to break apart such a distribution, we require the discriminator to learn a feature with spatial \textit{frequency} proportional to $g_2 / (g_1 - g_2)$ -- for small $g_2$, roughly, $k \sim g_2 / g_1$.

Assuming the discriminator has a frequency learning rate of $\gamma(k)$, the time required by the discriminator to learn such a feature scales as $T \sim \gamma(k)^{-1}$.
Plugging in an exponential learning rate, and using (Eq.~\ref{eq:freq}) gives us
$$
T \sim \exp(\alpha g_2 / g_1),
$$
suggesting that the number of discriminator steps ($n_{disc.}$) necessary to split this distribution should scale like
$$
n_{disc.} \sim \exp(\alpha g_2 / g_1),
$$
for some constant $\alpha$, matching the relationship observed in the figures. For ReLU, with a power law $\gamma(k)$, a similar result is found.

\section{Critical Regularization }{\color{red} why is this regularization ``critical"?}{\color{blue} In titling sections, I tried to choose names which would be salient to a physics audience. The analogy I'd like to draw in our regularization section is to critical damping, so I mentioned this in the title. This indeed could be a topic that Youssef might be able to speak to.}

Various regularization techniques have been applied to the problem of avoiding mode collapse \cite{veegan, mode_reg, youssef_mmd_paper}.
Here, we demonstrate that even in our simplified model GAN context, the effect of regularization on reducing mode collapse can be observed and probed to obtain a better understanding.

In \cite{youssef_mmd_paper}, a regularizer is defined which penalizes the kinetic energy of generator parameters.
The regularizer is defined by the term,
\begin{equation}
    \beta ||\nabla_\theta \langle D(G_\theta (z))\rangle ||^2 /2,
    \label{eq:reg1}
\end{equation}
and is introduced into the discriminator's loss function during training.
Since the velocities of the generator parameters ($\theta_k$) are given by the local gradients of the discriminator, this term controls their kinetic energy and
 effectively disincentivizes sharp gradients from being provided to the generator parameters; ideally encouraging the generator to take smoother paths to the target distribution. The effect on the GAN system can be seen to be similar to that of a damping term in physics, a connection made explicit in Appendix \ref{app:dirac_gan_reg}. Using this analogy as a conceptual starting-point, we can consider sweeping over $\beta$ to identify a region of `critical regularization'.

Despite the fact that our model GAN setup does without any reference to generator parameters (or, for that matter, a generator network), we may still incorporate such a term into training by expanding and recalling the definition of the NTK:
\begin{multline*}
    ||\nabla_\theta \langle D(G_\theta (z))\rangle ||^2 = \\
    \sum_{z, \ z' \sim q(z)} D'(G_\theta (z)) \nabla_\theta G_\theta (z) \nabla_\theta G_\theta (z') D'(G_\theta (z')) \\
    = \sum_{z, \ z' \sim q(z)} \nabla_{i} D(G(z)) \Gamma_\theta^{i, j}(z, z') \nabla_{j} D(G(z')).
\end{multline*}
The new form allows us to immediately apply the regularizer within the model GAN by including the term:
\begin{equation}
    \frac{\beta}{N^2} \sum_{a, b} \nabla_{i} D(X_a) \Gamma^{i, j}_{a, b} \nabla_{j} D(X_b).
    \label{eq:reg}
\end{equation}

{\color{blue}{
In this way, using only coupled particles in data-space, our model-GAN setup can incorporate regularizers on the generator parameters, whose very existence in the model is only implicit.
} \textit{improve phrasing here}}

{\color{red} What is the physical meaning of this regularization term shown above? Does it encourage the gradient vectors of the discriminator function (D) for different generator particles to anti-align? and how would such a regularization term help prevent mode collapsing in the generator particle setting?}

Inserting this into the same setup used to produce Figure \ref{fig:with_ntk_2d}, and repeating the identical procedure \footnote{Here, we set $\beta = 100/(\text{num. weights})$}, we now observe convergence (Fig. \ref{fig:reg_dyn_scatters}).

\begin{figure}[H]
    \centering
    \includegraphics[width=.5\textwidth]{regularized.pdf}
    \caption{Model GAN dynamics shown over time, taking $g_2/g_1 = 1/5$ and including the addition a regularizer (Eq.~\ref{eq:reg}). Although previously (as in Fig.~\ref{fig:with_ntk_2d}) the system oscillated from mode to mode, particles now converge steadily to the eight modes.}
    \label{fig:reg_dyn_scatters}
\end{figure}

Within our setup, we can experiment with the regularization parameter $\beta$. Running model GAN experiments using different $\beta$ values, we note regions which are under/over-regularized, and a critically regularized region (Fig. \ref{fig:reg_dyn}).

\begin{figure}[H]
    \centering
    \includegraphics[width=.5\textwidth]{relu_regularization_plot.pdf}
    \caption{
    The degree of mode collapse present within a model GAN when trained using a regularizer of strength $\beta$.
    In the upper plot, each point resents an experiment using a single-hidden layer ReLU discriminator, halted after 1600 iterations (similar to the experiments of Section \ref{sec:mode_experiments}). 200 points are used, with $g_2 / g_1 = 0.275$.
    The regularizer, which performs a role similar to that of a damping term, is seen to result in a reduction in mode collapse. We observe three regimes analogous to under, critically, and over-damped dynamics.
    Below, sampled generator particle configurations are shown for each of the three circled points within each regime. We observe that only the critically regularized example has converged.
    }
    \label{fig:reg_dyn}
\end{figure}

\section{Discussion}
Fill in / combine once body of paper is solid

\begin{enumerate}
    \item We have used an effective model to describe the dynamics of generator outoputs, without using a generator
    \item We have used this model to investigate mode collapse transition boundaries.

    \item (to my knowledge) the first connection between the f-principle and GAN training

    \item We have also shown that the model can represent the effect of regularizers.

    \item Discuss the importance of NTK's in neural networks. If NTK's were truly diagonal, neural networks would not train, as all points would need to be trained individually.

    \item We would like the particles we use to represent the distribution of the generator. [to be an unbiased estimator of the pushforward distribution] This means that the $g_1$ and $g_2$ values must be scaled as $N$ is varied.

\end{enumerate}

{\color{red}

\section{General comments -- YT}

\begin{enumerate}
    \item We need to justify more clearly the use of NTK to replace the generator.
    \item We need to justify the simplification of using two values $g_1$ and $g_2$ to approximate the NTK.
    \item We need to make the figures and figure captions more coherent and conform them to the format according to normal standard.
    \item We need to describe a bit more about the dynamics of the discriminator.
    \item We need to explain the "frequency" learning argument a bit more precisely.
    \item We need to provide an intuitive explanation and interpretation of how the "critical" regularizer works especially in the context of the NTK-based GAN used in our paper. Youseff's insights are needed here.
\end{enumerate}
}

\newpage

\section{\label{sec:level1}Introduction}
In the past decade, deep generative models have proven to be an impressive tool for sampling from complex distributions. In particular, generative adversarial networks (GANs) have been used to produce realistic data, and represent a powerful framework for training generative models \cite{gan_review, style_based_gan, scene_generation, face_generation}. Consequently, understanding and improving the training of GANs is of considerable interest.

GANs comprise two neural networks: one called the \textit{generator}, $G_\theta$, and the other called the \textit{discriminator}, $D_\phi$ (parameterized by $\theta$ and $\phi$, respectively).
\begin{align}
    \text{Generator }& G_\theta: \mathbb{R}^n \rightarrow \mathbb{R}^d \\
    \text{Discriminator }& D_\phi: \mathbb{R}^d \rightarrow \mathbb{R}
\end{align}
The generator is a function which maps randomly selected points in the latent space to points in data-space.
The discriminator assigns scores to these simulated data-points, as well as to genuine samples from the data-set. During training, the discriminator's goal is to distinguish real data from simulated data (through high and low scores, respectively), while the generator's goal is to increase the score assigned to its outputs by the discriminator \cite{gan_paper, deep_learning_book, wgan, mmd_gan}.

Although GANs are both powerful and popular, they are notoriously hard to train.
The adversarial nature of the dynamics distinguishes a GAN's objective, $\mathcal{L}(\phi, \theta)$, from a standard loss function -- one that is bounded from below and which the training algorithm seeks to minimize. Rather than living at the minimum, the ideal parameter settings here are at the saddle points of the loss landscape \cite{gan_paper}:
\begin{equation}
    \theta^* = \argmin_\theta \max_\phi \mathcal{L}(\phi, \theta).
\end{equation}
Convergence to such an equilibrium is difficult to attain, as it requires a careful balancing of the two competing networks during training.

One important form of non-convergence commonly encountered during GAN training is known as \textit{mode collapse} \cite{veegan, mode_reg}. Mode collapse occurs when samples from the generator fail to capture the full diversity of modes present in the data-set. Instead, the generator's output ``collapses" as it only produces samples from relatively few of the available modes in the data distribution.

When mode collapse occurs, during training the generator will focus its distribution on a small subset of the overall data-distribution. Eventually, the discriminator learns to identify the concentrated output of the generator, at which point the generator will switch from its current specialization to another \cite{veegan, mode_reg}. The generator's output switching from mode to mode, rather than converging on the distribution as a whole, is a key symptom of mode collapse.

Many practically useful training techniques for avoiding mode collapse have been proposed, often involving modified objective functions and novel regularizers \cite{veegan, wgan, gan_paper, gan_review}.
Here, rather than constructing empirical methods for reducing mode collapse,
we seek to understand this phenomenon from the perspective of dynamical systems, determine the physical meaning of competing factors, and derive principles to guide the training of GANs.

\iffalse
We motivate an abstracted model for GANs in which the generator is replaced by $N$ particles in $\mathbb{R}^d$, echoing the simplification applied within Sobolev Descent \cite{sobolev}. Here, however, we additionally incorporate a static neural tangent kernel (NTK) -- causing the particles to correspond to the outputs of an infinite width generator neural network.
\fi

The dynamics of learning in neural networks have been studied in weight space \cite{feng2020neural,feng2021phases}. Here, we map GANs to an effective model in which the output of the generator network is replaced by $N$ particles in $\mathbb{R}^d$. The learning dynamics in GANs can then be studied in the output space by following the motion of the $N$ ``output" particles, which descend the loss landscape set by both the discriminator's score function and the collective state of $N$  particles. We additionally incorporate a static neural tangent kernel (NTK) -- a \iffalse key \fi feature of realistic GANs using an infinite-width generator. Within our effective model, the NTK induces a dependence of the velocity of any particle on the discriminator gradient at the location of all particles.
As a result of the sampling procedure of generators and the form of common NTKs,
we show the presence of universality within a restricted set of neural network architectures;
many different types of infinite-width generator neural networks may lead to the same particle dynamics.
\iffalse
\textcolor{red}{As a result of the sampling procedure of generators and the form of common NTKs, we show the presence of universality -- many different types of infinite-width generator neural networks map to the same particle dynamics.}
\fi

We argue that this effective model provides a simplified and interpretable framework in which to understand mode collapse.
In particular, applying this model to a low-dimensional target distribution, we show a transition from convergence to mode collapse as a function of the NTK and the relative training time. We provide a physical interpretation which explains this transition in terms of learning characteristics of the discriminator.

Finally, we use this model to study GAN regularization --  modification of the training objective in order to promote convergence. We find that when a gradient regularizer \cite{youssef_mmd_paper} is introduced, it results in a reduction of mode collapse in our model GAN.
Additionally, by sweeping over regularization strengths, we are able to observe under-regularized, over-regularized, and critically-regularized regimes.
These regimes can be understood by analogy to the physics of a damped oscillator and its under, over, and critically damped cases.
The regularizer, which incentivizes a `smoother' generator, here plays the role of a damping term.

\iffalse
This further encourages the use of our model as a realistic proxy for true GANs.
\fi

\begin{figure}[H]
    \centering
    \includegraphics[width=.45\textwidth]{dynamics_diagram.pdf}
    \caption{
    \textbf{Mapping to an effective GAN model.}
    An illustration of how the input vectors $z$ in the seed space $\mathbf{Z}$ (sampled from a high-dimensional sphere) map to particles in the data space $\mathbf{X}$.
    During GAN training, generator parameters, $\theta_t$, evolve over time in the $\mathbf{\Theta}$ space (upper left, blue trajectory). As a result, given fixed inputs $\{ z \}$ (lower left), the set of data-space outputs, $\{X_t\}$, also evolves in time (lower right). It is the dynamics of these points in data space that our effective model directly describes.}
    \label{fig:dynamics}
\end{figure}

\section{Training and Failure}

In GANs, the generator is a neural network, $G_{\theta}$, which is fed random inputs, $z \in \mathbb{R}^n$, selected from some noise distribution $q(z)$.
 The generator outputs, $G_{\theta}(z) \in \mathbb{R}^d$, therefore represent samples from its implicit probability distribution in data-space, $p_{\theta}(X)$.
 Conceptually, the generator and discriminator seek to minimize and maximize an objective expressing the expected difference between the data-set and the generator's outputs:
\begin{equation}
    \mathcal{L}(\phi, \theta) = \langle D_\phi(x) \rangle_{x\sim p(x)} - \langle D_\phi(G_\theta(z)) \rangle_{z \sim q(z)}.
\end{equation}
The function $p(x)$ is the probability distribution of samples in the data-set, while $q(z)$ is the distribution from which seeds in the latent space are sampled.

In practice, however, the discriminator's objective is often modified to include regularization, restricting the magnitude of the discriminator network and promoting stability \cite{wgan, general_gan}.
Different GAN implementations exist, many with distinct objectives \cite{general_gan}.
Here we consider objective functions of the following form \cite{orig_mmd_gan, mmd_gan, wgan, general_gan} that characterize the discriminatory power under constraints:
\begin{align}\label{eq:loss_full}
    \mathcal{L}_D &\equiv \langle D_\phi(G_\theta(z)) \rangle_{z \sim q(z)} - \langle D_\phi(x) \rangle_{x\sim p(x)} + \lambda R(D_\phi, G_\theta),\\
    \mathcal{L}_G &\equiv \langle D_\phi(x) \rangle_{x\sim p(x)} - \langle D_\phi(G_\theta(z)) \rangle_{z \sim q(z)}.
\end{align}
$\mathcal{L_D}$ and $\mathcal{L_G}$ define the objectives for the discriminator and generator, respectively, where $R(D_\phi, G_\theta)$ represents a regularizer on the discriminator, limiting its magnitude under a norm of interest (here, we use an $L_2$-norm on the discriminator weights, $\phi$) with $\lambda \geq 0$ denoting the strength of the regularizer.

\iffalse Conceptually, \fi The discriminator parameters, $\phi$, evolve to maximize the expected difference between the discriminator's value on the real data and the generated data (Eq.~\ref{eq:loss_full}), while the generator parameters, $\theta$, evolve to minimize this difference.
\begin{equation}
    \dot{\phi} = -\alpha_D \frac{d \mathcal{L}_D}{d \phi}, \quad
    \dot{\theta} = -\alpha_G \frac{d \mathcal{L}_G}{d \theta}.
\end{equation}
The discriminator and generator evolution occurs at individual learning rates $\alpha_D$ and $\alpha_G$.

Practically, in neural networks, the loss function is defined using mini-batches of $N$ samples of real data and generated data, both of which are re-sampled at each training step:
\begin{align}
    \mathcal{L}_D^{(N)} &\equiv \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i)) - \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) + \lambda R(D_\phi, G_{\theta}), \\
    \mathcal{L}_G^{(N)} &\equiv \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i))
\end{align}
Training is performed in iterations. First, for $n_{disc.}$ steps, the discriminator is updated according to its stochastic gradient:
\begin{equation}
    \phi \xleftarrow[]{} \phi - \alpha_D \nabla_\phi \mathcal{L}_D^{(N)}.
\end{equation}
Then, for a single step, the generator is updated analogously with stochastic gradient descent:
\begin{equation}
    \theta \xleftarrow[]{} \theta - \alpha_G \nabla_\theta \mathcal{L}_G^{(N)}.
    \label{eq:generator_update}
\end{equation}
Alternating updates are repeated until convergence, or until training is halted after a large number of iterations.

\iffalse
\begin{algorithm}
\caption{The GAN training algorithm}\label{alg:gan_training}
\begin{algorithmic}
\For{iteration number}
    \For{$n_{disc.}$}
        \State $\bullet$ Sample $N$ data-points, $\{ x_i \}$
        \State $\bullet$ Sample $N$ seed points, $\{ z_i \}$ from the latent space according to $q(z)$

        \State$\bullet$  Compute
        $$
        \mathcal{L}_N = \frac{1}{N} \sum_{i=1}^N D_\phi(x_i) - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i)) - \lambda R(D_\phi)
        $$
        \State and update discriminator parameters by ascending its stochastic gradient
        $$
        \phi \gets \phi + \alpha_D \nabla_\phi \mathcal{L}_N
        $$
    \EndFor
    \State $\bullet$ Sample $N$ seed points, $\{ z_i \}$ from the latent space according to $q(z)$

    \State $\bullet$ Compute
    $$
    \mathcal{L}^{gen.}_N =  - \frac{1}{N} \sum_{i=1}^N D_\phi(G_\theta(z_i))
    $$

    \State and update generator parameters by descending its stochastic gradient
    $$
    \theta \gets \theta - \alpha_G \nabla_\theta \mathcal{L}^{gen.}_N
    $$
\EndFor
\end{algorithmic}
\end{algorithm}
\fi

Mode collapse occurs when the generator's outputs focus on a few of the available modes, rather than replicating the full data-distribution. During training, once the discriminator learns that the generator is focused at a particular mode, it assigns low scores to the data-points coming from this mode. The response of the generator is then to shift its output distribution to another mode.
Mode collapse is therefore characterized by the generator's distribution switching from mode to mode throughout training.

\section{Generator Particles and Universality}
Rather than following the dynamics of generator parameters (Eq.~\ref{eq:generator_update}), we study instead the time evolution of generator outputs treated as particles in data-space (Fig.~\ref{fig:dynamics}), an approach applied in \cite{sobolev} and later in \cite{ntk_perspective_gans}. While each generator parameter follows its own local (stochastic) gradient, as we will show below, the dynamics of generator outputs are explicitly correlated.

With a time-dependent vector of parameters, $\theta_t$, a fixed seed $z$ maps to a point in data-space at time $t$ according to
$$
X_t = G_{\theta_t}(z).
$$
This mapping relates updates in parameter space, $d\theta_t$, to updates in data-space, $d X_t$, by
\begin{equation}\label{eq:basic_data_dynamics}
    d X_t = \frac{d G_{\theta} (z)}{d \theta}^T|_{\theta=\theta_t} \frac{d \theta_t}{d t} dt.
\end{equation}
Under gradient dynamics, the generator parameters evolve according to
\begin{equation}\label{eq:sim_gan_gen}
     \dot{\theta}_t = -\alpha_G \ \frac{d \mathcal{L}_G}{d \theta_t}=\alpha_G \ \frac{d}{d \theta_t}\langle D_\phi(G_\theta(z))\rangle_{z\sim q(z)},
\end{equation}
and so
\begin{align}
    \frac{d \theta_t}{d t} &= \alpha_G \frac{d}{d \theta} \left(\int dz' q(z')  D(G_{\theta} (z'))\right)|_{\theta=\theta_t} \\
    &= \alpha_G \int dz' q(z')  \nabla_j D(G_{\theta_t} (z')) \frac{\partial G^j_{\theta} (z')}{\partial \theta}|_{\theta=\theta_t}.
\end{align}
To see the corresponding data-space dynamics, we plug this into Eq.~\ref{eq:basic_data_dynamics} and write
\begin{equation}
    dX_t^i = \alpha_G \ dt  \int dz' \ \Gamma^{i, j}_{\theta_t}(z, z') \nabla_j D(G_{\theta_t} (z')) q(z'), \label{eq:data_space_dynamics_1}
\end{equation}
where $i$ and $j$ index the components of the data vector $X_t$ and repeated indices are summed over.

Moreover, we have introduced the \textit{neural tangent kernel} (NTK) \cite{jacot_ntk}, $\Gamma^{i, j}_{\theta}(z, z')$, defined by
\begin{align}
    &\Gamma_\theta: \mathbb{R}^{n} \times \mathbb{R}^{n} \rightarrow \mathbb{R}^{d \times d},\\
    &\Gamma_\theta^{i, j}(z, z') = \sum_k \frac{\partial G^i (z)}{\partial \theta_k}\frac{\partial G^j (z')}{\partial \theta_k},
\end{align}
where $n$ and $d$ are the dimensions of the inputs and the data-space, respectively, and $\theta_k$ denotes the $k$th generator parameter.
Importantly, Eq.~\ref{eq:data_space_dynamics_1} makes clear that the NTK, $\Gamma_{\theta}$, couples the generator outputs; it specifies to what extent the dynamics of the generator particle at $X=G_\theta(z)$ is influenced by the discriminator gradients at the position $X'=G_\theta(z')$ of all other particles.
\iffalse
{\color{blue}
 More broadly, a network's NTK may be used to characterize its convergence during training as well as its inductive bias and generalization properties\cite{jacot_ntk}.
}\fi

In general, NTKs evolve during training. However, for larger-width networks, the weights, $\theta_t$, will asymptotically remain in the vicinity of their initial values, $\theta_0$. The network's NTK, which involves a sum over the network's weights, changes even less -- in the infinite-width limit becoming fixed at initialization \cite{jacot_ntk, deep_learning_theory_book} (see Appendix~\ref{sec:ntk_evolution} for an example of such large-width training dynamics). In this work, we will assume that the generator is in this infinite-width regime, and enforce that the generator NTK remains fixed during training: $\Gamma_{\theta_t} = \Gamma_{\theta_0}$.

The infinite-width regime is of particular interest, as the performance of neural networks has been observed to improve as their width is increased. Additionally, in this limit, it becomes possible to derive analytical results, as certain theoretical aspects of neural networks simplify \cite{ deep_learning_theory_book, Hanin_Finite, Halverson_2021, jacot_ntk, wide_networks_are_linear}. The exact form of the infinite-width NTK can be found for particular network architectures, such as those with a ReLU or Erf activation \cite{wide_networks_are_linear}.

\subsection{Mapping to model GANs}

In generative adversarial networks, random seeds are provided to the generator by sampling from a so-called noise distribution, $q(z)$, at each iteration. Usually, this is taken to be a high-dimensional Gaussian. Noting that points from $\mathcal{N}^n(0, 1)$ are approximately on a sphere of radius $\sqrt{n}$ in $n$ dimensions \footnote{
Note that a vector selected from $\mathcal{N}^n(0, \sigma^2)$ will have an average squared length of $n \sigma^2$, and the relative standard deviation of this estimate will drop as $\sqrt{\frac{2}{n}}$.
},
we take the noise distribution as a uniform selection from a $(n-1)$-sphere.

For certain activation (most prominently, ReLU), if input seeds have a fixed magnitude, then the infinite-width NTK will be a function only of the angle between inputs: $\Gamma(z, z') = \Gamma(\varphi_{z, z'})$ \cite{wide_networks_are_linear, kernel_methods_in_deep_learning}.
Additionally, these samples selected uniformly from a high-dimensional sphere will, with high probability, be nearly orthogonal \footnote{Two elements, $z$ and $z'$ uniformly selected from an $(n-1)$-sphere of radius $\sqrt{n}$ will have a dot product obeying
$$
 \cos(\varphi_{z, z'}) \sim \mathcal{N}(0, \sigma^2 = \frac{1}{n}).
 $$
}. Therefore, given such an NTK and high-dimensional inputs, it becomes possible to estimate the distribution of NTK values within a mini-batch.

\begin{figure}[ht]
\includegraphics[width=.5\textwidth]{ntk_histogram.pdf}
\caption[hist_caption]{
\textbf{Universality of NTK values.}
Two distinct network architectures result in similar distributions of sampled NTK values.
50 inputs, $\{ z_i \}$, are sampled from a unit sphere in 100 dimensions. Using two untrained networks, the NTK values for all pairs of inputs are computed.
The red histogram is obtained using a single hidden layer network (width 2048) with an Erf activation, while for the blue histogram a ReLU activation is used. Zero-mean normal distributions with distinct variances are used to initialize two networks' weights and biases.
Despite their differences, the two networks' NTK values are approximately characterized by the same two numbers: $\Gamma(\varphi_{z, z'} \approx \pi/2)$ for distinct inputs and $\Gamma(\varphi_{z, z} = 0)$ for pairs of the same input.
}
\label{fig:NTK_distribution}
\end{figure}

Using these observations, we propose a simplification of the GAN training protocol. Within our simplified model, we take the generator to be of large width with a static NTK. The noise distribution, $q(z)$, is taken to be a uniform distribution over a high-dimensional sphere.
Finally, (as in wide ReLU networks with inputs of fixed magnitude), we take the NTK to be a function of the dot-product of inputs only.

Our first assumption fixes the NTK at initialization \cite{wide_networks_are_linear, jacot_ntk}. The latter two concentrate the pairwise NTK values, obtained using one sample of inputs $\{z \}$, to two characteristic numbers $g_1$ and $g_2$.
The first number, $g_1 \equiv \Gamma(\varphi_{z, z} = 0)$, corresponds to evaluations involving the same point, and the second, $g_2 \equiv \Gamma(\varphi_{z, z'} = \pi/2)$, corresponds to pairs of distinct points chosen from the high-dimensional latent space.
The values of $g_1$ and $g_2$ are determined by the architecture of the network, but can also be modified by, for instance, the use of batch-normalization \cite{freeze_chaos_ntk, batchnorm}.

The fact that an entire generator neural network, with its activation functions and individual weight and bias distributions, can be to an extent characterized by just two numbers, suggests a sort of universality within this particular set of neural network architectures. Many different generator neural network architectures may be mapped on to the same system, parameterized only by ($g_1$, $g_2$). This universality can be observed in Fig.~\ref{fig:NTK_distribution}, in which two distinct networks (one with ReLU activation, the other with Erf, both using a single hidden layer with 2048 units \footnote{
The ReLU and Erf networks have respective weights sampled from $\mathcal{N}(\mu=0, \sigma^2 \approx 0.42 \text{ and } 1.18)$, and respective biases sampled from $\mathcal{N}(\mu=0, \sigma^2 \approx 1.17 \text{ and } 11.67)$.}) are observed to have very similar pairwise NTK values across a sample from a unit sphere in 100 dimensions.

Additionally, we consider a restricted version of GAN training.
Rather than re-sampling from the noise distribution (i.e. taking a new mini-batch from $q(z)$) at each training iteration, we instead train by effectively using one fixed set of $N$ generator inputs, $\{z\}$.

Based on these simplifications, we propose a coarse-grained NTK of the form
\begin{equation}\label{eq:general_ntk}
    \Gamma^{i, j}(z, z') = \delta_{i, j}\left(g_1 \delta_{z, z'} + g_2 (1-\delta_{z, z'})\right).
\end{equation}
This NTK is static throughout training,
and its two constant values, for diagonal ($g_1$) and off-diagonal ($g_2$) entries, characterize the NTK values for pairs of identical and distinct points, respectively.

\iffalse [We emphasize that within true GANs, training is not performed using a only single batch.
Rather, we take our cloud of particles to be an estimate of the full generator output distribution, with each point, for instance, representing the outputs from a region of the latent space (for example, all points $z'$ satisfying $\varphi_{z, z'}<\phi^*$).
We leave the mapping of generator networks to our model GAN setup for a future work, and here assume that such a mapping may be performed to obtain such a coarse-grained NTK.] \fi

\iffalse
correspond to the NTK value for $\varphi_{z, z}= 0$ and for $\varphi_{z, z'}\approx \pi/2$, respectively.
\fi

This NTK allows us to further simplify the effective model -- ignoring the latent space entirely, and instead explicitly correlating particles in data-space:
\begin{equation}\label{eq:general_ntk_2}
    \Gamma^{i, j}_{a, b} = \delta_{i, j}\left(g_1 \delta_{a, b} + g_2 (1-\delta_{a, b})\right).
\end{equation}
Here, $a$ and $b$ index particles, while $i$ and $j$ index components in data-space. Out front, $\delta_{i, j}$ can be understood as implying a lack of correlation between the gradients of output degrees of freedom of an infinite-width neural network.
$g_1$ and $g_2$ set the degree to which the discriminator gradient at local and distinct points, respectively, contribute to a generator particle's velocity.

Using this effective NTK, we can model the dynamics of data points in output-space by dynamics of coupled generator particles
\begin{equation}\label{eq:data_space_dynamics_2}
    \frac{d X^i_a}{dt} = \frac{\alpha_G}{N} \sum^{N, d}_{b, j} \Gamma^{i, j}_{a, b} \nabla_j D(X_b).
\end{equation}
These dynamics are reminiscent of flocking behavior, in which local velocities are found through a spatial average \cite{flock}. Here, however, the average is not over velocities, but over discriminator gradients. Additionally, the average is taken over all particles, rather than over a local region.

\subsection{Multi-Modal Target}
We now proceed with our simplified GAN training protocol, replacing the generator network with a collection of $N$ particles in data-space, and using the generator update rule of Eq.~(\ref{eq:data_space_dynamics_2}) rather than that of Eq.~(\ref{eq:data_space_dynamics_1}).

As a case study, we consider a canonical two-dimensional problem of training a GAN on a distribution of 8 Gaussians arranged in a circle of radius 2, each having a standard deviation 0.02. Since each Gaussian can naturally represent a distinct mode, this data distribution is used throughout GAN literature as a toy data-set for observing mode collapse \cite{veegan, mode_reg, youssef_mmd_paper}. Mode collapse in this context would correspond to a generator whose outputs are focused on one, or a subset, of the eight Gaussians. During training, mode collapse would cause the outputs to oscillate between distinct modes, without splitting to cover all eight.

The generator particles are taken to be 2000 parameterized points in the plane, initialized as a Gaussian distribution with $\sigma=0.5$, while the discriminator is a ReLU network with 4 hidden layers of width 512 \footnote{weights are initialized using a Glorot uniform distribution \cite{glorot}, and biases are initialized at zero}.
The discriminator parameters, $\phi$, and the generator points, $X_a$, are both updated during training according to their objective functions, following training routine described in Algorithm \ref{alg:gan_training_2}.

\iffalse
we now write its distribution as
\begin{equation}\label{eq:new_dist}
p_{X}(x) = \frac{1}{N}\sum^N_i \delta(x-X_i)
\end{equation}
 Such a `cloud of points' represents a proxy for the full distribution in data-space of a generator neural network\footnote{Meaning the pushforward of the distribution $q(z)$ through a generator $G$.}.
 Note that by working with the cloud itself, and applying the dynamics given by Eq.~(\ref{eq:data_space_dynamics_1}), we can ignore the details of the generator neural network while still learning about its properties\footnote{
 A similar point-cloud model was previously applied in \cite{sobolev}, where it was studied using a diagonal NTK.}. *\textbf{is this redundant?}*
 \fi

 \iffalse
 Here, instead of writing the NTK as a function of imagined inputs from a latent space, $z$, we will now express $\Gamma$ as a matrix with indices corresponding to points $X_i$ in data-space. Following the notation of (\ref{eq:general_ntk}), this takes the form
\begin{equation}\label{eq:updated_ntk_form}
    \Gamma^{k, l}(X_i, X_j) = \delta_{k, l}(g_1 \delta_{i, j} + g_2(1-\delta_{i, j}))
\end{equation}
By plugging Eqs.~(\ref{eq:new_dist}), (\ref{eq:updated_ntk_form}) into the generator dynamics given by Eq.~(\ref{eq:data_space_dynamics_1}), we find that $X_i$ evolve according to
\begin{equation}\label{eq:data_space_dynamics_2}
    \frac{d X_i}{dt} \propto \frac{1}{N} \sum^N_j \Gamma_{i, j} \nabla_x D(X_j)
\end{equation}
with the constant of proportionality dependent on the learning rate of the generator.
\fi

In Figures \ref{fig:no_ntk_2d} and \ref{fig:with_ntk_2d},
 we show time slices of the training progress. The generator particles are shown in white on a heat-map of the discriminator values.
We begin by running an experiment using a diagonal NTK ($g_2 / g_1 = 0$ \footnote{
Noting the dynamics described in Eq.~(\ref{eq:data_space_dynamics_1}), we normalize the generator's dynamics by the particle number, setting $g_1=2000$, $g_2=0$ so that $\frac{1}{N}\sum^N_{a, b}\Gamma(X_a, X_b) \nabla D(X_a) = \nabla D(X_b)$.
}).
In this case, the generator particles independently ascend the local gradient of the discriminator: $\dot{X}^i_a \propto \nabla_i D(X_a).$ Visually (Fig.~\ref{fig:no_ntk_2d}), this corresponds to each particle (in white) drifting up the color gradient (taking steps towards lighter regions). Meanwhile, the discriminator modifies its parameters to increase the difference between the expectation on the real data and the generator particles -- assigning higher values (brighter colors) to the eight data points in black, and lower values (darker colors) to the particles in white.

We observe the result of this dynamic in Fig.~\ref{fig:no_ntk_2d}. Initially, the discriminator assigns low values to the cluster of particles. However, the initial cloud of particles rapidly splits apart, and the adversarial dynamic results in informative gradients being passed to the generator particles, which quickly converge to the full multi-modal distribution.

In a second experiment, we begin with the same initialization, but instead use an NTK satisfying $g_2 / g_1 = 1/5$. Due to the non-trivial off-diagonal terms of the NTK, the velocities of the particles are correlated. As is shown in Fig.~\ref{fig:with_ntk_2d}, generator particles (in white) no longer split apart. Instead, they stay together as the entire cluster shifts from mode to mode indefinitely.
The discriminator repeatedly attempts to assign low values (darker colors) to the generator particle cluster's spatial region.

\iffalse
Whereas before (when $g_2/g_1=0$), this would result in particles splitting apart and ascending their local gradients, the particles' velocities are now overwhelmingly dominated by the mean discriminator gradient experienced by the ensemble. Therefore, the entire cloud of particles shifts away from the minimum, following the average experienced discriminator gradient. The result is that the cluster shifts from mode to mode, while the discriminator tries to `catch up', and place its minimum at the position of the cluster.
\fi

This behavior is a key signature of mode collapse, and suggests an understanding of this phenomenon through the lens of our model.
For the remainder of this work, we will identify the observed failure of convergence and switching between modes with mode collapse.
By varying $g_2 / g_1$, we will probe the onset of this failure mode, and investigate what training algorithms and discriminator characteristics would lead to improved performance.

\begin{figure}[ht]
    \centering
    \includegraphics[width=.5\textwidth]{convergence_version_2.pdf}
    \caption
    {\textbf{Convergence of the model GAN dynamics under a diagonal NTK ($g_2 = 0$)}. The generator particles are shown in white, while the discriminator values are color-coded, with higher values shown in brighter colors. Because the NTK is diagonal, the particles' velocities are not explicitly correlated, and they ascend their local gradients (from darker to brighter colors), $\dot{X} \propto \nabla D(X)$. Meanwhile, the discriminator attempts to maximize the difference between its expectation on the data (black points) and generator particles (in white). Initially, the discriminator assigns a low value to the cluster's location. Consequently, the cluster rapidly splits apart (step 100), with each point following the local discriminator gradient.
     The combined adversarial dynamics are seen to result in convergence to the eight modes.}
    \label{fig:no_ntk_2d}
\end{figure}

\begin{figure}[ht]
    \centering
    \includegraphics[width=.5\textwidth]{mode_collapse.pdf}
    \caption{
    \textbf{Model GAN dynamics exhibit mode collapse under all-to-all coupling by a non-diagonal NTK ($g_2 / g_1 = 1/5$).}
    As in Fig.~\ref{fig:no_ntk_2d}, dynamics are depicted over time.
    During training, the discriminator seeks to assign higher values (brighter colors) to real data-points (the eight points in black), and lower values (darker colors) to generated points (white dots).
    As had occurred in Fig.~\ref{fig:no_ntk_2d}, the discriminator initially places its minimum at the position of the generator particle cluster.
    Now, however, due to correlations in particle velocity, the cluster no longer splits apart. Instead, it shifts away from discriminator minima before splitting can happen.
    As a result, the cluster of generator particles switches from mode to mode, as the discriminator attempts to `catch up' -- a behavior indicative of mode collapse. }
    \label{fig:with_ntk_2d}
\end{figure}

\begin{algorithm}
\begin{algorithmic}
\For{iteration number}
    \For{$n_{disc.}$}
        \State $\bullet$ Sample $N$ data-points, $\{ x_i \}$, from the 8-Gaussian distribution.

        \State$\bullet$  Compute
        $$
        \mathcal{L}_D^{(N)} =  \frac{1}{N} \sum_{a=1}^N D_\phi(X_a) - \frac{1}{N} \sum_{l=1}^N D_\phi(x_l) + \frac{\lambda}{2}\sum_k \phi_k^2
        $$
        \State and update discriminator parameters by descending its stochastic gradient
        $$
        \phi \gets \phi - \alpha_D \nabla_\phi \mathcal{L}_D^{(N)}
        $$
    \EndFor
    \State $\bullet$ update $X_a$ according to Eq.~(\ref{eq:data_space_dynamics_2})
    $$
    X_a^i \gets X_a^i + \alpha_G \ \frac{1}{N} \sum^{N, d}_{b, j} \Gamma_{a, b}^{i, j} \nabla_j D(X_b)
    $$
\EndFor
\end{algorithmic}
\caption{The coarse-grained, $(g_1, g_2)$ GAN training algorithm.}
\label{alg:gan_training_2}
\end{algorithm}

\section{Model GAN Experiments | \newline The Mode Collapse Transition}
\label{sec:mode_experiments}

We have observed that the ratio $g_2/g_1$ may be increased to induce mode collapse. Apart from the architecture of the discriminator network, the remaining adjustable parameters in Algorithm \ref{alg:gan_training_2} concern the relative training dynamics of the discriminator and generator. The parameters $\alpha_D$ and $\alpha_G$ control the step-size of the discriminator and generator, respectively, while $n_{disc.}$ tunes the number of discriminator steps taken for each generator step.

We will therefore vary these parameters to examine the relationship between $g_2 / g_1$ and the discriminator's dynamics.
The latter can be varied in two ways: by modifying the learning rate $\alpha_D$, or by modifying the value of $n_{disc.}$ used in the algorithm. Here, we show the result of modifying $n_{disc.}$, leaving $\alpha_D$ experiments (which produce similar results) to Appendix~\ref{app:exp_details}.

\iffalse
As we vary $g_2/g_1$, we must take care to control for the overall effect that the NTK has on total velocity.
For example, if all points were initialized at the same location, $X$, then their velocity would obey
$$
\frac{d X_i}{dt} = \alpha_G \frac{1}{N}\sum_j \Gamma(X_i, X_j) D'(X_j) = \alpha_G \frac{D'(X)}{N} (g_1 + g_2 (N-1))
$$
To control for this effect, we vary $g_2/g_1$ while insisting that $g_1 + g_2 (N-1) = \text{const.}$ We take $g_1 + g_2 (N-1)=N$, so that $g_2 / g_1 = 0 \implies g_2=0, \ g_1 = N$, and $g_2 / g_1 = 1 \implies g_2= g_1 = 1$.
\fi

To characterize whether, at a given time, generator particles have converged or collapsed to a single mode, we define a metric based on the entropy of the distribution. Letting $P_i$ be the fraction of particles for which the $i^{th}$ mode is the nearest,
 we define the following:
\begin{equation}\label{eq:mode_collapse_metric}
\text{Mode Collapse Metric} = \log(8) + \sum_i P_i \log(P_i).
\end{equation}
Note that $P_i = 1/8$, $i=1,2,\cdots8$, would give a complete mode coverage with an even split, and have a value of $0$. On the other hand, $P_1 = 1$, $P_{i>1}=0$ would correspond to all generator points being nearest to a single mode, giving a mode collapse metric value of $\log 8$.

To further characterize the quality of convergence of the generator particles, we can compute the average log-likelihood, $\frac{1}{N}\sum_a \log(p(X_a))$, where $p(X)$ is the probability density of the multi-modal Gaussian distribution.
The combination of these two metrics (mode collapse and log-likelihood) indicates whether the generator points have both avoided mode collapse and successfully converged to the modes of the distribution.

\iffalse
Further, we compute the average Euclidean distance to the closest mode across each of the generator points. to quantify their convergence to individual modes.
The combination of these two metrics indicate whether the points have both overcome mode collapse, and converged to the distribution fully.
\fi

\iffalse
{\color{red}(wouldn't average distance to nearest mode be better? I think it might tell the story we're trying to convey more clearly. Below the threshold its low, above the threshold its low. Below the threshold it's close to the modes because of mode collapse, above the threshold its close to the modes because of convergence)}
\fi

\subsection{GAN Setup}

To maximize the interpretability of our results, we employ a simpler discriminator with a single wide hidden layer (2048 units):
\begin{equation}
    D(x) = \sqrt{\frac{2}{\text{width}}} a_i \sigma(w^j_i x^j + b_i).
\end{equation}
 Details of the initialization can be found in Appendix~\ref{app:exp_details}.
The activation function is set to ReLU, $\sigma(x) = \max(0, x)$. Experiments employing a Tanh activation were also performed and the results can be found in Appendix~\ref{app:fprince_support}.
The target data distribution is again taken to be the eight Gaussians.
A total of 200 generator particles are initialized at $(\frac{1}{\sqrt{2}}, \frac{1}{\sqrt{2}})$ with a standard deviation of $0.1$.

The discriminator loss function is defined as
\begin{multline}
  \mathcal{L}_D = \langle D(X) \rangle_{\text{gen.}} - \langle D(x) \rangle_\text{target}\\
  + \frac{1}{\text{width}} \left(\sum_{i, j} (w_i^j)^2 + \sum_i a_i^2 + \sum_i b_i^2 \right),
\end{multline}
where $\langle D(X) \rangle_{\text{gen.}}$ is the expectation of the discriminator on the generator distribution, $\frac{1}{N}\sum_a D(X_a)$, while $ \langle D(x) \rangle_\text{target}$ is its expectation on the data distribution (the eight Gaussians). The remaining terms represent an $L_2$-regularizer on the weights, placing an overall restriction on the discriminator.

Following Eq.~\ref{eq:data_space_dynamics_2}, particle velocities are given by
\begin{align}
    \frac{d X^i_a}{dt} &= \frac{\alpha_G}{N} \sum_{b, j}^{N, d} \Gamma^{i, j}_{a, b} \nabla_j D(X_{b}) \\
    &= \alpha_G\left( \frac{g_1-g_2}{N} \nabla_i D(X_{a}) + g_2 \langle \nabla_i D(X)\rangle \right).
    \label{eq:competition_eq}
\end{align}
Here the angular bracket indicates an average over all generator particles.
Hence, each particle, at position $X_a$, experiences a competition between the mean discriminator gradient over the ensemble, $g_2 \langle \nabla D(X)\rangle$, and the contribution from its local gradient, $\left( \frac{g_1 - g_2}{N}\right) \nabla D(X_a)$.

The entire system is trained using Algorithm \ref{alg:gan_training_2}.

\begin{figure}[t]
\includegraphics[width=.5\textwidth]{x_axis_revised_mode_collapse_relu_updated.png}
\caption{
\textbf{Phase diagram of the transition to mode collapse using a ReLU discriminator.}
Mode collapse metric values are shown as a function of generator NTK and discriminator training rate.
Brighter points indicate mode collapse while darker points correspond to an even distribution of generator particles over the target data.
The x-axis gives the value of $g_2 / g_1$ for a given experiment, while the y-axis indicates the number of training steps the discriminator takes at each iteration.
As $g_2 / g_1$ is increased, the discriminator requires more `time' (a greater $n_{disc.}$ value) in order to shift the training from mode collapse (bright points) to convergence (darker points).
Experimental results are taken after 5000 training iterations, and data is time-averaged, representing the mean of results taken at $5000 \pm n\times 20$, with $n = 0, 1, 2, 3$. A line is fit to $(1/2) \log 8$, indicating a power-law boundary.}
\label{fig:relu_data}
\end{figure}

\subsection{Results and Interpretation}
\label{interpretation}

We run the model GAN training algorithm for each ($g_2 / g_1$, $n_{disc.}$) pair considered.
After training for a fixed number of iteration
, and computing the mode collapse metric (Eq.~\ref{eq:mode_collapse_metric}) for all pairs, we can observe a clear transition from convergence (blue data-points) to mode collapse (yellow data-points), as shown in Fig~\ref{fig:relu_data}. In the following, we provide a heuristic argument to explain the observed characteristics of the mode-collapse transition.

During training, the discriminator seeks to maximize the difference between its expectation on the training data and on the generator distribution.
Suppose to this end, the discriminator has formed its minimum within a region (cluster) of generator particles.
 According to the equation of motion for the generator particles (Eq. 26), if the term involving local gradients dominates over $g_2 \langle \nabla D(X) \rangle$, then the particles in this cluster `split apart'; each particle follows its own local gradient, regardless of the location of the discriminator minimum within the region. As a result,  the generator particle cluster, which corresponds to mode collapse, can be split and the full targeted distribution can be recovered.

However, for sufficiently large $g_2$, the term $g_2 \langle \nabla D(X) \rangle$ may dominate. Now the location of the discriminator's minimum within the region becomes important, as it may determine both the magnitude and direction of $\langle \nabla D(X) \rangle$. As is depicted in Fig.~\ref{fig:slip_away}, if the discriminator obtains a minimum far from the center of a cluster, $|\langle \nabla D(X) \rangle|$ would become non-negligible, leading to an onset of instability, and causing the entire group of generator particles to ``slip away" -- including those on the opposing slope (panel a). In contrast, for a minimum closer to the cluster's center, the mean gradient experienced by the particle cloud becomes sufficiently small to allow the cluster to ``split apart" (panel b).

\begin{figure}[t]
\includegraphics[width=.5\textwidth]{revised_x_axis_updated_loglikelihood_transition_relu.png}

\caption{
\textbf{Precise convergence occurs well above the mode-collapse transition.}
The average log-likelihood of generator particles is depicted following the same procedure as in Fig.~\ref{fig:relu_data}.
The power-law transition of Fig.~\ref{fig:relu_data} is shown as a reference; only sufficiently far above this boundary would particles converge precisely to the eight modes.
}
\label{fig:relu_loglikelihood_data}
\end{figure}

In this way, the discriminator's precision in minimizing its value over a cluster of generator particles influences its ability to split apart the cluster. More spatially precise discriminators may yield smaller values of $|\langle \nabla D(X) \rangle|$, allowing local gradients to dominate particle dynamics.

The mode-collapse data resulting from using a ReLU discriminator is shown in Fig.~\ref{fig:relu_data} on a log-log scale. A dashed white line emphasizes a visible power-law boundary separating mode collapse from convergence. Examples of generator particle distributions sampled across this transition can be found in Appendix~\ref{app:sample_dists}.
Fig.~\ref{fig:relu_loglikelihood_data} plots the log-likelihood data and shows that only sufficiently far above the transition boundary would particles converge precisely to the target modes.

This power-law behavior matches another feature of wide ReLU networks, referred to as the frequency principle \cite{frequency_vs_training, convergence_vs_freq, linear_f_principle}. As networks learn, they tend to first learn lower-frequency functions, before including higher-frequency contributions. This behavior thus sets a rate $\gamma(k)$ at which a network can learn a feature of spatial frequency $k$.
For example, within wide ReLU networks, $\gamma(k)$
is expected to be power law,
whereas for wide Tanh networks, an exponential $\gamma(k)$ is predicted \cite{linear_f_principle}.
\iffalse
{\color{red}
Interestingly, using a ReLU discriminator, the mode-collapse metric value shows a power-law boundary separating mode collapse from convergence.
This is shown in Fig.~\ref{fig:relu_data} on a log-log scale.
Using a Tanh discriminator, however, we see a boundary which appears exponential, as is depicted on a log-linear plot in Fig.~\ref{fig:tanh_data}. Examples of generator particle distributions sampled across the transition can be found in Appendix~\ref{app:sample_dists}.
}\fi

\iffalse
\begin{figure}[H]
\includegraphics[width=.5\textwidth]{tanh_window_3_scatter_2500_time.png}

\caption{
\textbf{Phase diagram using a Tanh discriminator.}
Results are depicted as in Fig.~\ref{fig:relu_data}, although experimental results are taken after 2500 training iterations.
The line fit to the transition boundary depicts an exponential shape, differing from the power-law boundary observed in Fig.~\ref{fig:relu_data}.
\iffalse
The ratio, $g_2 / g_1$ is given on the x-axis for each experiment, while the y-axis gives the number of discriminator steps allowed at each iteration.
As for the ReLU results, data is time-averaged, showing the mean of results taken at $2500 \pm n\times 20$, with $n = 0, 1, 2, 3$. The magnitude of the mode collapse metric is shown in for each experiment, with bright/dark points indicating mode collapse/convergence. A line is fit to the values $(1/2) \log 8$, depicting a generally exponential transition. As in Fig.~\ref{fig:relu_data}, for greater $g_2 / g_1$, more training time (a greater $n_{disc.}$) is required in order for training to be shifted from mode collapse (bight markers) to convergence (darker markers).
\fi
}
\label{fig:tanh_data}
\end{figure}
\fi

If we identify spatial features of size $1/k$ as having a dominant spatial frequency of $k$, then simple arguments suggest (Appendix~\ref{section:precision_scaling_argument})
that in order to split a cluster, the maximum allowable spatial imprecision falls with increasing $g_2$ as $(g_1 - g_2)/g_2$.
This indicates that to break apart such a distribution, we require the discriminator to learn a feature with spatial \textit{frequency} proportional to $g_2 / (g_1 - g_2)$; for $g_2 \ll g_1$, roughly, $k \sim g_2 / g_1$.

Assuming that the discriminator has a frequency-dependent learning rate, $\gamma(k)$, then the time required to learn such a feature scales as $T \sim 1/\gamma(k)$. The necessary discriminator steps, $n_{disc}$ (and learning rate, $\alpha_D$), to overcome mode collapse would then scale as $1/\gamma(k)$. Interestingly, experiments involving Tanh discriminators (and other types) seem to support such a conclusion (Appendix~\ref{app:fprince_support}).
As such, the frequency principle suggests a connection between the relevant time-scales and length-scales of the experiment's learning objective.

\begin{figure}[ht]
\includegraphics[width=.48\textwidth]{updated_slip_away_diagram.pdf}
\caption{
\textbf{Increased discriminator precision may result in splitting of a cluster of generator particles}.
A cluster of generator particles is depicted as a uniform shaded region (with distribution $p_\theta$) between two modes of the target distribution ($p_\text{data}$). The value of the discriminator is shown as a simple function of the form $D_\phi(x) \equiv a |x-b|$. Arrows indicate the local velocities of particles on either side of the discriminator minimum. (a) Imprecision of the discriminator leads to a sub-minimal expectation, $\langle D(X) \rangle$, and a significant $\langle \nabla D(X) \rangle$ value, causing all particles to slip to the left due to all-to-all coupling through the NTK. (b) The discriminator has found the precise minimum at the center of the distribution, which kills off average gradients and allows particles to ascend their local slopes; the cluster is thus splitting apart.
}
\label{fig:slip_away}
\end{figure}

\section{Critical Regularization }
\label{sec:critreg}

Various regularization techniques have been applied to the problem of mode-collapse avoidance \cite{veegan, mode_reg, youssef_mmd_paper}.
Here, we demonstrate that even in our simplified model GAN, the effect of regularization on reducing mode collapse can be observed.

Following \cite{youssef_mmd_paper}, we introduce a gradient regularizer
\begin{equation}
    \beta ||\nabla_\theta \langle D(G_\theta (z))\rangle ||^2 /2
    \label{eq:reg1}
\end{equation}
into the discriminator's loss function during training.
Since the velocities of generator parameters are driven by local gradients of the discriminator, this term is analogous to the kinetic energy of these parameters. This regularization term penalizes sharp gradients and encourages generators to take smoother paths to the target distribution. The effect of the regularizer on the GAN system can be viewed in analogy to a damping term in physics (a connection made explicit in Appendix~\ref{app:dirac_gan_reg}), with oscillations from mode to mode corresponding to an under-damped regime, slow convergence to all available modes corresponding to over-damping, and a most efficient convergence corresponding to critical-damping. Using this analogy as a conceptual starting point, we can sweep over $\beta$ to identify a regime of ``critical regularization".

Despite the fact that our model GAN setup does not have any reference to generator parameters (or a generator network), we may still incorporate such a term into training via the effective NTK:
\begin{multline*}
    ||\nabla_\theta \langle D(G_\theta (z))\rangle ||^2 = \\
    \frac{1}{N^2}\sum_{z, \ z' \sim q(z)} \nabla_i D(G_\theta (z)) \nabla_\theta G^i_\theta (z) \nabla_\theta G^j_\theta (z') \nabla_j D(G_\theta (z')) \\
    = \frac{1}{N^2}\sum_{z, \ z' \sim q(z)} \nabla_{i} D(G(z)) \Gamma_\theta^{i, j}(z, z') \nabla_{j} D(G(z')).
\end{multline*}
This effective form immediately allows us to apply the regularizer within the model GAN by including the following term in Eq.~\ref{eq:loss_full}:
\begin{equation}
    \frac{\beta}{N^2} \sum_{a, b} \nabla_{i} D(X_a) \Gamma^{i, j}_{a, b} \nabla_{j} D(X_b).
    \label{eq:reg}
\end{equation}

We can understand the effect of such a regularizer by expressing (\ref{eq:reg}) in the form,
\begin{equation}
    \beta \left( g_2 |\langle \nabla D(X)\rangle|^2 + \frac{(g_1 - g_2)}{N}\langle |\nabla D(X)|^2 \rangle \right).
\end{equation}
The second term discourages sharp gradients from being provided to generator particles, leading to smoother paths to convergence.
The first term, directly proportional to $g_2$, can be seen to discourage the presence of large mean gradients over the ensemble of generator particles.
Incorporating this into the same setup used to produce mode collapse in Fig.~\ref{fig:with_ntk_2d} and repeating the procedure \footnote{Here, we set $\beta = 100$}, we now observe convergence instead (Fig. \ref{fig:reg_dyn_scatters}).

\iffalse
{\color{red}
\textbf{Question: }
Is there a way to connect this regularizer directly to the competition we pointed out between the terms $g_2 \langle \nabla D(X)\rangle$ and $\frac{g_1 - g_2}{N} \nabla D(X)$?

Or is that asking for too much?
}
\fi

\iffalse
 In Section~\ref{interpretation}, we implicated the competition between the terms $g_2 \langle \nabla D(X)\rangle$ and $\frac{g_1 - g_2}{N} \nabla D(X)$ in the presence of mode collapse.
By penalizing large average gradients, this regularizer helps address this source, and encourages velocities which decorrelate the generator particles.

{\color{red}\textbf{Question:} Doesnt this regularizer discourage both terms from being large? Whats the right way to explain this? Should I just leave out the part starting with ``In Section \ref{interpretation}, we implicated the competition..."?

Perhaps going too far into this would raise questions like ``why not simply penalize the particles' kinetic energy instead of the parameters' KE? That would give us
\begin{equation}
    \beta \left( (g_2^2 + 2 g_2 \frac{g_1-g_2}{N}) |\langle \nabla D(X)\rangle|^2 + (\frac{(g_1 - g_2)}{N})^2\langle |\nabla D(X)|^2 \rangle \right).
\end{equation}
For $g_2 > \frac{(g_1 - g_2)}{N}$, this penalizes the mean gradient relatively \textit{more} than the previous regularizer, right?
Wouldnt this be preferable?

Which I think are questions beyond the scope of this paper. What do you think?
}
\fi

\begin{figure}[ht]
    \centering
    \includegraphics[width=.5\textwidth]{regularized.pdf}
    \caption{
    \textbf{Model GAN dynamics with significant off-diagonal NTK values ($g_2 / g_1 = 1/5$) converge under regularization.}
    Model GAN dynamics is shown over time, taking $g_2/g_1 = 1/5$ and including a regularizer (Eq.~\ref{eq:reg}, with $\beta=100$). Despite that without the regularizer the system oscillates from mode to mode (Fig.~\ref{fig:with_ntk_2d}), now particles converge evenly and steadily to the eight modes.}
    \label{fig:reg_dyn_scatters}
\end{figure}

Within our setup, we can experiment with the regularization parameter $\beta$. Running model GAN experiments using different $\beta$ values, we note regions corresponding to under and over-regularization, and an intermediate regime of critical regularization. In this regime, convergence is most efficiently achieved (Fig. \ref{fig:reg_dyn}).

\begin{figure}[ht]
    \centering
    \includegraphics[width=.5\textwidth]{rescaled_relu_regularization_plot.pdf}
    \caption{
    \textbf{Critical regularization mitigates mode collapse within a model GAN.}
    Upper: each point represents an experiment using a single-hidden layer ReLU discriminator, halted after 1600 iterations (similar to the experiments of Section~\ref{sec:mode_experiments}). 200 generator points are used, with $g_2 / g_1 = 0.275$ and $n_{disc.} = 5$, placing the system within the regime of mode collapse.
    The regularizer, which plays a role similar to that of a damping term, results in a reduction of mode collapse. We observe three regimes analogous to under, critically, and over-damped dynamics.
    Lower: snapshots of sampled generator particle configurations for the three circled points in the upper diagram. Note that only the critically regularized example has converged.
    }
    \label{fig:reg_dyn}
\end{figure}

\section{Discussion}

In this paper, we consider a model GAN system constructed by incorporating limiting features present within real GANs.
The generator inputs are taken to be sampled uniformly from a sphere of high dimension.
Additionally, the generator is assumed to be of infinite width and to have a static NTK such that given inputs of fixed magnitude, the NTK is a function only of their dot product ($\Gamma(z, z') = \Gamma(z \cdot z')$, as is the case in infinite-width ReLU networks).
We also modify the training procedure, using a single fixed mini-batch of generator seeds throughout training.
Under these approximations, the outputs of an infinite-width generator may be represented as a cloud of particles, whose velocities are coupled through the generator network's NTK. Further, due to the nature of the assumed NTK and the high dimensional inputs, we argue that this coarse-grained NTK may be characterized using just two values.

\iffalse
In our limiting regime \textcolor{- say explicitly what 'limiting regime' means}, the outputs of an infinite-width generator are replaced by a cloud of particles, whose velocities are coupled through the generator network's NTK.
Further, we argue that under our approximations \textcolor{blue}{- say concisely what are the key approximations}, this coarse-grained NTK may be characterized using just two values.
\fi

Despite the simplicity of our model GAN system, we observe that it is able to exhibit the defining symptoms of mode collapse – generator outputs fail to become diverse. Indeed, the simplified particle-based setting allows for physical interpretation of the phenomenon through competition between the local gradient experienced by each particle, and the average discriminator gradient experienced by the cloud as a whole. When the latter dominates over the former, the particle cloud fails to split and hence cannot cover a target diversity of modes.
From this physically motivated effective description, we are able to connect the ratio of the two effective NTK values to the occurrence or avoidance of mode collapse.

Because the generator NTK values within the model GAN setup can be easily modified, our framework makes it possible to study learning dynamics over a broad range of generators.
Using simple discriminators with a single hidden layer, we investigated the onset of mode collapse as a function of the NTK parameters and the relative training rates of the discriminator and generator.
We were able to identify power-law and exponential relationships, and explain their presence by drawing a connection to the frequency principle; frequency-dependent learning rates alone suffice to explain the shape of the transition boundary. To our knowledge, this is the first time that the principle has been observed in the context of GANs.
As a consequence, for a given NTK matrix, mode collapse can be avoided by allowing sufficient time for the discriminator to learn finer features characterizing a multi-modal target distribution.

Would it be possible to reduce the training time while avoiding mode collapse?
We have experimented with a regularizer designed to reduce mode collapse in real GANs, the effect of which is to dampen the velocity of the generator parameters during training.
Despite the fact that our model contains no generator parameters, and focuses instead on the dynamics of the outputs, we show how it is possible to adapt such a regularizer to our particle-based setting.
We demonstrate that, in our effective model too, such a regularizer can encourage smooth paths to convergence.
Importantly, an analogy to a damped oscillator, made clear through examples, enables us to identify regimes analogous to over-damping, under-damping, and critical damping. Intermediate regularization strengths would allow most efficient convergence, suggesting
``critical regularization" as a potential means to cure mode collapse and shorten training time.

The problem of understanding GAN convergence is complex.
By essentializing key features of real GANs, we have probed GAN failure in a more physically interpretable setting, which allows for extensive experimentation.
However, the model's assumptions also suggest directions of future work in studying deviations from these limiting approximations, and in mapping the lessons learned to more realistic GAN settings.

We have, for instance, assumed a time-independent NTK with uniform values throughout data-space. In reality, however, for networks of finite width, the NTK evolves throughout training. Indeed, such dynamic corrections may be studied order by order (in $1 / \text{network width}$) \cite{deep_learning_theory_book} and incorporated into a more complete analysis.
An NTK function which develops spatial features during training
(that is, an NTK defined in data-space, $\Gamma(X, X')$)
 might yield dynamics showing closer parallels to flocking, in which individual birds look at spatial neighbors to update their velocities \cite{flock, toner_long-range_1995}.

We have also replaced the distribution of NTK values (comprising the evaluations of $\Gamma(z, z)$ and $\Gamma(z \neq z')$ within a mini-batch) with just two numbers: $g_1$ and $g_2$.
 By including some variance in the diagonal and off-diagonal NTK values, as is present in realistic settings of finite latent-space dimension (see, for example, Fig.~\ref{fig:NTK_distribution}), future work might broaden the scope of the noted universality.

 Finally, we have obtained our results using a modified training algorithm in which only one mini-batch of seeds is used throughout training. In order to extrapolate the lessons learned to a more realistic setting, we would like to better understand the implications of our results for contexts in which mini-batches of seeds $\{ z\}$ are continually resampled.

\section{Acknowledgments}
During the preparation of this work, Steven Durr was supported by funds from the Bhaumik Institute for Theoretical Physics at UCLA. This work used computational and storage services associated with the Hoffman2 Cluster hosted by the UCLA Institute for Digital Research and Education. Shenshen Wang is grateful for support from an NSF CAREER Award (Grant No. PHY-2146581).

\iffalse
- time varying NTKs,
Currently ntk values are taken to be static, and uniform across data-space.
the effect on localization of ntk values

- broader distributions of the effective ntk value for $z\neq z'$
Increase the scope of the universality

- Attempt to formalize a connection between the modified training algorithm, in which only one mini-batch of seeds is used, to a more realistic setting in which at each step a new mini-batch of seeds $\{ z\}$ is sampled. Approximating, for instance the space of inputs $g_1$ corresponds to pairs of inputs satisfying $\phi_{z, z'} < \phi_c$, for some $\phi_c$, while $g_1$ corresponds to pairs of inputs satisfying $\phi_{z, z'} < \phi_c$
\fi

\iffalse
For instance, assuming a finite-width generator leads to a time-evolving NTK. The resulting time-varying particle couplings may therefore be incorporated into our model.
\fi

\iffalse

Accounting for the widths of the distirbutions of the off diagonal component of effective ntks

generalize the degree of universality by using an ntk incorporating

Running experiments with real gans to attempt tp observe

\fi

\iffalse
In this work, we have developed a model for studying the dynamics of a form of GAN training. In the model, the outputs of the generator network are replaced by a cloud of particles, whose velocities are coupled to one another through the generator’s effective neural tangent kernel (NTK).
Taking the generator to be of infinite-width, and incorporating assumptions about the input space, the model GAN system essentializes the NTK to a coarse-grained form characterized by only two values, constants during training.
This simplification of the generator suggests a sort of universality, with a broad class of infinite width generator architectures mapping to a single pair of NTK parameters.
Despite these dramatic simplifications, this model is able to exhibit the features of mode collapse, a phenomenon observed here when the off-diagonal components of the effective NTK are made large in comparison to the diagonal values.

By varying the reduced number of parameters, we have observed the relationship between the relative training rate of the generator and discriminator, and the effective generator NTK -- identifying a phase of convergence, and a phase of mode collapse.
Moreover, we have shown a connection between the type of discriminator used, and the form of this convergence phase boundary – with the shape of the boundary being related to the frequency-dependent learning rate of the discriminator.

Finally, we have experimented with a regularizer designed to reduce mode collapse in real GANs, the effect of which is to dampen the velocity of the generator parameters during training.
Even though our model system does entirely without such generator parameters, we have shown how such a regularizer can be used within our particle-based setting, and demonstrated that here too, the presence of the regularizer can have the effect of encouraging smooth paths to convergence.
The connection to damping within an oscillator is made clear through examples, as we have identified regimes analogous to over-damping under-damping and critical damping.

The problem of encouraging convergence within GANs is a significant one.
While the approximations made here allow for  experimentation and interpretation, they also suggests a direction of future work in connecting lessons gleaned from this toy model to more realistic GAN environments.
\fi

\iffalse
In this work we have introduced a framework for modeling GAN dynamics based on properties present within realistic GANs.

By

In the model GAN system, the generator is replaced with a cloud of particles coupled by an effective NTK.

In the model GAN system, generator outputs are modeled by particles,

Fill in / combine once body of paper is solid

\begin{enumerate}
    \item We have used an effective model to describe the dynamics of generator outoputs, without using a generator
    \item We have used this model to investigate mode collapse transition boundaries.

    \item (to my knowledge) the first connection between the f-principle and GAN training

    \item We have also shown that the model can represent the effect of regularizers.

    \item Discuss the importance of NTK's in neural networks. If NTK's were truly diagonal, neural networks would not train, as all points would need to be trained individually.

    \item We would like the particles we use to represent the distribution of the generator. [to be an unbiased estimator of the pushforward distribution] This means that the $g_1$ and $g_2$ values must be scaled as $N$ is varied.

\end{enumerate}
\fi

\iffalse
{\color{red}

\section{General comments -- YT}

\begin{enumerate}
    \item We need to justify more clearly the use of NTK to replace the generator.
    \item We need to justify the simplification of using two values $g_1$ and $g_2$ to approximate the NTK.
    \item We need to make the figures and figure captions more coherent and conform them to the format according to normal standard.
    \item We need to describe a bit more about the dynamics of the discriminator.
    \item We need to explain the "frequency" learning argument a bit more precisely.
    \item We need to provide an intuitive explanation and interpretation of how the "critical" regularizer works especially in the context of the NTK-based GAN used in our paper. Youseff's insights are needed here.
\end{enumerate}
}
\fi

\newpage