One-element Batch Training by Moving Window

TL;DR

This paper introduces a novel training method for deep generative models that enables effective learning using only one-element mini-batches by splitting data into historical and current parts, reducing memory usage and allowing higher resolution training.

Contribution

The paper presents a generic approach to train models with one-element batches by splitting data in latent space, facilitating higher resolution training and broadening applicability.

Findings

Enables training with one-element mini-batches.

Reduces memory requirements for training.

Allows higher resolution image training.

Abstract

Several deep models, esp. the generative, compare the samples from two distributions (e.g. WAE like AutoEncoder models, set-processing deep networks, etc) in their cost functions. Using all these methods one cannot train the model directly taking small size (in extreme -- one element) batches, due to the fact that samples are to be compared. We propose a generic approach to training such models using one-element mini-batches. The idea is based on splitting the batch in latent into parts: previous, i.e. historical, elements used for latent space distribution matching and the current ones, used both for latent distribution computation and the minimization process. Due to the smaller memory requirements, this allows to train networks on higher resolution images then in the classical approach.

Tables

Table 1. Table 1: Comparison between the moving window mini-batch training procedure in two versions ( n = 64 𝑛 64 n=64 , k = 1 𝑘 1 k=1 and n = 64 𝑛 64 n=64 , k = 32 𝑘 32 k=32 ) with classical one ( n = k = 64 𝑛 𝑘 64 n=k=64 ). We used grid search over learning rate parameter for all models choosing optimal model in respect to WAE cost, i.e. sum of reconstruction error and the logarithm of WAE distance. For CELEB A, see FID scores in Fig.  3 .

Data set	Method	Learning	CWAE	WAE	Rec.	WAE	FID
		rate	distance	distance	error	cost	score
MNIST	WAE $n=64$, $k=64$	0.01	0.026	0.009	5.495	5.503	48.491
	WAE $n=64$, $k=32$	0.005	0.027	0.009	5.426	5.436	49.218
	WAE $n=64$, $k=1$	0.0025	0.154	0.306	4.894	5.200	54.753
CIFAR10	WAE $n=64$, $k=64$	0.0025	0.019	0.023	25.248	25.271	189.711
	WAE $n=64$, $k=32$	0.005	0.017	0.019	24.831	24.851	156.142
	WAE $n=64$, $k=1$	0.001	0.083	0.192	26.177	26.369	200.684
CELEB A	WAE $n=64$, $k=64$	0.005	0.016	0.008	116.873	116.881	62.234
	WAE $n=64$, $k=32$	0.005	0.017	0.009	117.256	117.266	60.605
	WAE $n=64$, $k=1$	0.001	0.064	0.124	116.306	116.431	74.880

Table 2. Table 2: Fashion MNIST, Adam - Comparison between the moving window mini-batch training procedure in two versions ( n = 64 𝑛 64 n=64 , k = 1 𝑘 1 k=1 and n = 64 𝑛 64 n=64 , k = 32 𝑘 32 k=32 ) with classical one ( n = k = 64 𝑛 𝑘 64 n=k=64 ).

Method	Learing	CWAE	WAE	Reconstruction	WAE	FID
	rate	distance	distance	error	cost	score
CWAE $n=64$, $k=64$	0.001	0.059	0.013	9.925	5.607	44.112
CWAE $n=64$, $k=32$	0.0001	0.055	0.009	9.166	4.511	46.557
CWAE $n=64$, $k=1$	0.00001	0.067	0.017	8.646	4.576	48.993
WAE $n=64$, $k=64$	0.001	0.057	0.008	9.627	4.822	50.656
WAE $n=64$, $k=32$	0.0001	0.058	0.009	8.931	4.312	51.940
WAE $n=64$, $k=1$	0.00001	0.072	0.027	8.587	4.980	53.799
SWAE $n=64$, $k=64$	0.001	0.065	0.014	9.641	5.429	53.122
SWAE $n=64$, $k=32$	0.0001	0.071	0.022	8.831	4.969	56.618
SWAE $n=64$, $k=1$	0.00001	0.078	0.031	8.594	5.123	52.453

Table 3. Table 3: Comparison between proposed MoW mini-batch training procedure in two versions ( n = 64 𝑛 64 n=64 , k = 1 𝑘 1 k=1 and n = 64 𝑛 64 n=64 , k = 32 𝑘 32 k=32 ) with classical one ( n = k = 64 𝑛 𝑘 64 n=k=64 ). Grid search over learning rate parameter was used choosing optimal model in respect to CWAE cost (consisting of the sum of reconstruction error and the logarithm of CWAE distance). See Fig.  5 for FID scores.

Data set	Method	Learing	CWAE	WAE	Rec.	CWAE	FID
		rate	distance	distance	error	cost	score
