

Abstract

Autoencoders, or nonlinear factor models parameterized by neural networks, have become an indispensable tool for generative modeling and representation learning in high dimensions. Imposing structural constraints such as conditional independence on the latent variables (representation, or factors) in order to capture invariance or fairness with autoencoders has been attempted through adding ad hoc penalties to the loss function mostly in the variational autoencoder (VAE) context, often based on heuristic arguments. In this paper, we demonstrate that Wasserstein autoencoders (WAEs) are highly flexible in embracing structural constraints. Well-known extensions of VAEs for this purpose are gracefully handled within the framework of the seminal result by Tolstikhin et al. (2018). In particular, given a conditional independence structure of the generative model (decoder), corresponding encoder structure and penalties are induced from the functional constraints that define the WAE. This property of WAEs opens up a principled way of penalizing autoencoders to impose structural constraints. Utilizing this generative model structure, we present results on fair representation and conditional generation tasks, and compare them with other preceding methods.

1. INTRODUCTION

The ability to learn informative representation of data with minimal supervision is a key challenge in machine learning (Tschannen et al., 2018) , toward obtaining which autoencoders have become an indispensable toolkit. An autoencoder consists of the encoder, which maps the input to a lowdimensional representation, and the decoder, that maps a representation back to a reconstruction of the input. Thus an autoencoder can be considered a nonlinear factor analysis model as the latent variable provided by the encoder carries the meaning of "representation" and the decoder can be used for generative modeling of the input data distribution. Most autoencoders can be formulated as minimizing some "distance" between the distribution P X of input random variable X and the distribution g ♯ P Z of the reconstruction G = g(Z), where Z is the latent variable or representation having distribution P Z and g is either deterministic or probabilistic decoder (in the latter case g is read as the conditional distribution of G given Z), which is variationally described in terms of an encoder Q Z|X . For instance, the variational autoencoder (VAE, Kingma & Welling, 2014) minimizes D VAE (P X , g ♯ P Z ) = inf Q Z|X ∈Q E P X [D KL (Q Z|X ∥P Z ) -E Q Z|X log g(Z)] (1) over the set of probabilistic decoders or conditional densities g of G given Z, where D KL is the Kullback-Leibler (KL) divergence, and the Wasserstein autoencoder (WAE, Tolstikhin et al., 2018) minimizes D WAE (P X , g ♯ P Z ) = inf Q Z|X ∈Q E P X E Q Z|X d p (X, g(Z)) over the set of deterministic decoders g, where d is the metric in the space of input X and p ≥ 1. Set Q restricts the search space for the encoder. In VAEs, a popular choice is a class of normal distributions Q = {Q Z|X regular conditional distribution : Z|{X = x} ∼ N (µ(x), Σ(x)), (µ, Σ) ∈ N N } where N N is a class of functions parametrized by neural networks. In WAEs, the choice Q = {Q Z|X regular conditional distribution : Q Z ≜ E P X Q Z|X = P Z } makes the left-hand side of Eq. ( 2) equal to the (p-th power of) the p-Wasserstein distance between P X and g ♯ P Z (Tolstikhin et al., 2018, Theorem 1); Q Z is called an aggregate posterior of Z. If Q is 1

