Doubly Stochastic Adversarial Autoencoders

• The paper introduces a doubly stochastic mechanism that replaces a fixed discriminator with random feature sampling to smooth gradients and mitigate mode collapse.
• The methodology combines reconstruction loss with an adversarial penalty derived from stochastic random feature maps, bridging GAN and MMD approaches.
• Experimental evaluations on MNIST show that DS-AAE improves sample diversity and latent space exploration compared to traditional AAEs and MMD-AEs.

A Doubly Stochastic Adversarial Autoencoder (DS-AAE) is a probabilistic autoencoder architecture in which the conventional adversary of an Adversarial Autoencoder (AAE) is replaced by a stochastic function sampled from a space of random frequency features. This introduces a second source of algorithmic randomness, which smooths gradients, regularizes the adversarial training, and mitigates mode collapse. DS-AAE interpolates between Maximum Mean Discrepancy Autoencoders (MMD-AE) and classical AAEs/GANs depending on the choice of random feature distribution and parameterizations (Azarafrooz, 2018).

1. Motivation and Background

Variational Autoencoders (VAE) and AAEs both enforce a prescribed prior on the latent code. VAEs achieve this via a closed-form KL divergence penalty matching the latent aggregated posterior $Q_\theta(z)=\int Q_\theta(z|x)p_{\text{data}}(x)dx$ to the prior $p(z)$, but this can yield sample blurriness and under-exploration of multimodal posteriors since the KL term pushes each $q_\theta(z|x)$ towards $p(z)$ individually. AAEs replace the KL penalty with a GAN-style discriminator, yielding objectives of the form

$$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$

AAEs can generate sharper samples, but adversarial training may suffer from mode collapse: the discriminator can quickly distinguish “fake” codes, pushing the encoder to collapse to a few latent modes to fool the discriminator.

DS-AAE replaces the deterministic discriminator $D$ with a space of stochastic functions, injecting additional randomness. This mechanism

• Smooths gradients for the encoder,
• Prevents overfitting of the adversary,
• Encourages the generator to explore more latent modes,
• Reduces mode collapse.

2. Model Components and Stochastic Adversary

The DS-AAE comprises an encoder $q_\theta(z|x)$, which deterministically maps input $x$ to latent code $z$; a decoder/generator $p_\phi(x|z)$; and an imposed prior $p(z)$0 (commonly $p(z)$1).

Conventional AAEs use a single neural network discriminator. DS-AAE instead defines a function class: $p(z)$2 where $p(z)$3 is a “doubly stochastic gradient feature” constructed by sampling a random frequency $p(z)$4 from a measure $p(z)$5—often Gaussian for RBF kernels—and defining a random feature map

$p(z)$6

with $p(z)$7, for example, set to $p(z)$8. Each function $p(z)$9 utilized during training depends on a fresh random draw $q_\theta(z|x)$0, ensuring the stochasticity of the adversary.

3. Doubly Stochastic Minimax Objective

DS-AAE is formulated as a doubly stochastic saddle point problem: $q_\theta(z|x)$1 with objective

$q_\theta(z|x)$2

where the stochastic divergence is

$q_\theta(z|x)$3

Maximizing over $q_\theta(z|x)$4 gives

$q_\theta(z|x)$5

For large numbers of random features $q_\theta(z|x)$6, this converges to the MMD divergence; if $q_\theta(z|x)$7 is parameterized by a deep network and $q_\theta(z|x)$8 is degenerate, the formulation recovers the GAN divergence.

The overall loss includes:

• The usual reconstruction penalty $q_\theta(z|x)$9,
• The adversarial penalty $p(z)$0,
• A constraint $p(z)$1 or $p(z)$2 regularization on $p(z)$3.

4. Training Algorithm

Training in DS-AAE alternates between batch sampling of data and batch sampling of random features. The two sources of randomness are crucial for doubly stochastic optimization. The training loop proceeds as follows:

1. Sample a minibatch of data $p(z)$4.
2. Encode: $p(z)$5.
3. Sample prior codes $p(z)$6.
4. Sample $p(z)$7 random features $p(z)$8, compute features $p(z)$9.
5. Compute the stochastic gradient terms:
   • $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$0,
   • $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$1.
6. Update the adversary ($$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$2) by ascending the objective, including regularization.
7. Update the encoder and decoder parameters $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$3, using gradients from the reconstruction loss and the adversarial penalty.

Increasing the number of feature samples $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$4 or using control variates mitigates variance introduced by random feature sampling.

5. Theoretical Properties

The convergence of the algorithm leverages results from stochastic optimization in Reproducing Kernel Hilbert Spaces (RKHS) [6], demonstrating that, when the adversary step size $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$5 is small, iterates $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$6 remain in the RKHS defined by $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$7. The minimax optimization thus converges to a stationary Nash point.

Introducing randomness via $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$8 prevents the adversary from perfectly overfitting to the current generator, which compels the generator to explore the latent space more broadly, encouraging the discovery of additional latent modes and supporting a more uniform coverage of $$\min_{\theta,\phi} \mathbb{E}_x[\ell_{\text{recon}}(x;\theta,\phi)] + \lambda \cdot \sup_D \Big[ \mathbb{E}_{z\sim p(z)}\log D(z) + \mathbb{E}_x \log(1-D(q_\theta(x))) \Big]$$9.

6. Experimental Evaluation

Experiments were conducted primarily on MNIST (28×28), with additional preliminary results on CIFAR-10. Network architectures employed three fully-connected layers (1024→512→216), with ReLU activations and a final sigmoid for the decoder. The latent dimension was set to 6 for DS-AAE, with 20% input dropout and an Adam optimizer (learning rate 0.001). Minibatch and random feature batch sizes were both set to 1000 (RBF, $D$0).

Parzen-window log-likelihoods for 10,000 MNIST samples are summarized below:

Model	Parzen LL (mean ± std)
GAN [3]	225 ± 2
GMMN+AE [5]	282 ± 2
AAE [1]	340 ± 2
MMD-AE [5]	228 ± 1.6
DS-AAE	243.2 ± 1.7

DS-AAE samples display higher visual diversity, with more heterogeneous digit styles across generated panels compared to the relatively homogeneous samples produced by standard AAE or MMD-AE. Latent space interpolations are sharp and cover multiple Gaussian modes, indicating a reduction in mode collapse and improved coverage of the latent manifold.

7. Extensions and Open Questions

DS-AAE demonstrates that replacing a fixed discriminator with a stochastic function space regularizes minimax training and promotes exploration, yielding loosely a continuous spectrum between GAN-based and MMD-based regularization. The doubly stochastic scheme (data plus random features) is integral to this effect.

However, DS-AAE is sensitive to the size of data and feature batches; using small batch sizes diminishes exploration. Its convergence may be slower than that of deterministic AAEs due to the additional randomness. Directions for further research include convolutional DS-AAEs, adaptive random feature sampling strategies, improved variance reduction, and establishing generalization bounds for the doubly stochastic adversary (Azarafrooz, 2018).