MNIST	CWAE $n=64$, $k=64$	0.01	0.024	0.008	5.617	1.903	48.201
	CWAE $n=64$, $k=32$	0.0075	0.024	0.008	5.324	1.604	43.796
	CWAE $n=64$, $k=1$	0.0025	0.032	0.025	4.989	1.568	36.960
CIFAR10	CWAE $n=64$, $k=64$	0.001	0.010	0.006	35.921	31.341	195.161
	CWAE $n=64$, $k=32$	0.0025	0.009	0.005	37.285	32.651	159.279
	CWAE $n=64$, $k=1$	0.001	0.014	0.015	26.681	22.435	133.820
CELEB A	CWAE $n=64$, $k=64$	0.005	0.015	0.006	118.584	114.425	59.919
	CWAE $n=64$, $k=32$	0.0075	0.014	0.004	117.466	113.204	58.577
	CWAE $n=64$, $k=1$	0.00075	0.017	0.009	116.060	112.025	58.234

Table 4. Table 4: Comparison between the MoW mini-batch training procedure for ( n = 64 𝑛 64 n=64 , k = 1 𝑘 1 k=1 ) and ( n = 64 𝑛 64 n=64 , k = 32 𝑘 32 k=32 ) with classical one ( n = k = 64 𝑛 𝑘 64 n=k=64 ) for SWAE model and various data sets. We used grid search over learning rate parameter and chose the optimal model in respect to SWAE cost, i.e. a sum of reconstruction error and the logarithm of SWAE distance. See Fig.  5 for FID scores.

Data set	Method	Learing	CWAE	WAE	Rec.	WAE	FID
		rate	distance	distance	error	cost	score
MNIST	SWAE $n=64$, $k=64$	0.01	0.030	0.014	5.658	1.394	51.987
	SWAE $n=64$, $k=32$	0.001	0.031	0.016	7.291	3.164	79.401
	SWAE $n=64$, $k=1$	0.0025	0.095	0.156	4.871	3.018	43.398
CIFAR10	SWAE $n=64$, $k=64$	0.01	0.025	0.026	26.659	23.036	149.705
	SWAE $n=64$, $k=32$	0.001	0.049	0.058	25.295	22.457	228.323
	SWAE $n=64$, $k=1$	0.001	0.044	0.088	26.204	23.775	154.941
CELEB A	SWAE $n=64$, $k=64$	0.001	0.026	0.020	121.953	118.072	73.150
	SWAE $n=64$, $k=32$	0.001	0.026	0.020	118.462	114.588	77.356
	SWAE $n=64$, $k=1$	0.00075	0.042	0.066	115.683	112.970	79.518

Equations

$$\mathrm{cost}(\theta ;q_1,\dots ,q_n)=\mathcal{F}({\mathcal{E}}_{\theta }(q_1),\dots ,{\mathcal{E}}_{\theta }(q_n);V)+\sum _{i=1}^n{\mathcal{G}}_{\theta }(q_i),$$

$$\mathcal{F}(z_{\sigma (1)},\dots ,z_{\sigma (n)};V)=\mathcal{F}(z_1,\dots ,z_n;V)$$

$$z_j={\mathcal{E}}_{{\theta }_0}(x_j)\text{mbox}forj=1,\dots ,(n-k).$$

$${\mathrm{cost}}_{l+1}(\theta )={\mathcal{F}}_{\theta }(\underset{n-k}{\underbrace{z_{m-n+k+1},\dots ,z_m}},\underset{k}{\underbrace{{\mathcal{E}}_{\theta }(x_{m+1}),\dots ,{\mathcal{E}}_{\theta }(x_{m+k})}};V_{l+1})+\sum _{j=m+1}^{m+k}{\mathcal{G}}_{\theta }(x_j).$$

$${\theta }_{l+1}={\theta }_l+\eta \nabla {\mathrm{cost}}_{l+1}({\theta }_l).$$

$$z_j={\mathcal{E}}_{{\theta }_{l+1}}(x_j)\text{mbox}forj=(m+1),\dots ,(m+k).$$

$${\mathrm{cost}}_{l+1}(\theta )={\mathcal{E}}_{\theta }(q_{n-k+1}),\dots ,{\mathcal{E}}_{\theta }(q_n);V)+\sum _{j=n-k+1}^n{\mathcal{G}}_{\theta }(q_j)$$

$$C(\theta )=\mathbb{E}[\mathrm{cost}(\theta ;q_1,\dots ,q_n,V)|(q_i)\sim Q(D),V\sim \mathbb{V}].$$

$$S^{\prime }(t)=-\frac{k}{n}\nabla C(S(t)).$$

$$\underset{\eta \to 0}{\lim }{S}_{\eta }=S.$$

