Learning Implicit Generative Models Using
Differentiable Graph Tests

The problem of two-sample testing for distributional equality has received significant interest in statistics. For example, the celebrated Kolmogorov-Smirnov test compares two one dimensional distributions by taking the maximal difference of the empirical CDFs. Another one-dimensional test is the runs test of Wald and Wolfowitz [1], which has been extended to the multivariate case by Friedman and Rafsky [2] (FR). It is exactly this test, together with $k$-NN test originally suggested in [3] that we analyze. These tests have been analyzed in more detail by Henze and Penrose [4], and Henze [5], Schilling [6] respectively. Their asymptotic efficiency has been discussed by Bhattacharya [7]. Chen and Zhang [8] considered the problem of tie breaking when applying the FR tests to discrete data and suggested averaging over all minimal spanning trees, which can be seen as as special case of our test in the low-temperature setting. A very prominent test that has been more recently developed is the kernel maximum mean discrepancy (MMD) test of Gretton et al. [9], which we compare with in Section 5. The test statistic is differentiable and has been used for learning implicit models by Li et al. [10], Dziugaite et al. [11]. Sutherland et al. [12] consider the problem of learning the kernel by creating a $t$-statistic using a variance estimator. Moreover, they also pioneered the idea of using tests for model criticism — for two fixed distributions, one optimizes over the parameters of the test (the kernel used). The energy test of Székely and Rizzo [13], a special case of the MMD test, has been used by Bellemare et al. [14].

Other approaches for learning implicit models that do not depend on two sample tests have been developed as well. For example, one approach is by estimating the log-ratio of the distributions [15]. Another approach, that has recently sparked significant interest, and can be also seen as estimating the log-ratio of the distributions, are the generative adversarial networks (GAN) of Goodfellow et al. [16], who pose the problem as a two player game. One can, as done in [12], combine GANs with two sample tests by using them as feature matchers at some layer of the generating network [17]. Nowozin et al. [18] minimize an arbitrary $f$-divergence [19] using a GAN framework, which can be related to our approach, because the limit of our tests converge to specific $f$-divergences, as explained in Section 2. For an overview of various approaches to learning implicit models we direct the reader to Mohamed and Lakshminarayanan [20].

This test, developed by Friedman and Rafsky [2], uses the minimum-spanning tree (MST) of $\mathcal{G}(X)$ as the neighbourhood structure $U^{*}$, which can be computed using the classical algorithms of Prim [21] and Kruskal [22] in time $O(n^{2}\log n)$. If we use $d(\mathbf{x}_{i},\mathbf{x}_{j})=\|\mathbf{x}_{i}-\mathbf{x}_{j}\|$, the problem is also known as the Euclidean spanning tree problem, and in this case Henze and Penrose [4] have proven that the test is consistent and has the following asymptotic limit.

As noted by Berisha and Hero [23], after some algebraic manipulation of the right hand side of the above equation, we obtain that $1-T_{\pi^{*}}(U^{*})\frac{n_{1}+n_{2}}{2n_{1}n_{2}}$ converges almost surely to the following $f$-divergence [19]

In [23] it is also noted that if $n_{1}=n_{2}$, then $\alpha=1/2$ and in that case $D_{1/2}$ is equal to $2\int\frac{(p(\mathbf{x})-q(\mathbf{x}))^{2}}{p(\mathbf{x})+q(\mathbf{x})}d\mathbf{x}$, which is known as the symmetric $\chi^{2}$ divergence.

Maybe the most intuitive way to construct a neighbourhood structure is to connect each point $\mathbf{x}_{j}\in X$ to its $k$ nearest neighbours. Specifically, we will add the edge $\mathbf{x}_{i}\to\mathbf{x}_{j}$ to $U^{*}$ iff $\mathbf{x}_{i}$ is one of the $k$ closest neighbours of $\mathbf{x}_{j}$ as measured by $d(\mathbf{x},\mathbf{x}^{\prime})$. If one uses the Euclidean norm, then the asymptotic distribution and the consistency of the test have been proven by Schilling [6]. These results has been extended to arbitrary norms by Henze [5], who also proved the limiting behaviour of the statistic as $n\to\infty$.

