Sliced Wasserstein Autoencoders (SWAE)

• Sliced Wasserstein Autoencoders (SWAE) are generative models that use the sliced Wasserstein distance to align encoded data with a chosen latent prior.
• They employ Monte Carlo estimation and sorted one-dimensional projections to compute efficient, differentiable approximations of the Wasserstein distance.
• Empirical results on image benchmarks like MNIST and CelebA demonstrate competitive reconstruction quality and realistic sample generation without adversarial training.

Sliced Wasserstein Autoencoders (SWAE) are generative autoencoder models that utilize the sliced Wasserstein distance (SWD) as a regularization mechanism to align the distribution of encoded data in latent space with a target prior distribution. By leveraging the computational tractability of one-dimensional projections in optimal transport, SWAEs enable flexible latent matching and efficient sampling-based estimation of distributional discrepancies. The core innovation is the replacement of classical regularizers such as adversarial losses or $f$-divergences (e.g., KL divergence) with the SWD, facilitating efficient training, simple implementation, and the ability to impose arbitrary sample-based priors in latent space (Kolouri et al., 2018).

1. Mathematical Formulation and Loss

Let $x \in \mathcal{X}$ represent data sampled from $p_{\rm data}$, with encoder $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$ (parameters $\phi$) and decoder $\psi_\theta: \mathcal{Z} \to \mathcal{X}$ (parameters $\theta$). Denoting the encoded data distribution $p_Z$ and a target latent prior $q_Z$, the SWAE objective function is

$$L(\theta, \phi) = \mathbb{E}_{x \sim p_{\rm data}}[c(x, \psi_\theta(\varphi_\phi(x)))] + \lambda \, SW_2(p_Z, q_Z)$$

where $x \in \mathcal{X}$0 is the pointwise reconstruction cost (often squared $x \in \mathcal{X}$1), $x \in \mathcal{X}$2 is a regularization parameter, and $x \in \mathcal{X}$3 is the sliced Wasserstein distance of order 2 (Kolouri et al., 2018).

The $x \in \mathcal{X}$4-Wasserstein distance between distributions $x \in \mathcal{X}$5 and $x \in \mathcal{X}$6 on $x \in \mathcal{X}$7 is

$x \in \mathcal{X}$8

The sliced Wasserstein distance is defined as

$x \in \mathcal{X}$9

where $p_{\rm data}$0 is the Radon transform (projected 1D marginal) of $p_{\rm data}$1 along direction $p_{\rm data}$2.

2. Efficient Slicing and Monte Carlo Estimation

The critical computational advantage of SWAE originates from the closed-form solution available for 1D Wasserstein distance. In practical terms, $p_{\rm data}$3 is approximated via Monte Carlo integration by sampling $p_{\rm data}$4 random directions $p_{\rm data}$5 uniformly from $p_{\rm data}$6:

$p_{\rm data}$7

For each slice,

• Project latent codes and prior samples onto $p_{\rm data}$8;
• Sort both sets, then compute the $p_{\rm data}$9-Wasserstein cost via matched ranks:

$\varphi_\phi: \mathcal{X} \to \mathcal{Z}$0

This formulation enables the SWD penalty to be differentiable almost everywhere and amenable to standard stochastic optimization (Kolouri et al., 2018, Knop et al., 2018, Wu et al., 2017).

3. Autoencoder Architecture, Prior Matching, and Variants

SWAEs utilize standard encoder-decoder frameworks, where both networks are typically deterministic. The encoder is often a convolutional network for image data, mapping into a latent code (e.g., $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$1-dimensional), while the decoder reconstructs data from latent codes.

A key design aspect is the flexibility of the target prior $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$2:

• SWAE can match any samplable latent prior—Gaussian, uniform, discrete mixture, or structured ring—enabling latent-distribution "shaping" beyond the typical Gaussian imposed in VAEs (Kolouri et al., 2018).
• Extended frameworks such as Encoded Prior SWAE (EPSWAE) incorporate a learnable prior-encoder, mapping a simple input distribution via a neural network to a flexible prior, which can preserve topological and geometric features of data manifolds (Krishnagopal et al., 2020).

Several implementation and architectural enhancements are introduced:

• Stacked SWD blocks with parameterized orthogonal projections (trained on the Stiefel manifold) can efficiently transport code distributions to the target prior in higher dimensions (Wu et al., 2017).
• No adversarial or kernel machinery is required, distinguishing SWAE from WAE-GAN and WAE-MMD.