$$\frac{S_{\eta }(t+h)-S_{\eta }(t)}{h}=-\frac{k}{n}\nabla C(S_{\eta }(t))+o(h),$$

$$\begin{array}[]{l}\frac{\partial}{\partial\theta}\mathrm{cost}_{l+1}(\theta)=\sum\limits_{j=n-k+1}^{n}\frac{\partial\mathcal{F}_{\theta}}{\partial r_{j}}(z_{m-n+k+1},\ldots,z_{m},\mathcal{E}_{\theta}(x_{m+1}),\ldots,\mathcal{E}_{\theta}(x_{m+k});V_{l+1})\frac{\partial\mathcal{E}_{\theta}}{\partial\theta}(x_{j})\\[8.61108pt] \hskip 65.00009pt+\sum\limits_{j=m+1}^{m+k}\frac{\partial G_{\theta}}{\partial\theta}(x_{j}).\end{array}$$

$$\begin{array}[]{l}d_{n}=\sum\limits_{j=n-k+1}^{n}\frac{\partial\mathcal{F}_{\theta}}{\partial r}(\mathcal{E}_{\theta}x_{m-n+k+1},\ldots,\mathcal{E}_{\theta}x_{m+k};V_{l+1})\frac{\partial\mathcal{E}_{\theta}}{\partial\theta}(x_{j})+\sum\limits_{j=m+1}^{m+k}\frac{\partial G_{\theta}}{\partial\theta}(x_{j}),\end{array}$$

$$C^{\prime }(\theta )=\mathbb{E}\left[\sum _{i=1}^n\frac{\partial {\mathcal{F}}_{\theta }}{\partial r}({\mathcal{E}}_{\theta }{q}_1,\dots ,{\mathcal{E}}_{\theta }{q}_n;V)\frac{\partial {\mathcal{E}}_{\theta }}{\partial \theta }+\sum _{i=1}^n\frac{\partial G_{\theta }}{\partial \theta }(x_j)|(q_i)\sim Q(D),V\sim \mathbb{V}\right].$$

$$cost(X;\mathcal{E},\mathcal{D})=mse(X;\mathcal{E},\mathcal{D})+\lambda \cdot d_{WAE}(\mathcal{E}(X),Z),$$

$$cost(X;\mathcal{E},\mathcal{D})=mse(X;\mathcal{E},\mathcal{D})+\lambda \cdot d_{CWAE}(\mathcal{E}(X),N(0,I)).$$

$$cost(X;\mathcal{E},\mathcal{D})=mse(X;\mathcal{E},\mathcal{D})+\lambda \cdot \frac{1}{k}{\sum }_{i=1}^k{d}_{SWAE}(v_i^T\mathcal{E}(X),N(0,1)),$$

Full text

We propose a generic approach to training such models using one-element mini-batches. The idea is based on splitting the batch in latent into parts: previous, i.e. historical, elements used for latent space distribution matching and the current ones, used both for latent distribution computation and the minimization process. Due to the smaller memory requirements, this allows to train networks on higher resolution images then in the classical approach.

1 Introduction

In recent years a number of deep neural network models which use cloud of points/samples and invariant with respect to the permutation order function in training process was constructed. One distinctive class of such consists of Wasserstein autoencoder WAE (Tolstikhin et al.,, 2017) or Cramer-Wold autoencoder CWAE (Tabor et al.,, 2018) models. Both use elegant geometric properties of the Wasserstein (Arjovsky et al.,, 2017) and the Maximum Mean Discrepancy MMD (Tolstikhin et al.,, 2017) distances. The metrics are often used to measure the distance between two samples – reprehension of data-set and sample from prior distribution.

Distances between samples might also be used in adversarial networks. Generative moment matching network GMMN (Li et al.,, 2015) is a deep generative model differs from generative adversarial model GAN (Goodfellow et al.,, 2014) by replacing the discriminator with a two-sample test based on MMD. MMD GAN (Arbel et al.,, 2018; Binkowski et al.,, 2018) is a modification of GMMN and classical GAN by introducing adversarial kernel learning techniques, as the replacement of a fixed Gaussian kernel in the original GMMN.

Invariancy to the permutation order which works on cloud of point/samples is also used in the case of processing sets (Maziarka et al.,, 2018; Zaheer et al.,, 2017) or 3D point cloud (Qi et al.,, 2017; Zamorski et al.,, 2019). Each point is processed by a neural network and than a simple symmetric function is used to aggregate the information from all the points.

Similar approach is used in sliced generative models (Knop et al.,, 2018). In a Sliced Wasserstein Autoencoder (SWAE) (Kolouri et al.,, 2018) authors use sliced-Wasserstein distance one dimensional projections. The methodology is to take the mean of Wasserstein distances between one-dimensional projections of data-set.