As for the FR test, we can also re-write the limit as an $f$-divergence³³3This $f$ does not vanish at one, but we can simply shift it. corresponding to $f(t)=(\alpha^{2}t^{2}+(1-\alpha))/(\alpha t+(1-\alpha))$. Moreover, if we compare the integrands in $D^{\textrm{FR}}_{\alpha}$ and $D^{\textrm{NN}}_{\alpha}$, we see that they are related and they differ by the term $2\alpha(1-\alpha)p(\mathbf{x})q(\mathbf{x})$ in the numerator. The fact that they are closely related can be also seen from Figure 1, where we plot the corresponding $f$-functions for the $n_{1}=n_{2}$ case.

Even though one can directly use the smoothed test statistic $T_{\pi^{*}}^{\lambda}$ as an objective when learning implicit models, it does not necessarily mean that lower values of this statistic result in higher $p$-values. Remember that to compute a $p$-value, one has to run a permutation test by computing quantiles of $T_{\pi}^{\lambda}$ under random draws of the permutation $\pi\sim H_{0}$. However, as this procedure is not smooth and can be costly to compute, we suggest as an alternative that does not suffer from these problems the following $t$-statistic

The same strategy has been undertaken for the FR and $k$-NN tests in [2, 4, 6]. Before we show to compute the first two moments under $H_{0}$, we need to define the matrix $\Pi$ holding the second moments of the variables $\Delta_{\pi}(e)$.

While the computation of the mean is trivial, it seems that the computation of the variance needs $O(|E|^{2})$ operations. However, we can simplify its computation to $O(|E|)$ using the following result.

To better motivate the use of a $t$-statistic, we can, similarly to the arguments in [2, 4, 6], show that it is is close to a normal distribution by casting it as a generalized correlation coefficient [25, 3]. Namely, these are tests whose statistics are the form form $\kappa=\sum_{i=1}^{n}\sum_{j=1}^{n}\overline{\mu}_{i,j}b_{i,j}$, and whose critical values are computed using the distribution of $\sum_{i=1}^{n}\sum_{j=1}^{n}\overline{\mu}_{i,j}b_{\pi(i),\pi(j)}$, where $\pi$ is a random permutation on $\{1,2,\ldots,n\}$. It is easily seen that we can fit the suggested tests in this framework if we set $\overline{\mu}_{i,j}=\frac{1}{2}(\bm{\mu}(\mathbf{d}/\lambda)_{i\to j}+\bm{\mu}(\mathbf{d}/\lambda)_{j\to i})$ and $b_{i,j}=\Delta_{\pi^{*}}(\{i,j\})$. Then, using the conditions of Barbour and Eagleson [26], we obtain the following bound on the deviation from normality.

Let us analyze the above bound in the setting that we will use it — when $n_{1}=n_{2}$. First, let us look at the variance, as formulated in 2. The last term can be ignored as it is always non-negative because $\chi_{2}\geq 4\chi_{1}$ (shown in the appendix). Because $\sum_{e\in\delta(v)}\mu_{e}\geq 1$, it follows that the variance grows as $\Omega(n)$. Thus, without any additional assumption on the growth of the neighbourhoods, we have asymptotic normality as $n\to\infty$ if the numerator is of order $o(n^{1.5})$. For example, that would be satisfied if the largest neighbourhood $\max_{i}\sum_{e\in\delta(i)}\overline{\mu}_{e}$ grows as $o(n^{1/6})$. Note that in the low temperature setting (when $\lambda\to 0$), the coordinates of $\bm{\mu}$ will be very close to either zero or one. As observed by Friedman and Rafsky [2], in this case $S_{2}=O(1)$ as the nodes of both the $k$-NN and MST graphs have nodes whose degree is bounded by a constant independent of $n$ as $n\to\infty$ [27]. We also observe experimentally in Section 5 that the distribution gets closer to normality as $\lambda$ decreases.

The constraint $\nu(\cdot)$ in this case requires the total number of edges in $U$ incoming at each node to be exactly $k$. First, note that the problem completely separates per node, i.e., the marginals of edges with different target vertices are independent. Formally, if we denote by $U_{i}$ the set of edges incoming at vertex $i$, then $U_{i}$ and $U_{j}$ are independent for $i\neq j$. Hence, for each node $i$ separately, have to perform inference in