4. Empirical Performance and Benchmarking

SWAEs have been evaluated on image generation benchmarks such as MNIST, CIFAR-10, and CelebA, demonstrating the following:

• On CelebA (64×64), typical Fréchet Inception Distance (FID): SWAE $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$3 48.9, outperforming VAE (66.8) and WAE-MMD (59.1), and comparable to AAE/WAE-GAN (49.3) (Wu et al., 2017).
• Reconstruction MSE is competitive with other generative autoencoding methods (Kolouri et al., 2018).
• Visual analysis reveals that SWAE produces sharp, realistic samples, smooth latent-space interpolations, and the imposed prior geometry is accurately reflected in the latent code scatterplots.
• Ablation shows performance saturates at $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$4 SWD blocks and is robust to histogram bin choices in the primal SWD implementation (Wu et al., 2017).

Cramer–Wold Autoencoders (CWAE) replace the sampled SWD penalty by a Cramer–Wold metric that admits a smooth, deterministic closed form, improving optimization stability and reducing sampling variance at the cost of $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$5 complexity for batch size $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$6 (Knop et al., 2018).

5. Theoretical Properties and Regularization Effects

The use of SWD as a regularizer has several notable effects:

• It enforces global distributional alignment in latent space but, unlike per-sample $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$7-divergence regularization (e.g., KL), does not suppress information-theoretic mutual information $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$8 between data $\varphi_\phi: \mathcal{X} \to \mathcal{Z}$9 and latent $\phi$0 (Adhya et al., 16 Jul 2025).
• The Cramér–Wold theorem underpins the identifiability guarantee for sliced approaches, as two probability measures are equal if and only if their 1D projections coincide for all directions (Knop et al., 2018).
• Sliced penalties avoid the degeneracies of kernel-based moment matching (MMD) and adversarial loss instability.
• For hyperspherical latent spaces (e.g., text topic modeling), the Spherical Sliced Wasserstein Autoencoder (S2WTM) uses a SSW penalty to align the aggregated posterior to a prior (von Mises–Fisher or uniform on the sphere), demonstrating improved mutual information and avoiding KL-driven posterior collapse (Adhya et al., 16 Jul 2025).

6. Algorithmic and Computational Aspects

A typical SWAE training loop comprises:

1. Sample a mini-batch of data $\phi$1 and encode to latent codes $\phi$2;
2. Sample a mini-batch from the latent prior $\phi$3;
3. For each of $\phi$4 directions $\phi$5:
   • Project both sets,
   • Sort projections,
   • Compute 1D $\phi$6-Wasserstein loss,
   • Accumulate SWD penalty;
4. Compute reconstruction loss,
5. Update encoder and decoder parameters via backpropagation (Kolouri et al., 2018, Wu et al., 2017).

SWAE complexity scales as $\phi$7 per batch (with $\phi$8 samples, $\phi$9 directions). For moderate batch sizes and $\psi_\theta: \mathcal{Z} \to \mathcal{X}$0, this is computationally efficient. CWAE, with its closed-form penalty, is $\psi_\theta: \mathcal{Z} \to \mathcal{X}$1 but often faster per batch for $\psi_\theta: \mathcal{Z} \to \mathcal{X}$2 (Knop et al., 2018).

7. Applications, Extensions, and Limitations

SWAEs have been successfully applied in image generative modeling, anomaly detection (including hyperspectral data, where sliced projections facilitate anomaly scoring in high dimensions) (Chen et al., 2021), multidimensional manifold learning, and topic modeling on hyperspheres (Adhya et al., 16 Jul 2025).

Limitations include:

• Monte Carlo estimation requires sufficiently many projections to capture data geometry in high-dimensional space; insufficient slices can lead to incomplete matching (Kolouri et al., 2018).
• The sorting operation can become computationally expensive for very large batches or high numbers of projections.
• No analytic density for generated data; only a generator for sampling is available.
• CWAE offers stable and deterministic training but incurs quadratic batch complexity (Knop et al., 2018).

Potential research directions include adaptive or learned slicing directions, hierarchical slicing for multi-scale regularization, combining with normalizing flows to enable tractable density evaluation, and generalizations to non-Euclidean latent manifolds or data modalities (Kolouri et al., 2018, Krishnagopal et al., 2020, Adhya et al., 16 Jul 2025).