All above methods have an important limitation. That is, one cannot train the model by taking small (in extreme – one element) batches, due to the fact that samples are to be compared. Contrarily, in case of adversarial training (e.g. GAN (Goodfellow et al.,, 2014), WGAN (Arjovsky et al.,, 2017)), flow-based generative models (Dinh et al.,, 2014; Kingma & Dhariwal,, 2018)) approaches training is essentially element-independent due to some additional function used: log-likelihood in flow models, discriminator in GAN, or the Lipschitz function in WGAN. In some sense, that function accumulates knowledge over previous batches.

At the same time, the impact of mini-batch size on generalization is being examined in depth. Experiments show that, for a classification problem, using small batch sizes give preferable training stability and generalization performance (Masters & Luschi,, 2018). In all experiments presented, the best results have been obtained with mini-batch sizes $m=32$ or smaller, often as small as $2$ or $4$.

In this paper we show how one can train such models using one-element mini-batches. Each batch in latent is treated as a group of previous (historical) elements and the current mini-batch. The minimization process is applied only to the current elements group, while latent space normality (or any other prior set) is based on all. Solution is based on batches constructed using moving windows and in each new batch we have $k$ new elements and $n-k$ elements from previous batch, see Fig. 1. Then, there is no need to propagate historical elements through the network as very close approximations of latent values corresponding to them are available from last iteration, provided careful training. The network is therefore trained using only $k$ elements, where it might even denote a single element. Consequently, our neural network is trained by using only $k$ elements (in extreme – a single data element).

In experiments using autoencoder generative models, we show that such learning process gives similar, often better, results to the classical approaches, while making it possible to train the network using large inputs, e.g. large pictures.

Our contributions are the following:

we construct an optimizer MoW which allows to use one-element batches during optimization in latent-models, 2. 2.

construct auto-encoder models using the defined MoW optimizer, 3. 3.

show experiments of proposed learning using high definition images.

2 Related work

Generative models introduced, both of variational and adversarially trained (e.g. (Tolstikhin et al.,, 2017; Arjovsky et al.,, 2017; Goodfellow et al.,, 2014; Kingma & Welling,, 2014; Makhzani et al.,, 2015)) have become very important, having proved to be highly highly usable in numerous fields, e.g. (Gregor et al.,, 2015; Heljakka et al.,, 2019; Isola et al.,, 2017; Dosovitskiy et al.,, 2015; Brock et al.,, 2016). The maximum mean discrepancy learning MMD algorithm, ,which we approach here,was defined and dealt in depth, for several architectures, e.g. in (Binkowski et al.,, 2018; Dziugaite et al.,, 2015; Li et al.,, 2015).

Work on high and super-fine resolutions require lots of memory to exploit a GPU efficiently, see e.g. (Kingma & Dhariwal,, 2018; Heljakka et al.,, 2018; Larsen et al.,, 2015; Ledig et al.,, 2016). Our objective is to propose a new approach to the optimization process such that this memory requirements could be handled with. While big image approaches like in (Heljakka et al.,, 2019) try to deal with large images by iterative multiplication of them learning from low to high resolutions, the amounts of memory needed are not handled with.

The motivation of methodology in this paper is to propose an optimizer approach which could handle high memory requirements enabling high GPU bandwidth. The key is to take advantage of historical data computed. Several learning approaches use this paradigm. One is, naturally, the momentum optimization extending the regular Gradient Descent optimizer. Another is the Batch Normalization (Ioffe & Szegedy,, 2015). Scale $\gamma$ and shift $\beta$ parameters are optimized using estimates of mini-batch means and variances computed during training. With tiny batch sizes, these estimates might become less accurate approximations of the true mean and variance used for testing (Ioffe,, 2017). It has been shown to significantly improve training performance and has now become a standard component of most state-of-the-art networks.

A specific example of one example per minibatch is that in Glow network (Kingma & Dhariwal,, 2018). Instead of batch normalization, authors use an activation normalisation, to perform an affine transformations of activations using scale and bias parameters per channel initialized such that the post-actnorm activations per channel have zero mean and unit variance given initial mini-batch being a data-dependent initialization (Salimans & Kingma,, 2016).

3 Moving window in latent models

We shall now present the moving window model in a general setting. Consider an $n$-element vector of elements $(q_{i})_{i=1..n}\sim{}Q(D)$, where $Q$ can be thought of as a data generator producing a queue of examples from some data set $D$. We consider a cost function for $(q_{i})$

[TABLE]