which is a special case of the cardinality potentials considered by Tarlow et al. [28], Swersky et al. [29]. Swersky et al. [29] consider the same model, and note that we can compute all marginals in time $O(nk)$ using the algorithm in [28], which works by re-writing the model as a chain CRF and running the classical forward-backward algorithm. Hence, the total time complexity to compute the vector $\bm{\mu}(\mathbf{d}/\lambda)$ is $O(n^{2}k)$. Moreover, as marginalization requires only simple operations, we can compute the derivatives with any automatic differentiation software, and we thus do not provide formulas for the second-order moments. In [29] the authors provide an approximation for the Jacobian, which we did not use in our experiments, but instead we differentiate through the messages of the forward-backward algorithm.

As a concrete example, let us work out the simplest case — the $k$-NN test with $k=1$. In this case, the smoothed statistic reduces to

where $s_{i}(\mathbf{x}_{1},\ldots,\mathbf{x}_{n})=\texttt{softmax}(-\otimes_{l\neq i}\|\mathbf{x}_{i}-\mathbf{x}_{l}\|/\lambda)$. In other words, for each $i$ you compute the softmax of the distances to all other points using $s_{i}$, and then sum up only those positions that correspond to points from the other sample. One interpretation of the loss is the following — maximize the number of incorrect predictions if we are to estimate the label $\pi(i)$ from $\mathbf{x}_{i}$ using a soft $1$-nearest neighbour approach.

Furthermore, we can also make a clear connection between the smooth $1$-NN test and neighbourhood component analysis (NCA) [30]. Namely, we can see NCA as learning a mapping $h\colon\mathbf{x}\to A\mathbf{x}$ so that the test distinguishes (by minimizing $T^{\lambda}_{\pi^{*}}$) the two samples as best as possible after applying $h$ on them. The extension of NCA to $k$-NN [31] can be also seen as minimizing the test statistic for a particular instance of their loss function.

The model that we have to perform inference in for this test seems extremely complicated and intractable at first because the constraint has the form $\nu(U)=\llbracket U\textrm{ forms a spanning tree}\rrbracket$. First, note that if $\mathbf{d}/\lambda$ had all entries equal to a constant $\gamma$, we have that $A(-\mathbf{d}/\lambda)=(1-n)\gamma+\log c_{G(X)}$, where $c_{\mathcal{G}(X)}$ is the number of spanning trees in the graph $\mathcal{G}(X)$, and can be computed using Kirchoff’s (also known as the matrix-tree) theorem. To treat the weighted case, we use the approach of Lyons [32], who has showed that the above model is a determinantal point process (DPP), so that marginalization can be done exactly as follows. First, create the incidence matrix $A\in\{-1,0,+1\}^{(n-1)\times|E|}$ of the graph $\mathcal{G}(X)$ after removing an arbitrary vertex, and construct its Laplacian $L=A\texttt{diag}\big{[}\exp(-\mathbf{d}/\lambda)\big{]}A^{T}$. Then, if we compute $H=L^{-1/2}A\texttt{diag}\big{[}\exp(-\mathbf{d}/(2\lambda))\big{]}$, the distribution $P(U)$ is a DPP with kernel matrix $K=H^{T}H$, implying that for every $W\subseteq E$

where $K_{W}$ is the $|W|\times|W|$ submatrix of $K$ formed by the rows and columns indexed by $W$. Thus, we can easily compute all marginals and the smoothed test statistic and its derivatives using (3) as

where $\mathbf{u}_{i}$ is the vector with coordinates equal to zero, except the $i$-th coordinate which is one. Note that if we first compute the inverse $L^{-1}$, all quantities of the form $L^{-1}(\mathbf{u}_{i}-\mathbf{u}_{j})$ can be computed in time $O(n)$ as the vectors $\mathbf{u}_{i}$ have a single non-zero entry, for a total complexity of $O(n^{3})$.

