<figure><img src="image_7.png" /><figcaption>Figure 6: Evolution of the functional equation 6 along the WGF $\mu_t$ and the learned SWGF $\hat{\mu}_t$. We observe a dilation of parameter 2 between the WGF and the SWGF.</figcaption></figure>

<figure><img src="image_8.png" /><figcaption>Figure 7: Evolution of the mean taking into account the dilation parameter. $\mu$ denotes the true mean of WGF, $\hat{\mu}$ the mean obtained through SW-JKO (17) with $\tau = 0.05$ and $\mu_*$ the mean of the stationary measure. We observe that the mean of approximated measure obtained through SW-JKO seems to follow the one of the WGF.</figcaption></figure>

For each target Gaussian, we run the SW-JKO dilated scheme (18) with $\tau = 0.05$ for a RealNVP normalizing flow. We compare it with JKO-ICNN with also $\tau = 0.05$ and with Euler-Maruyama with $10^3$, $10^4$ and $5 \cdot 10^4$ particles and a step size of $10^{-3}$. For JKO-ICNN, we use, as Mokrov et al. (2021), DenseICNN with convex quadratic layers introduced in (Korotin et al., 2019) and available at https://github.com/iamalexkorotin/Wasserstein2Barycenters. For the JKO-ICNN scheme, we use our own implementation.

We compute the symmetric Kullback-Leibler divergence between the ground truth of WGF $\mu^*$ and the distribution $\hat{\mu}$ approximated by the different schemes at times $t = 0.5$ and $t = 0.9$. The symmetric Kullback-Leibler divergence is obtained as

$$ \mathrm{SymKL}(\mu^*, \hat{\mu}) = \mathrm{KL}(\mu^* || \hat{\mu}) + \mathrm{KL}(\hat{\mu} || \mu^*). \qquad (71) $$

To approximate it, we generate $10^4$ samples of each distribution and evaluate the density at those samples.