where $Z$ denotes the latent space, $\mathcal{E}:\mathbb{R}^{N}\to Z$ and $V$ is a sample vector generated by a given fixed random vector $\mathbb{V}$, while $\theta$ are weight parameters. Since the batch is sampled randomly, without loss of generality111Replacing $\mathcal{F}$ with $\tilde{\mathcal{F}}$ given by its mean over all permutations: $\tilde{\mathcal{F}}(z_{1},\ldots,z_{n};V)=\frac{1}{n!}\sum_{\sigma}\mathcal{F}(z_{\sigma(1)},\ldots,z_{\sigma(n)};V)$. we may reduce to the case when $F$ is invariant with respect to permutations

[TABLE]

where $\sigma$ is an arbitrary permutation, while $z_{i}=\mathcal{E}(q_{i})$. The above setting is correct, but to be effective with respect to memory savings, the dimension of the latent needs to be smaller then that of the input space.

Remark 1

One can observe, see next section, that typical autoencoder based generative models can be written in the above setting (1), where $\mathcal{G}$ is the reconstruction error and $\mathcal{F}$ is used to compare the sample constructed by the generator with the prior latent distribution (see the discussion in the next section).

E.g, a Wasserstein GAN(Arjovsky et al.,, 2017), sampling examples $x$ and vector $V$, using a generator $\mathcal{D}$ to decode both and compare (the $\mathcal{F}$ function), to train both generator and discriminator with $\mathcal{G}$ cost defined with its rate of correct recognition, can be just described as above. Similarly an MMD GAN (Binkowski et al.,, 2018).

We are ready to proceed with the introduction of the $\mbox{\rm{MoW}}_{k}(n)$ optimization procedure, where $n$ as before denotes the theoretical mini-batch size, as in the standard procedure, and $k$ is the real batch size, i.e. the number of examples read at each step and kept in the memory. Thus $\mbox{\rm{MoW}}_{1}(n)$ will be using in practice one-element batches. By $\eta$ we denote the learning rate.

Now let $X=(x_{i})_{i=1,\ldots,\infty}$ denote a fixed infinite random sample generated from the data distribution (in practice it comes by drawing with repetitions from the original data-set). In that sense we can identify $X$ with the dataset generator $Q(D)$ output. We also choose a random sequence $(V_{i})$ generated from some given prior $P_{\mathcal{Z}}$.

Our objective is to define inductively the sequence $(\theta_{l})$ of weight parameter vectors and of “historical” latent elements $(z_{l})\in{}Z$.

STEP 1 (initialization). As is common, we start with some (typically randomly chosen) initial weights $\theta_{0}$. Given $\theta_{0}$, we define the first $(n-k)$ elements of the sequence $(z_{j})$ using formula

[TABLE]

STEP 2 (recursive step $l\to l+1$). Suppose that we have already defined $\theta_{j}$ for $j=0,\ldots,l$ and $z_{j}$ for $j=1,\ldots,(n-k+lk)$, i.e. $k$ elements for each of the past $l$ iterations. Put $m=n-k+lk$. We consider now the modification of the cost function (1) so that all gradients on elements $(x_{j})$ for $j=1,\ldots,m$ are frozen. Namely, we put

[TABLE]

Notice that we have dropped from the original cost function the first $(n-k)$ elements of the sum $\sum_{j=m-n+k+1}^{m+k}\mathcal{G}_{\theta}(x_{j})$ as the reconstruction error shall not be minimized for the historical elements now. We minimize the above with an arbitrary standard gradient descent procedure starting at $\theta_{l}$, e.g.:

[TABLE]

We update the sequence $(z_{j})$ by putting

[TABLE]

The MoW pseudo-code is given in Algorithm 1.

In the following theorem we show that the MoW procedure gives the correct approximation of the gradient descent method.

Theorem 1

Let $C:\theta\to\mathbb{R}$ denote the expected value of the cost function, where $\theta$ denote the weight space, i.e.:

[TABLE]

Let $S:[0,T]\to\theta$ denote the exact solution of the gradient optimization process starting from $\theta_{0}$, $S(0)=\theta_{0}$:

[TABLE]

Let $S_{\eta}:[0,T]_{\eta}\to\theta$ denote the solution given by $\mbox{\rm{MoW}}(k,n)$ with step size $\eta$ with

• •

fixed random choice of sequence $(x_{j})$ from the data set generated by $Q(D)$,

• •

random sample $(V_{j})$ from the random variable,

• •

initial starting weight $\theta_{0}$: $S_{\eta}(0)=\theta_{0}$,

where $[0,T]_{\eta}:=\{k\eta\,|\,k\in\mathbb{N},k\eta\leq T\}$ is the discretization of time (with respect to step size $\eta$). Then

[TABLE]

*Proof. * We first extend $S_{\eta}$ in an affine way to the whole interval $[0,T]$. Now, for an arbitrary $t\in[0,T)$ we are going to show that

[TABLE]