To speed up this computation we can leverage the existing theory on fast solvers of Laplacian systems. Let us first create from $\mathcal{G}(X)$ the graph $e^{\mathcal{G}}(X)$ that has the same structure as $\mathcal{G}(X)$, but with edge weights $e^{-d(e)/\lambda}$ instead of $d(e)$. Hence, in this graph, a large weight between $\mathbf{x}$ and $\mathbf{x}^{\prime}$ indicates that these two points are similar to one another. In $e^{\mathcal{G}}(X)$, the marginals $\bm{\mu}_{e}$ are also known as effective resistances⁵⁵5For additional properties of the effective resistances see [33].. Spielman and Srivastava [34] provide a method to compute all marginals at once in time that is $\tilde{O}(rn^{2}/\varepsilon^{2})$, where $\varepsilon$ is the desired relative precision and $r=\frac{1}{\lambda}(\max_{e}d(e)-\min_{e}d(e))$. The idea is to first solve for $Z^{T}=L^{-1}A\texttt{diag}\big{[}\exp(-\mathbf{d}/2\lambda)\big{]}R$ where $R\in\{-1/\sqrt{k},+1/\sqrt{k}\}^{|E|\times p}$ is a random projection matrix with elements chosen uniformly from $\{-1/\sqrt{k},+1/\sqrt{k}\}$ and $p=O(\log n/\varepsilon^{2})$. Then, the suggested approximation is $\mu_{i\to j}\approx\|Z(\mathbf{u}_{i}-\mathbf{u}_{j})\|^{2}$. While computing $Z$ naïvely would take $O(n^{3}+n^{2}p)$, one achieves the claimed bound with the Laplacian solver of Spielman and Teng [35].

As an extra benefit, the above connection provides an alternative interpretation of the smoothed FR test. Namely, assume that we want to create a spectral sparsifier [36] of $e^{\mathcal{G}}(X)$, which should contain significantly less edges, but be a good summary of the graph by having a similar spectrum. Spielman and Srivastava [34] provide a strategy to create such a sparsifier by sampling edges randomly, where edge $e$ is sampled proportional to $\mu_{e}$. Hence, by optimizing $T^{\lambda}_{\pi^{*}}$ we are encouraging the constructed sparsifier of $e^{\mathcal{G}}(X)$ to have in expectation as many edges as possible connecting points from $X_{1}$ with points from $X_{2}$.

In our first experiment we analyze the effect of the smoothing strength on the power of our differentiable tests. In addition to the classical FR and $k$-NN tests, we have considered the unbiased MMD test [9] with the squared exponential kernel (as implemented in Shogun [37] using the code from [12]), and the energy test [13]. The problem that we consider, which is challenging in high dimensions, is that of differentiating the distribution $\mathcal{N}(\mathbf{0},I)$ from $\mathcal{N}((\mu,0,\ldots,0),\texttt{diag}(\sigma^{2},1,\ldots,1))$. This setting was considered to be fair in [38], as the KL divergence between the distribution is constant irrespective of the dimension. To set the smoothing strength and the bandwidth of the MMD kernel (in addition to the median heuristic) we used the same strategy as in [38] by setting $\lambda=d^{\gamma}$ for varying $\gamma\in[0,1]$. The results are presented in Figure 2, where can observe that (i) our test have similar results with MMD for shift-alternatives, while performing significantly better for scale alternatives, and (ii) by varying the smoothing parameter we can significantly increase the power of the test. In the third column we present only the best performing MMD, while we present the remaining results in Appendix B. Note that we expect the power to go to zero as the dimension increases [7, 38].

As we have already hinted in the introduction, we stochastically optimize

using the Adam [39] optimizer. To optimize, we draw at each round $n_{1}$ samples from the true distribution $P$, $n_{2}=n_{1}$ samples from the base measure $Q_{0}$, and then plug them in into the smoothed $t$-statistic.

The first experiment we perform, with the goal of understanding the effects of $\lambda$, is on the toy two moons dataset from scikit-learn [40]. We show the results in Figure 3. From the second row, showing the estimated $p$-value versus the correct one (from 1000 random permutations) at several points during training, we can indeed see that the permutation null gets closer to normality as $\lambda$ decreases. Most importantly, note that the relationship is monotone, so that we would expect the optimization to not be significantly harmed if we use the approximation. Qualitatively, we can observe that the solutions have the general structure of $P$, and that they improve as we decrease $\lambda$ — the symmetry is better captured and the two moons get better separated.

Finally, we have trained several models on the MNIST [41] dataset, which we present in Figure 4. We can observe that despite the high (784) dimensional data and the fact that we use the distance directly on the pixels, the learned models generate digits that look mostly realistic and are competitive with those obtained using MMD [10, 11].

The research was partially supported by ERC StG 307036 and a Google European PhD Fellowship.