If we note $g_\theta$ a normalizing flows, $p_Z$ the distribution in the latent space and $\rho = (g_\theta)_{\#p_Z}$, then we can evaluate the log density of $\rho$ by using the change of variable formula. Let $x = g_\theta(z)$, then

$$ \log(\rho(x)) = \log(p_Z(z)) - \log|\det J_{g_\theta}(z)|. \qquad (72) $$
<figure><img src="image_9.png" /><figcaption>Figure 8: Evolution of the components of the covariance matrix taking into account the dilation parameter. Σ denotes the true covariance matrix of WGF, Σ̂ the covariance matrix obtained through SW-JKO (17) with τ = 0 . 05 and Σ ∗ the covariance matrix of the stationary distribution. We observe some difference between WGF and SWGF.</figcaption></figure>

We choose RealNVPs (Dinh et al., 2016) for the simplicity of the transformations and the fact that we can compute efficiently the determinant of the Jacobian (since we have a closed-form). A RealNVP flow is a composition of transformations T of the form

$$
\forall z \in \mathbb{R}^d, x = T(z) = (z^1, \exp(s(z^1)) \odot z^2 + t(z^1)) \quad (73)
$$

where we write z = ( z $^{1}$, z$^{2}$) and with s and t some neural networks. To modify all the components, we use also swap transformations ( i.e. ( z $^{1}$, z$^{2}$) ↦→ ( z $^{2}$, z$^{1}$)). This transformation is invertible with log det J$_{T}$ ( z ) = ∑ i s ( z 1 $_{i}$).

For JKO-ICNN, we choose strictly convex ICNNs, and can hence invert them as well as compute the density. In this case, we do not have access to a closed-form for the Jacobian. Therefore, we used backpropagation to compute it. As this experiment is in low dimension, the computational cost is not too heavy. However, there exist stochastic methods to approximate it in greater dimension. We refer to (Huang et al., 2020) and (Alvarez-Melis et al., 2021) for more explanations.

We approximate the functional by using Monte-Carlo approximation as in Section 3.3.

For Euler-Maruyama, as in (Mokrov et al., 2021), we use kernel density estimation in order to approximate the density. We use the scipy implementation (Virtanen et al., 2020) “gaussian_kde” with the Scott’s rule to choose the bandwidth.

Finally, we report on the Figure 9 the mean of the log of the symmetric Kullback-Leibler divergence over 15 Gaussians in each dimension and the 95% confidence interval.

For the training of the neural networks, we use an Adam optimizer (Kingma & Ba, 2014) with a learning rate of 10 − 4 for RealNVP (except for the 1st iteration where we take a learning rate of 5 · 10 − $^{3}$) and of 5 · 10 − 3 for JKO-ICNN. At each inner optimization step, we start from a deep copy of the last neural network, and optimize RealNVP for 200 epochs and ICNNs for 500 epochs, with a batch size of 1024.

We see on Figure 9 that the results are better than the particle schemes obtained with Euler-Maruyama (EM) with a step size of 10 − 3 and with either 10 $^{3}$, 104 or 5 · 10 4 particles in dimension higher than 2. However, JKO-ICNN obtained better results.
<figure><img src="image_9.png" /><figcaption>Figure 8: Evolution of the components of the covariance matrix taking into account the dilation parameter. Σ denotes the true covariance matrix of WGF, Σ̂ the covariance matrix obtained through SW-JKO (17) with τ = 0 . 05 and Σ ∗ the covariance matrix of the stationary distribution. We observe some difference between WGF and SWGF.</figcaption></figure>

We choose RealNVPs (Dinh et al., 2016) for the simplicity of the transformations and the fact that we can compute efficiently the determinant of the Jacobian (since we have a closed-form). A RealNVP flow is a composition of transformations T of the form

$$
\forall z \in \mathbb{R}^d, x = T(z) = (z^1, \exp(s(z^1)) \odot z^2 + t(z^1)) \quad (73)
$$

where we write z = ( z $^{1}$, z$^{2}$) and with s and t some neural networks. To modify all the components, we use also swap transformations ( i.e. ( z $^{1}$, z$^{2}$) ↦→ ( z $^{2}$, z$^{1}$)). This transformation is invertible with log det J$_{T}$ ( z ) = ∑ i s ( z 1 $_{i}$).

For JKO-ICNN, we choose strictly convex ICNNs, and can hence invert them as well as compute the density. In this case, we do not have access to a closed-form for the Jacobian. Therefore, we used backpropagation to compute it. As this experiment is in low dimension, the computational cost is not too heavy. However, there exist stochastic methods to approximate it in greater dimension. We refer to (Huang et al., 2020) and (Alvarez-Melis et al., 2021) for more explanations.

We approximate the functional by using Monte-Carlo approximation as in Section 3.3.

For Euler-Maruyama, as in (Mokrov et al., 2021), we use kernel density estimation in order to approximate the density. We use the scipy implementation (Virtanen et al., 2020) “gaussian_kde” with the Scott’s rule to choose the bandwidth.

Finally, we report on the Figure 9 the mean of the log of the symmetric Kullback-Leibler divergence over 15 Gaussians in each dimension and the 95% confidence interval.

For the training of the neural networks, we use an Adam optimizer (Kingma & Ba, 2014) with a learning rate of 10 − 4 for RealNVP (except for the 1st iteration where we take a learning rate of 5 · 10 − $^{3}$) and of 5 · 10 − 3 for JKO-ICNN. At each inner optimization step, we start from a deep copy of the last neural network, and optimize RealNVP for 200 epochs and ICNNs for 500 epochs, with a batch size of 1024.

We see on Figure 9 that the results are better than the particle schemes obtained with Euler-Maruyama (EM) with a step size of 10 − 3 and with either 10 $^{3}$, 10 4 or 5 · 10 4 particles in dimension higher than 2. However, JKO-ICNN obtained better results.
<figure><img src="image_10.png" /><figcaption>Figure 9: On the left: SymKL divergence at time 0.5 and 0.9 between the groundtruth of the Fokker-Planck equation at time t and the solution of the SW Gradient Flow at time t. On the right: Absolute error between the functionals evaluated for WGF and SWGF at time t.</figcaption></figure>

<figure><img src="image_11.png" /><figcaption>Figure 10: Symmetric KL divergence between the learned distribution at time t = 8 and the true stationary solution on Gaussians</figcaption></figure>

## D.2 Convergence to stationary distribution

Here, we want to demonstrate that, through the SW-JKO scheme, we are able to find good minima of functionals using simple generative models.

**Gaussian.** For this experiment, we place ourselves in the same setting of Section 4.1. We start from $\mu_0 = \mathcal{N}(0, I)$ and use a step size of $\tau = 0.1$ for 80 iterations in order to match the stationary distribution. In this case, the functional is

$$ \mathcal{F}(\mu) = \int V(x)d\mu(x) + \mathcal{H}(\mu) \qquad (74) $$

with $V(x) = -\frac{1}{2}(x-b)^T A(x-b)$, and the stationary distribution is $\rho^*(x) \propto e^{-V(x)}$, hence $\rho^* = \mathcal{N}(b, A^{-1})$.

We generate 15 Gaussians for $d$ between 2 and 12, and $d \in \{20, 30, 40, 50, 75, 100\}$. Due to the length of the diffusion, and to numerical unstabilities, we do not report results obtained with JKO-ICNN. In Figure 2, we showed the results in low dimension (for $d \in \{2, \dots, 12\}$) and the instability of JKO-ICNN. We report on Figure 10 the SymKL also in higher dimension.

We use 200 epochs of each inner optimization and an Adam optimizer with a learning rate of $5 \cdot 10^{-3}$ for the first iteration and $10^{-3}$ for the rest. We also use a batch size of 1000 sample.

**Bayesian logistic regression.** For the Bayesian logistic regression, we have access to covariates $s_1, \dots, s_n \in \mathbb{R}^d$ with their associated labels $y_1, \dots, y_n \in \{-1, 1\}$. Following (Liu & Wang, 2016; Mokrov et al., 2021), we put as prior on the regression weights $w$, $p_0(w|\alpha) = \mathcal{N}(w; 0, \frac{1}{\alpha})$ with $p_0(\alpha) = \Gamma(\alpha; 1, 0.01)$.
<figure><img src="image_10.png" /><figcaption>Figure 9: On the left: SymKL divergence at time 0.5 and 0.9 between the groundtruth of the Fokker-Planck equation at time t and the solution of the SW Gradient Flow at time t. On the right: Absolute error between the functionals evaluated for WGF and SWGF at time t.</figcaption></figure>

<figure><img src="image_11.png" /><figcaption>Figure 10: Symmetric KL divergence between the learned distribution at time t = 8 and the true stationary solution on Gaussians</figcaption></figure>

## D.2 Convergence to stationary distribution

Here, we want to demonstrate that, through the SW-JKO scheme, we are able to find good minima of functionals using simple generative models.

**Gaussian.** For this experiment, we place ourselves in the same setting of Section 4.1. We start from $\mu_0 = \mathcal{N}(0, I)$ and use a step size of $\tau = 0.1$ for 80 iterations in order to match the stationary distribution. In this case, the functional is

$$ \mathcal{F}(\mu) = \int V(x)d\mu(x) + \mathcal{H}(\mu) \qquad (74) $$

with $V(x) = -\frac{1}{2}(x-b)^T A(x-b)$, and the stationary distribution is $\rho^*(x) \propto e^{-V(x)}$, hence $\rho^* = \mathcal{N}(b, A^{-1})$.

We generate 15 Gaussians for $d$ between 2 and 12, and $d \in \{20, 30, 40, 50, 75, 100\}$. Due to the length of the diffusion, and to numerical unstabilities, we do not report results obtained with JKO-ICNN. In Figure 2, we showed the results in low dimension (for $d \in \{2, \dots, 12\}$) and the instability of JKO-ICNN. We report on Figure 10 the SymKL also in higher dimension.

We use 200 epochs of each inner optimization and an Adam optimizer with a learning rate of $5 \cdot 10^{-3}$ for the first iteration and $10^{-3}$ for the rest. We also use a batch size of 1000 sample.

**Bayesian logistic regression.** For the Bayesian logistic regression, we have access to covariates $s_1, \dots, s_n \in \mathbb{R}^d$ with their associated labels $y_1, \dots, y_n \in \{-1, 1\}$. Following (Liu & Wang, 2016; Mokrov et al., 2021), we put as prior on the regression weights $w$, $p_0(w|\alpha) = \mathcal{N}(w; 0, \frac{1}{\alpha})$ with $p_0(\alpha) = \Gamma(\alpha; 1, 0.01)$.
<table><caption>Table 3: Number of features, of samples and batch size of each dataset.</caption><thead><tr><th></th><th>covtype</th><th>german</th><th>diabetis</th><th>twonorm</th><th>ringnorm</th><th>banana</th><th>splice</th><th>waveform</th><th>image</th></tr></thead><tbody><tr><th>features</th><td>54</td><td>20</td><td>8</td><td>20</td><td>20</td><td>2</td><td>60</td><td>21</td><td>18</td></tr><tr><th>samples</th><td>581012</td><td>1000</td><td>768</td><td>7400</td><td>7400</td><td>5300</td><td>2991</td><td>5000</td><td>2086</td></tr><tr><th>batch size</th><td>512</td><td>800</td><td>614</td><td>1024</td><td>1024</td><td>1024</td><td>512</td><td>512</td><td>1024</td></tr></tbody></table>

<table><caption>Table 4: Hyperparameters for SWGFs with RealNVPs. nl: number of coupling layers in RealNVP, nh: number of hidden units of conditioner neural networks, lr: learning rate using Adam, JKO steps: number of SW-JKO steps, Iters by step: number of epochs for each SW-JKO step, τ: the time step, batch size: number of samples taken to approximate the functional.</caption><thead><tr><th></th><th>covtype</th><th>german</th><th>diabetis</th><th>twonorm</th><th>ringnorm</th><th>banana</th><th>splice</th><th>waveform</th><th>image</th></tr></thead><tbody><tr><th>nl</th><td>2</td><td>2</td><td>2</td><td>2</td><td>2</td><td>2</td><td>5</td><td>5</td><td>2</td></tr><tr><th>nh</th><td>512</td><td>512</td><td>512</td><td>512</td><td>512</td><td>512</td><td>128</td><td>128</td><td>512</td></tr><tr><th>lr</th><td>2e<sup>-5</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-4</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-5</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-4</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-5</sup></td></tr><tr><th>JKO steps</th><td>5</td><td>5</td><td>10</td><td>20</td><td>5</td><td>5</td><td>5</td><td>5</td><td>5</td></tr><tr><th>Iters by step</th><td>1000</td><td>500</td><td>500</td><td>500</td><td>1000</td><td>500</td><td>500</td><td>500</td><td>500</td></tr><tr><th>τ</th><td>0.1</td><td>10<sup>-6</sup></td><td>5 · 10<sup>-6</sup></td><td>10<sup>-8</sup></td><td>10<sup>-6</sup></td><td>0.1</td><td>10<sup>-6</sup></td><td>10<sup>-8</sup></td><td>0.1</td></tr><tr><th>batch size</th><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>512</td><td>1024</td></tr></tbody></table>

Therefore, we aim at learning the posterior $p(w, \alpha|y)$:

$$p(w, \alpha|y) \propto p(y|w, \alpha)p_0(w|\alpha)p_0(\alpha) = p_0(\alpha)p_0(w|\alpha) \prod_{i=1}^{n} p(y_i|w, \alpha)$$

where $p(y_i|w, \alpha) = \sigma(w^T s_i)^{\frac{1+y_i}{2}} (1 - \sigma(w^T s_i))^{\frac{1-y_i}{2}}$ with $\sigma$ the sigmoid. To evaluate $V(\mu) = \int V(x) d\mu(x)$, we resample data uniformly.

In our context, let $V(x) = -\log(p_0(\alpha)p_0(w|\alpha)p(y|w, \alpha))$, then using $F(\mu) = \int V d\mu + H(\mu)$ as functional, we know that the limit of the stationary solution of Fokker-Planck is proportional to $e^{-V} = p(w, \alpha|y)$.

Following Mokrov et al. (2021); Liu & Wang (2016), we use the 8 datasets of Mika et al. (1999) and the covtype dataset (https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html).

We report in Table 3 the characteristics of the different datasets. The datasets are loaded using the code of Mokrov et al. (2021) (https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows). We split the dataset between train set and test set with a 4:1 ratio.

We report in Table 4 the hyperparameters used for the results reported in Table 1. We also tuned the time step $\tau$ since for too big $\tau$, we observed bad results, which the SW-JKO scheme should be a good approximation of the SWGF only for small enough $\tau$.

Moreover, we reported in Table 1 the mean over 5 training. For the results obtained with JKO-ICNN, we used the same hyperparameters as Mokrov et al. (2021).

## D.3 Influence of the number of projections

It is well known that the approximation of Sliced-Wasserstein is subject to the curse of dimensionality through the Monte-Carlo approximation (Nadjahi et al., 2020). We provide here some experiment to quantify this influence. However, first note that the goal is not to minimize the Sliced-Wasserstein distance, but rather the functional, SW playing mostly a regularizer role. Experiments on the influence of the number of experiments to approximate the SW have already been conducted (see e.g. Figure 2 in (Nadjahi et al., 2020) or Figure 1 in (Deshpande et al., 2019)).
<table><caption>Table 3: Number of features, of samples and batch size of each dataset.</caption><thead><tr><th></th><th>covtype</th><th>german</th><th>diabetis</th><th>twonorm</th><th>ringnorm</th><th>banana</th><th>splice</th><th>waveform</th><th>image</th></tr></thead><tbody><tr><th>features</th><td>54</td><td>20</td><td>8</td><td>20</td><td>20</td><td>2</td><td>60</td><td>21</td><td>18</td></tr><tr><th>samples</th><td>581012</td><td>1000</td><td>768</td><td>7400</td><td>7400</td><td>5300</td><td>2991</td><td>5000</td><td>2086</td></tr><tr><th>batch size</th><td>512</td><td>800</td><td>614</td><td>1024</td><td>1024</td><td>1024</td><td>512</td><td>512</td><td>1024</td></tr></tbody></table>

<table><caption>Table 4: Hyperparameters for SWGFs with RealNVPs. nl: number of coupling layers in RealNVP, nh: number of hidden units of conditioner neural networks, lr: learning rate using Adam, JKO steps: number of SW-JKO steps, Iters by step: number of epochs for each SW-JKO step, τ: the time step, batch size: number of samples taken to approximate the functional.</caption><thead><tr><th></th><th>covtype</th><th>german</th><th>diabetis</th><th>twonorm</th><th>ringnorm</th><th>banana</th><th>splice</th><th>waveform</th><th>image</th></tr></thead><tbody><tr><th>nl</th><td>2</td><td>2</td><td>2</td><td>2</td><td>2</td><td>2</td><td>5</td><td>5</td><td>2</td></tr><tr><th>nh</th><td>512</td><td>512</td><td>512</td><td>512</td><td>512</td><td>512</td><td>128</td><td>128</td><td>512</td></tr><tr><th>lr</th><td>2e<sup>-5</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-4</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-5</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-4</sup></td><td>1e<sup>-4</sup></td><td>5e<sup>-5</sup></td></tr><tr><th>JKO steps</th><td>5</td><td>5</td><td>10</td><td>20</td><td>5</td><td>5</td><td>5</td><td>5</td><td>5</td></tr><tr><th>Iters by step</th><td>1000</td><td>500</td><td>500</td><td>500</td><td>1000</td><td>500</td><td>500</td><td>500</td><td>500</td></tr><tr><th>τ</th><td>0.1</td><td>10<sup>-6</sup></td><td>5 · 10<sup>-6</sup></td><td>10<sup>-8</sup></td><td>10<sup>-6</sup></td><td>0.1</td><td>10<sup>-6</sup></td><td>10<sup>-8</sup></td><td>0.1</td></tr><tr><th>batch size</th><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>1024</td><td>512</td><td>1024</td></tr></tbody></table>

Therefore, we aim at learning the posterior $p(w, \alpha|y)$:

$$p(w, \alpha|y) \propto p(y|w, \alpha)p_0(w|\alpha)p_0(\alpha) = p_0(\alpha)p_0(w|\alpha) \prod_{i=1}^{n} p(y_i|w, \alpha)$$

where $p(y_i|w, \alpha) = \sigma(w^T s_i)^{\frac{1+y_i}{2}} (1 - \sigma(w^T s_i))^{\frac{1-y_i}{2}}$ with $\sigma$ the sigmoid. To evaluate $V(\mu) = \int V(x) d\mu(x)$, we resample data uniformly.

In our context, let $V(x) = -\log(p_0(\alpha)p_0(w|\alpha)p(y|w, \alpha))$, then using $F(\mu) = \int V d\mu + H(\mu)$ as functional, we know that the limit of the stationary solution of Fokker-Planck is proportional to $e^{-V} = p(w, \alpha|y)$.

Following Mokrov et al. (2021); Liu & Wang (2016), we use the 8 datasets of Mika et al. (1999) and the covtype dataset (https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html).

We report in Table 3 the characteristics of the different datasets. The datasets are loaded using the code of Mokrov et al. (2021) (https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows). We split the dataset between train set and test set with a 4:1 ratio.

We report in Table 4 the hyperparameters used for the results reported in Table 1. We also tuned the time step $\tau$ since for too big $\tau$, we observed bad results, which the SW-JKO scheme should be a good approximation of the SWGF only for small enough $\tau$.

Moreover, we reported in Table 1 the mean over 5 training. For the results obtained with JKO-ICNN, we used the same hyperparameters as Mokrov et al. (2021).

## D.3 Influence of the number of projections

It is well known that the approximation of Sliced-Wasserstein is subject to the curse of dimensionality through the Monte-Carlo approximation (Nadjahi et al., 2020). We provide here some experiment to quantify this influence. However, first note that the goal is not to minimize the Sliced-Wasserstein distance, but rather the functional, SW playing mostly a regularizer role. Experiments on the influence of the number of experiments to approximate the SW have already been conducted (see e.g. Figure 2 in (Nadjahi et al., 2020) or Figure 1 in (Deshpande et al., 2019)).
<figure><img src="image_12.png" /><figcaption>Figure 11: Density over time of the solution of the aggregation equation learned over the discrete grid.</figcaption></figure>

<figure><img src="image_13.png" /><figcaption>Figure 12: Density over time of the solution of the aggregation equation by learning particles.</figcaption></figure>

Here, we take the same setting of Section 4.1, i.e. we generate 15 random Gaussians, and then vary the number of projections and report the Symmetric Kullback-Leibler divergence on Figure 3. We observe that the results seem to improve with the number of projections until it reaches a certain plateau. The plateau seems to be attained for a bigger number of dimensions in high dimension.

### D.4 Aggregation equations

Here, we use as functional

$$W(\mu) = \iint W(x - y)d\mu(x)d\mu(y). \qquad (75)$$

Carrillo et al. (2021) use a repulsive-attractive interaction potential, for $a > b \ge 0$,

$$W(x) = \frac{\|x\|^a}{a} - \frac{\|x\|^b}{b} \qquad (76)$$

using the convention $\|x\|_0 = \ln(\|x\|)$. For some values of $a$ and $b$, there is existence of stable equilibrium state (Balagué et al., 2013).

**Dirac Ring.** First, for the Dirac ring example, we take $a = 4$ and $b = 2$. Then, $W$ is a repulsive-attractive interaction potential (repulsive in the short range, and attractive in the long range because $b < a$ (Balagué
<figure><img src="image_12.png" /><figcaption>Figure 11: Density over time of the solution of the aggregation equation learned over the discrete grid.</figcaption></figure>

<figure><img src="image_13.png" /><figcaption>Figure 12: Density over time of the solution of the aggregation equation by learning particles.</figcaption></figure>

Here, we take the same setting of Section 4.1, i.e. we generate 15 random Gaussians, and then vary the number of projections and report the Symmetric Kullback-Leibler divergence on Figure 3. We observe that the results seem to improve with the number of projections until it reaches a certain plateau. The plateau seems to be attained for a bigger number of dimensions in high dimension.

### D.4 Aggregation equations

Here, we use as functional

$$W(\mu) = \iint W(x - y)d\mu(x)d\mu(y). \qquad (75)$$

Carrillo et al. (2021) use a repulsive-attractive interaction potential, for $a > b \ge 0$,

$$W(x) = \frac{\|x\|^a}{a} - \frac{\|x\|^b}{b} \qquad (76)$$

using the convention $\|x\|^0 = \ln(\|x\|)$. For some values of $a$ and $b$, there is existence of stable equilibrium state (Balagué et al., 2013).

**Dirac Ring.** First, for the Dirac ring example, we take $a = 4$ and $b = 2$. Then, $W$ is a repulsive-attractive interaction potential (repulsive in the short range, and attractive in the long range because $b < a$ (Balagué
<figure><img src="image_14.png" /><figcaption>Figure 13: Density over time of the solution of the aggregation equation approximated with a FCNN.</figcaption></figure>

<figure><img src="image_15.png" /><figcaption>Figure 14: Density over time of the solution of the aggregation equation approximated with JKO-ICNN.</figcaption></figure>

et al., 2013)), we show the densities over time obtained on a discrete grid (Figure 11), by learning the position of the particles (Figure 12) and with the FCNN (Figure 13). For the last two, the density reported is obtained with kernel density estimation where we chose by hand the bandwidth to match well the sampled points. We also report the evolution of the density for the JKO-ICNN scheme. In Figure 14, we show the evolution of the density obtained by the change of variable formula. We observe that values seem to explode which may due to the fact that the stationary solution may not have a density or to numerical unstabilities. We report in Figure 15 the densities obtained with a kernel density estimation. We also report the particles generated through a FCNN, the particle evolution and JKO-ICNN in Figure 16. We observe that the particles match perfectly the ring. For the FCNN, there seem to be some noise. JKO-ICNN recovers also well the ring but particles seem to not be uniformly distributed over the ring.

For the SW-JKO scheme, we take $\tau = 0.05$ and run it for 200 steps (from $t = 0$ to $t = 10$) starting from $\mu_0 = \mathcal{N}(0, I)$. For JKO-ICNN, we choose $\tau = 0.1$ and run it for 100 steps as the diffusion was really long.

We take a grid of 50 × 50 samples on [-1, 1]². To optimize the weights, we used an SGD optimizer with a momentum of 0.9 with a learning rate of $10^{-4}$ for 300 epochs by inner optimization scheme. For particles, we optimized 1000 particles (sampled initially from $\mu_0$) with an SGD optimizer with a momentum of 0.9, a learning rate of 1 and 500 epochs by JKO step. We take a FCNN composed of 5 hidden layers with 64 units and leaky relu activation functions. They are optimized with an Adam optimizer with a learning rate of $10^{-4}$ (except for the first iteration where we take a learning rate of $5 \cdot 10^{-3}$) for 400 epochs by JKO step. For
<figure><img src="image_14.png" /><figcaption>Figure 13: Density over time of the solution of the aggregation equation approximated with a FCNN.</figcaption></figure>

<figure><img src="image_15.png" /><figcaption>Figure 14: Density over time of the solution of the aggregation equation approximated with JKO-ICNN.</figcaption></figure>

et al., 2013)), we show the densities over time obtained on a discrete grid (Figure 11), by learning the position of the particles (Figure 12) and with the FCNN (Figure 13). For the last two, the density reported is obtained with kernel density estimation where we chose by hand the bandwidth to match well the sampled points. We also report the evolution of the density for the JKO-ICNN scheme. In Figure 14, we show the evolution of the density obtained by the change of variable formula. We observe that values seem to explode which may due to the fact that the stationary solution may not have a density or to numerical unstabilities. We report in Figure 15 the densities obtained with a kernel density estimation. We also report the particles generated through a FCNN, the particle evolution and JKO-ICNN in Figure 16. We observe that the particles match perfectly the ring. For the FCNN, there seem to be some noise. JKO-ICNN recovers also well the ring but particles seem to not be uniformly distributed over the ring.

For the SW-JKO scheme, we take $\tau = 0.05$ and run it for 200 steps (from $t = 0$ to $t = 10$) starting from $\mu_0 = \mathcal{N}(0, I)$. For JKO-ICNN, we choose $\tau = 0.1$ and run it for 100 steps as the diffusion was really long.

We take a grid of 50 × 50 samples on [-1, 1]². To optimize the weights, we used an SGD optimizer with a momentum of 0.9 with a learning rate of $10^{-4}$ for 300 epochs by inner optimization scheme. For particles, we optimized 1000 particles (sampled initially from $\mu_0$) with an SGD optimizer with a momentum of 0.9, a learning rate of 1 and 500 epochs by JKO step. We take a FCNN composed of 5 hidden layers with 64 units and leaky relu activation functions. They are optimized with an Adam optimizer with a learning rate of $10^{-4}$ (except for the first iteration where we take a learning rate of $5 \cdot 10^{-3}$) for 400 epochs by JKO step. For
<figure><img src="image_16.png" /><figcaption>Figure 15: Density over time of the solution of the aggregation equation approximated with JKO-ICNN and kernel density estimator.</figcaption></figure>

<figure><img src="image_17.png" /><figcaption>Figure 16: Samples of the stationary distribution</figcaption></figure>

JKO-ICNN, we choose the same parameters as for the previous experiments. In any case, we take $n_{\theta} = 10^4$ projections to approximate $SW_2$ and a batch size of 1000.

On Figure 17, we plot on the two first columns the stationary density learned by the different methods. On the third column, we plot the evolution of the functional along the flows.

The functionals seem to converge towards the same value. However, JKO-ICNN seems to not be able to capture the right distribution. We also note that the curve for the FCNN is not very smooth, which can probably be explained by the fact that we take independent samples at each step.

On Figure 18, we show the evolution of the functionals along different learned flows. We observe that for the discretized grid, we obtain the worse results which is understandable since we have discretization error. The particle is the most stable as particle's position do not move anymore once the stationary state is reached. For the FCNN, we observe oscillations which are due to the fact that we take independent samples at each time t. Finally, the JKO-ICNN scheme seems to converge toward the same value as the SW-JKO scheme with FCNNs.

We also tried to use normalizing flows with this functional. We observed that the training seems harder than with the FCNN. Indeed, with simple flows such as RealNVP, the model has a lot of troubles of learning the Dirac ring, probably because of the hole. More generally, since normalizing flows are bijective transformations, they must preserve topological properties, and therefore do not perform well when the
<figure><img src="image_16.png" /><figcaption>Figure 15: Density over time of the solution of the aggregation equation approximated with JKO-ICNN and kernel density estimator.</figcaption></figure>

<figure><img src="image_17.png" /><figcaption>Figure 16: Samples of the stationary distribution</figcaption></figure>

JKO-ICNN, we choose the same parameters as for the previous experiments. In any case, we take $n_{\theta} = 10^4$ projections to approximate $SW_2$ and a batch size of 1000.

On Figure 17, we plot on the two first columns the stationary density learned by the different methods. On the third column, we plot the evolution of the functional along the flows.

The functionals seem to converge towards the same value. However, JKO-ICNN seems to not be able to capture the right distribution. We also note that the curve for the FCNN is not very smooth, which can probably be explained by the fact that we take independent samples at each step.

On Figure 18, we show the evolution of the functionals along different learned flows. We observe that for the discretized grid, we obtain the worse results which is understandable since we have discretization error. The particle is the most stable as particle's position do not move anymore once the stationary state is reached. For the FCNN, we observe oscillations which are due to the fact that we take independent samples at each time t. Finally, the JKO-ICNN scheme seems to converge toward the same value as the SW-JKO scheme with FCNNs.

We also tried to use normalizing flows with this functional. We observed that the training seems harder than with the FCNN. Indeed, with simple flows such as RealNVP, the model has a lot of troubles of learning the Dirac ring, probably because of the hole. More generally, since normalizing flows are bijective transformations, they must preserve topological properties, and therefore do not perform well when the
<figure><img src="image_18.png" /><figcaption>Figure 17: Steady state and evolution of the functional of the aggregation equation for $a = 4$, $b = 2$.</figcaption></figure>

standard distribution and the target distribution do not share the same topology (e.g. do not have the same number of connected components or "holes" as it is explained in (Cornish et al., 2020)). For CPF, it worked slightly better, but was not able to fully recover the ring (at least at time $t = 5$) as we can see on Figure 19. Moreover, the training time for CPF was really huge compared to the FCNN.

**Other functionals.** As in section 3 of Carrillo et al. (2021), we also tried to use

$$W(x) = \frac{\|x\|^2}{2} - \log(\|x\|) \qquad (77)$$

as interaction potential with a FCNN. We find well that the steady state is an indicator function on the centered disk of radius 1 (Figure 20).
<figure><img src="image_18.png" /><figcaption>Figure 17: Steady state and evolution of the functional of the aggregation equation for $a = 4$, $b = 2$.</figcaption></figure>

standard distribution and the target distribution do not share the same topology (e.g. do not have the same number of connected components or "holes" as it is explained in (Cornish et al., 2020)). For CPF, it worked slightly better, but was not able to fully recover the ring (at least at time $t = 5$) as we can see on Figure 19. Moreover, the training time for CPF was really huge compared to the FCNN.

**Other functionals.** As in section 3 of Carrillo et al. (2021), we also tried to use

$$W(x) = \frac{\|x\|^2}{2} - \log(\|x\|) \qquad (77)$$

as interaction potential with a FCNN. We find well that the steady state is an indicator function on the centered disk of radius 1 (Figure 20).
<figure><img src="image_19.png" /><figcaption>Figure 18: Evolution of the aggregation functional along different flows.</figcaption></figure>

<figure><img src="image_20.png" /><figcaption>Figure 19: Density learned for the aggregation equation for CPF and JKO-ICNN.</figcaption></figure>

Another possible functional without using internal energies is to add a drift term

$$ \int V(x)\rho(x)dx. \qquad (78) $$

Then, the Wasserstein gradient flow is solution to

$$ \partial_t \rho_t = \text{div}(\rho \nabla (W * \rho)) + \text{div}(\rho \nabla V). \qquad (79) $$

Carrillo et al. (2021) use $W(x) = \frac{\|x\|^2}{2} - \log(\|x\|)$ and $V(x) = -\frac{\alpha}{\beta} \log(\|x\|)$ with $\alpha = 1$ and $\beta = 4$. Then, it can be shown (see Carrillo et al. (2021); Chen & Kolokolnikov (2014); Carrillo et al. (2015)) that the steady state is an indicator function on a torus of inner radius $R_i = \sqrt{\frac{\alpha}{\beta}}$ and outer radius $R_o = \sqrt{\frac{\alpha}{\beta} + 1}$ which we observe on Figure 21 and Figure 22.

For this last experiment, we observed some unstably issues in the training phase. It may be due to non-locality of the interaction potential $W$ as it is stated in (Carrillo et al., 2021).

## D.5 Sliced-Wasserstein flows

In this experiment, we aim at minimizing the following functional:

$$ \mathcal{F}(\mu) = \frac{1}{2} SW_2^2(\mu, \nu) + \lambda \mathcal{H}(\mu) \qquad (80) $$
<figure><img src="image_19.png" /><figcaption>Figure 18: Evolution of the aggregation functional along different flows.</figcaption></figure>

<figure><img src="image_20.png" /><figcaption>Figure 19: Density learned for the aggregation equation for CPF and JKO-ICNN.</figcaption></figure>

Another possible functional without using internal energies is to add a drift term

$$ \int V(x)\rho(x)dx. \qquad (78) $$

Then, the Wasserstein gradient flow is solution to

$$ \partial_t \rho_t = \text{div}(\rho \nabla (W * \rho)) + \text{div}(\rho \nabla V). \qquad (79) $$

Carrillo et al. (2021) use $W(x) = \frac{\|x\|^2}{2} - \log(\|x\|)$ and $V(x) = -\frac{\alpha}{\beta} \log(\|x\|)$ with $\alpha = 1$ and $\beta = 4$. Then, it can be shown (see Carrillo et al. (2021); Chen & Kolokolnikov (2014); Carrillo et al. (2015)) that the steady state is an indicator function on a torus of inner radius $R_i = \sqrt{\frac{\alpha}{\beta}}$ and outer radius $R_o = \sqrt{\frac{\alpha}{\beta} + 1}$ which we observe on Figure 21 and Figure 22.

For this last experiment, we observed some unstably issues in the training phase. It may be due to non-locality of the interaction potential $W$ as it is stated in (Carrillo et al., 2021).

## D.5 Sliced-Wasserstein flows

In this experiment, we aim at minimizing the following functional:

$$ \mathcal{F}(\mu) = \frac{1}{2} SW_2^2(\mu, \nu) + \lambda \mathcal{H}(\mu) \qquad (80) $$
<figure><img src="image_21.png" /><figcaption>Figure 20: Steady state and evolution of the functional for $W(x) = \frac{\|x\|^2}{2} - \log(\|x\|)$.</figcaption></figure>

<figure><img src="image_22.png" /><figcaption>Figure 21: Density over time of the solution of the aggregation-drift equation approximated with a FCNN.</figcaption></figure>

where $\nu$ is some target distribution from which we have access to samples. In this section, we use MNIST (LeCun & Cortes, 2010), FashionMNIST (Xiao et al., 2017) and CelebA (Liu et al., 2015).

We report the Fréchet Inception Distance (FID) (Heusel et al., 2017) for MNIST and FashionMNIST between $10^4$ test samples and $10^4$ generated samples. As they are gray images, we duplicate the gray levels into 3 channels. Moreover, we use the code of Dai & Seljak (2021) (available at https://github.com/biweidai/SINF) and reported their result for SWF (in the ambient space) and SWAE in Table 2. For SWF in the latent space, we used our own implementation.

**With a pretrained autoencoder.** First, we optimize the functional in the latent space of a pretrained autoencoder (AE). We choose the same AE as Liutkus et al. (2019) which is available at https://github.com/aliutkus/swf/blob/master/code/networks/autoencoder.py. We report results for a latent space of dimension $d = 48$.

In this latent space, we applied a FCNN and a RealNVP. In either case, we used $n_\theta = 10^3$ projections to approximate $SW_2$. We chose a batch size of 128 samples. The FCNN was chosen with 5 hidden layers of 512 units with leaky relu activation function. The RealNVP was composed of 5 coupling layers, with FCNN as scaling and shifting networks (with 5 layers and 100 hidden units). We trained the networks at each inner optimization step with a learning rate of $5 \cdot 10^{-3}$ for RealNVP (and $10^{-2}$ for the first iteration) and of $10^{-3}$ for the FCNN (and $5 \cdot 10^{-3}$ for the first iteration) during 1000 epochs.
<figure><img src="image_21.png" /><figcaption>Figure 20: Steady state and evolution of the functional for $W(x) = \frac{\|x\|^2}{2} - \log(\|x\|)$.</figcaption></figure>

<figure><img src="image_22.png" /><figcaption>Figure 21: Density over time of the solution of the aggregation-drift equation approximated with a FCNN.</figcaption></figure>

where $\nu$ is some target distribution from which we have access to samples. In this section, we use MNIST (LeCun & Cortes, 2010), FashionMNIST (Xiao et al., 2017) and CelebA (Liu et al., 2015).

We report the Fréchet Inception Distance (FID) (Heusel et al., 2017) for MNIST and FashionMNIST between $10^4$ test samples and $10^4$ generated samples. As they are gray images, we duplicate the gray levels into 3 channels. Moreover, we use the code of Dai & Seljak (2021) (available at https://github.com/biweidai/SINF) and reported their result for SWF (in the ambient space) and SWAE in Table 2. For SWF in the latent space, we used our own implementation.

**With a pretrained autoencoder.** First, we optimize the functional in the latent space of a pretrained autoencoder (AE). We choose the same AE as Liutkus et al. (2019) which is available at https://github.com/aliutkus/swf/blob/master/code/networks/autoencoder.py. We report results for a latent space of dimension $d = 48$.

In this latent space, we applied a FCNN and a RealNVP. In either case, we used $n_\theta = 10^3$ projections to approximate $SW_2$. We chose a batch size of 128 samples. The FCNN was chosen with 5 hidden layers of 512 units with leaky relu activation function. The RealNVP was composed of 5 coupling layers, with FCNN as scaling and shifting networks (with 5 layers and 100 hidden units). We trained the networks at each inner optimization step with a learning rate of $5 \cdot 10^{-3}$ for RealNVP (and $10^{-2}$ for the first iteration) and of $10^{-3}$ for the FCNN (and $5 \cdot 10^{-3}$ for the first iteration) during 1000 epochs.
<figure><img src="image_23.png" /><figcaption>Figure 22: Steady state and evolution of the functional for the aggregation-drift equation.</figcaption></figure>

<figure><img src="image_24.png" /><figcaption>Figure 23: Generated sample obtained through a pretrained decoder + FCNN.</figcaption></figure>

Following Liutkus et al. (2019), we choose $\tau = 0.5$ and $\lambda = 0$. We run it for 10 outer iterations.

We report in Figure 23 the result obtained with FCNN on MNIST and FashionMNIST. Overall, the results seem quite comparable and our method seems to perform well compared to other methods applied in latent spaces.

On Figure 24, we report samples obtained with RealNVP on MNIST, FashionMNIST and CelebA. For CelebA, we choose the same autoencoder with a latent space of dimension 48 and we ran the SW-JKO scheme for 20 steps with $\tau = 0.1$.

We can also optimize directly particles in the latent space and we show it on CelebA on Figure 25 for $\tau = 0.1$ and for 10 steps.

**In the Original Space.** We report here the results obtained by running the SW-JKO scheme in the original spaces of images, which are very high dimensional ($d = 784$ for MNIST and FashionMNIST). We obtained worse results than in the latent space. Notice that we used only 1000 projections, and ran it for 50 outer iterations with a step size of $\tau = 5$.

On Figure 26, we use a RealNVP and add a uniform dequantization (Ho et al., 2019) and learned it in the logit space as it is done in (Dinh et al., 2016; Papamakarios et al., 2017) because using normalizing flows need continuous data. The RealNVP is composed here of 2 coupling layers with FCNN also composed of 2 layers and 512 hidden units. We choose $\tau = 5$, and run it for 20 steps. At each inner optimization problem, the neural networks are optimized with an Adam optimizer with a learning rate of $10^{-2}$ and for 500 epochs.
<figure><img src="image_23.png" /><figcaption>Figure 22: Steady state and evolution of the functional for the aggregation-drift equation.</figcaption></figure>

<figure><img src="image_24.png" /><figcaption>Figure 23: Generated sample obtained through a pretrained decoder + FCNN.</figcaption></figure>

Following Liutkus et al. (2019), we choose $\tau = 0.5$ and $\lambda = 0$. We run it for 10 outer iterations.

We report in Figure 23 the result obtained with FCNN on MNIST and FashionMNIST. Overall, the results seem quite comparable and our method seems to perform well compared to other methods applied in latent spaces.

On Figure 24, we report samples obtained with RealNVP on MNIST, FashionMNIST and CelebA. For CelebA, we choose the same autoencoder with a latent space of dimension 48 and we ran the SW-JKO scheme for 20 steps with $\tau = 0.1$.

We can also optimize directly particles in the latent space and we show it on CelebA on Figure 25 for $\tau = 0.1$ and for 10 steps.

**In the Original Space.** We report here the results obtained by running the SW-JKO scheme in the original spaces of images, which are very high dimensional ($d = 784$ for MNIST and FashionMNIST). We obtained worse results than in the latent space. Notice that we used only 1000 projections, and ran it for 50 outer iterations with a step size of $\tau = 5$.

On Figure 26, we use a RealNVP and add a uniform dequantization (Ho et al., 2019) and learned it in the logit space as it is done in (Dinh et al., 2016; Papamakarios et al., 2017) because using normalizing flows need continuous data. The RealNVP is composed here of 2 coupling layers with FCNN also composed of 2 layers and 512 hidden units. We choose $\tau = 5$, and run it for 20 steps. At each inner optimization problem, the neural networks are optimized with an Adam optimizer with a learning rate of $10^{-2}$ and for 500 epochs.
<figure><img src="image_25.png" /><figcaption>Figure 24: Generated sample obtained through a pretrained decoder + RealNVP.</figcaption></figure>

<figure><img src="image_26.png" /><figcaption>Figure 25: Particle through SW-JKO with $\tau = 0.1$ for 10 steps.</figcaption></figure>

We also report results obtained using a convolutional neural network (CNN) in Figure 27. The idea here is that we can capture inductive bias, as it is well known that CNNs are efficient for image-related tasks. We obtained a slightly better FID on MNIST, but worse results on FashionMNIST. In terms of quality of image, the generated samples do not seem better.

For the CNN, we choose a latent space of dimension 100, and we first apply a linear layer into a size of 128 × 7 × 7. Then we apply 3 convolutions layers of (kernel_size, stride, padding) being respectively (4, 2, 1), (4, 2, 1), (3, 1, 1). All layers are followed by a leaky ReLU activation, and a sigmoid is applied on the output.
<figure><img src="image_25.png" /><figcaption>Figure 24: Generated sample obtained through a pretrained decoder + RealNVP.</figcaption></figure>

<figure><img src="image_26.png" /><figcaption>Figure 25: Particle through SW-JKO with $\tau = 0.1$ for 10 steps.</figcaption></figure>

We also report results obtained using a convolutional neural network (CNN) in Figure 27. The idea here is that we can capture inductive bias, as it is well known that CNNs are efficient for image-related tasks. We obtained a slightly better FID on MNIST, but worse results on FashionMNIST. In terms of quality of image, the generated samples do not seem better.

For the CNN, we choose a latent space of dimension 100, and we first apply a linear layer into a size of 128 × 7 × 7. Then we apply 3 convolutions layers of (kernel_size, stride, padding) being respectively (4, 2, 1), (4, 2, 1), (3, 1, 1). All layers are followed by a leaky ReLU activation, and a sigmoid is applied on the output.
<figure><img src="image_27.png" /><figcaption>Figure 26: Generated sample obtained in the original space with RealNVP.</figcaption></figure>

<figure><img src="image_28.png" /><figcaption>Figure 27: Generated sample obtained in the original space with CNN.</figcaption></figure>

Nevertheless, samples obtained in the space of image look better than those obtained through the particle
scheme induced in (Liutkus et al., 2019).
<figure><img src="image_27.png" /><figcaption>Figure 26: Generated sample obtained in the original space with RealNVP.</figcaption></figure>

<figure><img src="image_28.png" /><figcaption>Figure 27: Generated sample obtained in the original space with CNN.</figcaption></figure>

Nevertheless, samples obtained in the space of image look better than those obtained through the particle
scheme induced in (Liutkus et al., 2019).