for sufficiently small $\eta$. If (4) is valid, $S_{\eta}$ (for $\eta$ sufficiently close to zero) is an approximate solution to the discretization (2). Since $S_{\eta}(0)=\theta_{0}$, by the uniqueness of the solutions of differential equations, we obtain (3), and consequently the assertion of the theorem.

Let us now proceed to the proof of (4). We take $h$ small enough so that the changes in the weights are minimal, i.e. $S_{\eta}(s)\approx S_{\eta}(t)$ for $s\in[T,T+h]$. Let us consider the procedure defining $S_{\eta}$ described in the first part of the section, where we assume that $\eta\ll h$, i.e. we assume that $\eta=Kh$ for some large $K$. We investigate the iterative process defining $S_{\eta}$. Let us first observe, that (we apply the notation used in the introduction to the model):

[TABLE]

Since $\theta$ does not, almost, change during the optimization process for the time in the interval $[t,t+h]$, $\mathcal{E}_{\theta}x_{j}\approx z_{j}$, and therefore we can approximate the above by

[TABLE]

where we use the notation $\frac{\partial F_{\theta}}{\partial r}$ to denote an arbitrary $\frac{\partial F_{\theta}}{\partial r_{j}}$ (they are all equal by the assumptions). To obtain approximation of the mean derivative over the interval $[t,t+h]$ we take the mean $\frac{1}{K}\sum\nolimits_{i=n}^{n+N-1}d_{n}$.

Now the derivative of the cost function $C$ is given by

[TABLE]

One can easily observe, that since in (5) every component has only $k$ factors, while in the above formula every sum has $n$ factors, taking $K$ large we obtain (4). $\Box$

4 Experiments

In this section we empirically validate the proposed training based on moving windows222The code is available https://github.com/gmum/MoW. We use three datasets: CELEB A, CIFAR-10, and MNIST (for the Fashion MNIST experiments see B). We show the properties of our training strategy in the case of three methods: CWAE (Tabor et al.,, 2018), WAE-MMD (Tolstikhin et al.,, 2017) and SWAE (Kolouri et al.,, 2018). As we shall see, it is possible to train such models, by using batches containing only one element.

For convenience of the reader we start form short description of CWAE (Tabor et al.,, 2018), WAE-MMD (Tolstikhin et al.,, 2017) and SWAE (Kolouri et al.,, 2018). Let $X=(x_{i})_{i=1..n}\subset\mathbb{R}^{N}$ be a given data set, which can be considered as sample from true, though unknown, data distribution $P_{X}$. The basic aim of an autoencoder is to transport the data to a (typically, but not necessarily) less dimensional latent space $\mathcal{Z}=\mathbb{R}^{D}$ with reconstruction error as small as possible. Thus, we search for an encoder $\mathcal{E}\colon\mathbb{R}^{N}\to\mathcal{Z}$ and decoder $\mathcal{D}\colon\mathcal{Z}\to\mathbb{R}^{N}$ functions, which minimize some reconstruction error, e.g. $mse(X;\mathcal{E},\mathcal{D})=\frac{1}{n}\sum_{i=1}^{n}\|x_{i}-\mathcal{D}(\mathcal{E}x_{i})\|^{2}$.

An autoencoder based generative model extends AE by introducing a cost function that makes the model generative, i.e. ensures that the data transported to the latent space $\mathcal{Z}$ conforms to some given (frequently Gaussian) prior distribution $P_{\mathcal{Z}}$. A usual way to ensure it is through adding to $mse(X;\mathcal{E},\mathcal{D})$ a regularization (using appropriate $\lambda>0$) term that penalizes dissimilarity between the distribution of the encoded data $P_{\mathcal{E}(X)}$ and the prior $P_{\mathcal{Z}}$.

WAE (Tolstikhin et al.,, 2017) is a classical autoencoder model which uses the Wasserstein metric to measure the distance between two samples – latent representation of dataset and sample from prior distribution:

[TABLE]

where $Z$ is a sample form the prior distribution $P_{\mathcal{Z}}$.

Another autoencoder based model which uses a similar approach is the CWAE (Tabor et al.,, 2018). In CWAE authors use Cramer-Wold distance between latent representation and prior distribution $N(0,I)$:

[TABLE]

SWAE (Kolouri et al.,, 2018) is a modification of WAE relying on the use of sliced Wasserstein distance. It takes the mean of the Wasserstein distances between one-dimensional projections of $\mathcal{E}(X)$ and $Z$ (sample form prior $P_{\mathcal{Z}}$, usually $N(0,I)$). Note that SWAE, similarly to WAE, also needs sampling from $P_{\mathcal{Z}}$. Consequently in SWAE two types of sampling are applied: sampling over one-dimensional projections and sampling from the prior distribution

[TABLE]

for $k$ one-dimensional projections on the spaces spanned by the unit vectors $v_{i}\in\mathbb{R}^{D}$ for $i=1,\ldots,k$ and one-dimensional Wasserstein distance $d_{SWAE}$.

In the experiment we use two basic architecture types. Experiments on MNIST use a feed-forward network for both encoder and decoder, and a 20 neuron latent layer, all using ReLU activations. For CIFAR-10, and CELEB A data sets we use convolution-deconvolution architectures. Please refer to Supplementary materials, Section A for full details. Additional experiments on FASHION MNIST set are available in Supplementary materials, Section C.

In our experiments we use classical Gradient descent minimization. This is because advanced models, like Adam, in itself use history to evaluate next steps. Thus, the results would not be clear to separate the results of the MoW and the basic optimizer. On the other hand, results using Adam optimizer are presented in Supplementary materials, Section B.

We evaluated three types of batches, all with $n=64$ overall elements. First, we use classical batches $k=n$ (classical method). Second, we use batches containing only one new element $k=1$. The third model uses $k=32$.

The quality of a generative model is typically assessed by examining generated samples and by interpolating between samples in the hidden space. We present such a comparison between our strategies in WAE-MMD architecture in Fig. 2.

For each model we consider: interpolation between two random examples from the test set (leftmost in Fig. 2), reconstruction of a random example from the test set (middle column in Fig. 2), and a sample reconstructed from a random point sampled from the prior distribution (right column in Fig. 2). The experiment shows that there are no perceptual differences between batches containing one element ($n=64$, $k=1$) and classical approach ($n=k=64$).

In order to quantitatively compare between our moving window mini-batch training procedure with the classical one, we use the Fréchet Inception Distance (FID) (Heusel et al.,, 2017). WAE-MMD, CWAE and SWAE methods are used. Thus, in addition to FID, we will also report the reconstruction error and WAE distances. Results for WAE are presented in Fig. 3 and Tab. 1. The results for CWAE and SWAE w presented in Supplementary materials, Section C.

In Fig. 3 we report for CELEB A data set: the FID score, reconstruction error and WAE distances during learning process. No essential differences between batches containing one element ($n=64$, $k=1$) and classical approach ($n=64$, $k=0$) can be seen. Moreover, adding information from previous batches allows us to obtain better FID score on CIFAR 10 dataset. The method using only one new element in batches gives comparable results in grid search procedure.

Next, we show that our approach allows us to train our model on high resolution images from Flickr-Faces-HQ (FFHQ) database. FFHQ consists of 70,000 high-quality PNG images at $1024\times 1024$ resolution and contains considerable variation in terms of age, ethnicity and image background.

The quality of our generative model is presented in Fig. 4: interpolation between two random examples from the test set, reconstruction of a random example from the test set, and a sample reconstructed from a random point sampled from the prior latent distribution.

The experiment shows that it is possible to train our model on high resolution images directly. Therefore, we do not need any modification of the architecture like in progressively growing strategy (Heljakka et al.,, 2018, 2019) where the neural network architecture is increased in each iteration. In the case of Pioneer network authors trained the network progressively through each intermediate resolution until reach the target resolution ($64\times 64$, $128\times 128$, or $256\times 256$) on classical CelebA-HQ. We show that our approach can process high-resolution images directly.

5 Conclusions

Training generative networks for high-dimensional problems, e.g. images, can be cumbersome. The dimensionality has a particular impact on the GPU memory used and a given card may not be able to store more than one example. On the other hand, the training methods which compare the latent distribution $P_{\mathcal{E}(X)}$ with some prior $P_{\mathcal{Z}}$ require that several training examples are used in a mini-batch. The examples occupy a lot of memory.

On the other hand, the latent dimension $D$ is usually, except for some very specific models, much smaller than the input example size. Therefore we proposed to use a buffer of some recent latent activation vectors together with a input single example. Past activations together with the current one would be used for minimizing given distribution distance metric $d(\mathcal{E}(X),P_{\mathcal{Z}})$, while the current example would drive the reconstruction error down.

We have proposed the MoW optimizer, that can work with as little as single input example. We have shown that such procedure approximate classical mini-batch strategy. To show that the proposed approach is useful in practice, we have performed several experiments, together with those on high-dimensional FFHQ images. The results are comparable to those while training on bigger mini-batches.

Appendix A Architecture details

The following feedforward (MNIST) and convolution-deconvolution autoencoder architectures were used:

MNIST:

• input $\mathbb{R}^{28\times 28}$

• encoder

• $3\times$ fully connected ReLU layers, 200 neurons each.

• $1\times$ fully connected layer with 20 neurons.

• latent 20-dimensional

• decoder

• $2\times$ fully connected ReLU layers, 200 neurons each.

• $1\times$ fully connected layer with 200 neurons and sigmoid activation.

CelebA:

• input $\mathbb{R}^{64\times 64\times 3}$

• encoder

• $4\times$ convolution layers with $4\times 4$ filters and $2\times 2$ strides (consecutively 32, 32, 64, and 64 output channels), all ReLU activations,

• $2\times$ fully connected layers (1024 and 256 ReLU neurons)

• $1\times$ fully connected layer with 32 neurons.

• latent 32-dimensional

• decoder

• $2\times$ fully connected layers (256 and 1024 ReLU neurons),

• $3\times$ transposed-convolution layers with $4\times 4$ filters with $2\times 2$ strides (consecutively 64, 32, 32 channels) with ReLU activation,

• $1\times$ transposed-convolution $4\times 4$ with $2\times 2$ stride, 3 channels, and sigmoid activation.

CIFAR-10

• input $\mathbb{R}^{32\times 32\times 3}$

• encoder

• $4\times$ convolution layers with $2\times 2$ filters, the second one with $2\times 2$ strides, other non-strided (3, 32, 32, and 32 channels) with ReLU activation,

• $1\times$ fully connected ReLU layer with 128 neurons,

• $1\times$ fully connected layer with 64 neurons.

• latent 64-dimensional

• decoder

• $2\times$ fully connected $ReLU$ layers with $128$ and $8192$ neurons,

• $2\times$ transposed-convolution layers with $2\times 2$ filters (32 and 32 channels) and ReLU activation,

• $1\times$ a transposed convolution layer with $3\times 3$ filter and $2\times 2$ strides (32 channels) and ReLU activation,

• $1\times$ a transposed convolution layer with $2\times 2$ filter (3 channels) and sigmoid activation.

Fashion MNIST:

• input $\mathbb{R}^{28\times 28}$

• encoder

• $4\times$ convolution layers with $4\times 4$ filters, $2\times 2$ strides and consecutively 128, 256, 512, and 1024 channels followed by ReLU activation,

• $1\times$ a fully connected layer with 8 neurons.

• latent 8-dimensional

• decoder

• $1\times$ a fully connected ReLU layer with $7\times 7\times 1024$ neurons,

• $2\times$ transposed-convolution layers with $4\times 4$ filters, $2\times 2$ strides and consecutively 512 and 256 channels,

• $1\times$ a transposed-convolution layer with $4\times 4$ filter and 1 channel and $tanh$ activation.

FFHQ:

• input $\mathbb{R}^{1024\times 1024\times 3}$

• encoder

• $4\times$ convolution layers with $5\times 5$ filters, $4\times 4$ strides and consecutively 128, 256, 512, and 1024 channels followed by ReLU activation,

• $1\times$ a fully connected layer with 64 neurons.

• latent 64-dimensional

• decoder

• $1\times$ a fully connected ReLU layer with $16\times 16\times 1024$ neurons,

• $3\times$ transposed-convolution layers with $5\times 5$ filters, $4\times 4$ strides and consecutively 512, 256, and 128 channels,

• $1\times$ a transposed-convolution layer with $5\times 5$ filter and 3 channel with $tanh$ activation.

Appendix B Experiments using ADAM optimizer

In this section we empirically validate the proposed training based on moving windows. We use Fashion MNIST dataset. We show the properties of our training strategy for three methods: CWAE (Tabor et al.,, 2018), WAE-MMD (Tolstikhin et al.,, 2017) and SWAE (Kolouri et al.,, 2018). It turns out that it is possible to train these models using single-element batches.

Convolution-deconvolution architectures (see Supplementary materials, Section A) together with Adam optimizer were used. We also used grid search for a learning rate parameter. We chose the optimal model in respect to WAE (and respectively CWAE and SWAE) cost, which consist of the sum of reconstruction error and the logarithm of WAE (and respectively CWAE and SWAE) distance.

We evaluated three types of batches, see Fig. 5 and Tab. 2.

Appendix C Additional experiments on CWAE and SWAE architecture

In this section we empirically validate the proposed MoW training based on moving windows. We use three datasets: CELEB A, CIFAR-10, MNIST. We show the properties of our training strategy in the case of CWAE (Tabor et al.,, 2018) and SWAE (Kolouri et al.,, 2018).

Two basic architecture types are used. Experiments on MNIST use a feed-forward network for both encoder and decoder, and a 20 neuron latent layer, all using ReLU activations. For CIFAR-10, and CELEB A data sets we use convolution-deconvolution architectures (see Supplementary materials, Section A).

In our experiments we use classical Gradient descent minimization. We chose the optimal model in respect to CWAE (and respectively SWAE) cost, which consist of the sum of reconstruction error and the logarithm of CWAE (and respectively SWAE) distance.

We evaluated three types of batches, see Fig. 5, Tabs. 3 and 4.