Statistical Regeneration Guarantees of the Wasserstein Autoencoder with Latent Space Consistency

1 Introduction . Unsupervised data generation has been a time-hallowed problem in machine learning theory . Though not its strongest suit , Variational Autoencoders ( VAE ) [ 1 ] have been implemented for generative purposes to a great extent . Once emerged as a crucial unsupervised method for representation learning , its poor performances in this particular area somewhat depreciate its appeal . On the contrary , generative adversarial networks ( GANs ) [ 2 ] have proven to be superior when it comes to the quality of the generated image samples . The theoretical elegance of the adversarial training method can be duly credited for this achievement . In an attempt to adopt this merit , such ‘ adversarial loss ’ functions were incorporated into VAE objectives [ 3 , 4 ] . The Wasserstein metric , also being used in GANs [ 5 ] previously , was an inevitable choice , since it brings along the beckoning of Transportation Theory . Using the same to measure the reconstruction loss in a VAE , Wasserstein Autoencoder ( WAE ) [ 6 ] was conceived . By approaching generative modeling from an Optimal Transport ( OT ) point of view , it not only achieved remarkable reconstruction standards but also strengthened the theoretical backbone . An appropriate WAE framework consists of two components , namely an encoder ( E ) and a decoder ( D ) , which are represented by neural networks in practice . Samples from an input distribution ( PX ) are fed into the encoder , which aims at resulting in a target law ( PZ ) in the latent space ( Z ) . The decoder deals with the reconstruction of the original distribution from this ‘ encoded ’ signal . Mathematically , the objective of WAEs is to minimize L = R ( D , E ) + Ω ( E ) over the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . set of probabilistic encoders E. Here , R ( ) ensures the reconstruction of samples put in whereas , Ω ( ) ensures the attainment of the latent distribution . Needless to say , R ( ) is represented by the Wasserstein discrepancy . As the name suggests , our primary interest lies in investigating whether a suitably characterized WAE can learn a distribution , which represents its latent behaviour , without losing the memory of the initial signal it was fed . When it comes to the question of memorizing the input distribution in the form of reconstruction , theoretical guarantees in the context of GANs , have been provided rigorously . In the absence of a flexible discriminator , WAEs need more care and closer observation . Moreover , the concurrent task of sculpting a desired latent distribution adds nuances to its mechanism . Most of the existing literature on deep generative models , under adversarial loss , misses out on this very point . Against such a backdrop , this first of its kind study substantiates the simultaneous performances of WAE , from a theoretical viewpoint . Our analysis starts with a special family of WAEs , namely f -Wasserstein Autoencoders ( f -WAE ) [ 7 ] , which arises due to the specific choice of Ω ( ) as f -divergences ( 1 ) . We follow a non-parametric approach in defining both the networks involved and the distributions under consideration . The transformations induced by the encoder and the decoder are conceived as measurable functions between spaces . The nature of an ‘ ideal ’ encoder in our language , is a network that preserves information . Whereas , the decoder is characterized by a map that propagates the information retained , precisely . Our statistical assumptions on the same reflect these very ideas . As a result , the two concurrent objectives of f -WAE boil down to non-parametric estimation of laws , under vastly different circumstances . Our major findings can be summarized in the following way , • Given a non-negative margin of error , f -WAEs can learn latent space distributions with a convergence rate of O ( n− 12 ) , under fairly gentle assumptions on the input and target law ( 1 ) . The proof involves a realistic characterization of the transformation , induced by the encoder network . • Under a geometry preserving representation of the decoder map , Hölder densities can be successfully reconstructed by f -WAEs up to a constant margin ( 2 ) . Our work also produces deterministic concentration bounds on the regeneration error and hence , the rate of convergence of the empirical laws . • A particular instance of our analysis goes on to show that the input distribution converges strongly to its latent counterpart under encoding , in the form of empirical measures . Also , f -WAEs can exactly reproduce samples from smooth densities in the sense of weak convergence ( 4 ) . • Inferences , identical to that in case of f -WAE , can be drawn for the original WAE [ 6 ] , under additional assumptions on the decoder ( 2 ) . 2 Related Work . In recent times , scrutiny of the empirical achievements of deep generative models from a theoretical point of view has gained great momentum . GANs [ 2 ] , being at the helm , have received most of the attention [ 8 , 9 ] . Possibly , the cornerstone of success for a GAN lies in its adversarial training and the accompanying loss functions . A plethora of GAN variants have been devised by incorporating different losses , such as Wasserstein distance ( WD ) [ 5 ] , f -divergence [ 10 ] , Maximum Mean Discrepancy [ 11 ] and Sobolev distances [ 12 ] , to name a few . The generative aspect , in particular , has attracted a lot of research in the form of distributional estimation studies [ 13 , 14 ] . Additionally , [ 15 ] has proved weak convergence of distributions under several of the aforementioned adversarial losses . In a recent study , [ 16 ] has also discussed the convergence of the empirical measures in the case of WGAN , under parametrized network assumptions . VAEs , despite being used vastly as generative tools , have not experienced such a level of independent theoretical investigation , which comprehends its latent space mechanism . In a novel approach to unify the deep generative techniques , Adversarial AEs [ 3 ] were introduced . Subsequent work to establish resemblance [ 17 , 4 ] , has only strengthened that foundation . Eventually , the formulation of WAE [ 6 ] bridged the gap significantly by matching regeneration standards . The equivalence in performance is a direct reflection of the similarity of its objective to that of GANs ’ [ 7 ] . This ‘ primal-dual relationship ’ prepares a basis to explore the generative facets of WAE , in particular , now with the collective experience provided by GANs . Typically , WD serves as the measure of discrepancy used to represent the reconstruction loss in WAEs . Non-parametric density estimation under this metric has seen the establishment of minimax bounds based on covering numbers , under fairly general assumptions [ 18 , 19 ] . The same has also been achieved by [ 20 ] , for unbounded metric spaces , assuming that some moments are finite . Singh et al . [ 21 ] established minimax convergence rates in the same context , under smoothness assumptions on the networks employed in the generation . Despite being influential in their own ways , none of the above aim at discussing the coinciding objectives of WAEs . In a recent empirical study , Diffusion VAE [ 22 ] was introduced , which attempts at capturing arbitrary latent manifolds using kernels of Brownian motion . However , without precise conceptual insight . As such , the theoretical questions regarding the parallel machinery of WAEs remain largely unanswered . Our analysis is a humble attempt at filling this void of statistical understanding . 3 Preliminaries . 3.1 Notations . We start with denoting the input space by X , which can be assumed to be Polish ( i.e. , separable and completely metrizable ) . We will use Z for the latent space . Both may be thought of as Euclidean spaces , since it is a typical characterization . The space of probability measures ( distributions ) defined on X is denoted by P ( X ) . The set of all measurable functions from X to Z is refered to as F ( X , Z ) . For example , observe that a probabilistic encoder E belongs to F ( X , P ( Z ) ) . For any function T : X → Z , we denote the pushforward measure of µ as T # µ ∈ P ( Z ) , i.e . for any measurable set ω ⊂ Z we have T # µ ( ω ) = µ ( { x ∈ X : T ( x ) ∈ Z } ) = µ ( T−1 ( ω ) ) . To denote absolute continuity of a measure µ with respect to λ , we use the notation µ λ . The Support of µ is denoted by Supp ( µ ) . To indicate the Closure of X we use X̄ and Int ( X ) denotes the set of its Interior Points . Given S ⊆ X , the S-distance between a pair of functions f and g is given by , ‖f − g‖S = supx∈S ∣∣f ( x ) − g ( x ) ∣∣ . When the sup is taken over the whole space , we call it the uniform norm . The total variation distance between probability measures µ and γ ∈P ( X ) is denoted by‖µ− γ‖TV . Now we introduce some useful concepts that are made use of later . Definition 1 ( f -Divergence ) . Let P , Q ∈ P ( X ) such that P Q . For a convex function f with f ( 1 ) = 0 , the f-divergence of P from Q is defined as Df ( P , Q ) = ∫ X f ( dP dQ ) dQ . Definition 2 ( Yatracos Family [ 23 ] ) . For a class of real-valued functions F , defined on Z , the Yatracos family associated to F is a collection of subsets of Z defined as Y ( F ) = { x ∈ Z : f ( x ) ≥ g ( x ) ; f , g ∈ F } . Definition 3 ( Empirical Yatracos Minimizer [ 24 ] ) . The Empirical Yatracos Minimizer is a function LP ( Z ) : ∞⋃ n=1 Zn −→P ( Z ) defind as LP ( Z ) n = argmin q∈P ( Z ) ‖q − p̂n‖Y ( P ( Z ) ) , where p̂n is an empirical distribution based on a sequence of observations { Zi } ni=1 from Z . Definition 4 ( -Closeness ) . A pair of distributions f , g ∈P ( X ) are said to be -close to each other with respect to the metric c if c ( f , g ) ≤ . Definition 5 ( Quasi-Isometric Embedding ) . A map f : ( Z , d1 ) −→ ( X , d2 ) is called a quasiisometric embedding if there exists a constant A ≥ 1 such that , 1 A d1 ( x , y ) ≤ d2 ( f ( x ) , f ( y ) ) ≤ Ad1 ( x , y ) ∀ x , y ∈ Z . Definition 6 ( Covering Number [ 18 ] ) . Let S ⊆ X . Then , the minimum integer m such that there existsm closed balls { Bi } mi=1 of diameter each 3 S ⊆ ⋃ 1≤i≤mBi , is called the -covering number of S. We denote it by N ( S ) .

Summary: This paper analyzes the statistical properties of WAE, providing theoretical guarantees for both sample reconstruction and latent distribution tasks. To do so, the authors use the family of so-called f-Wasserstein autoencoders due to the specific choice of discrepancy selected for measuring the target and latent distribution. Here total variation distance is chosen. After setting up necessary preliminaries, and assumptions the authors prove the first of two main theorems, consistency in estimating the latent space distribution. Under suitable conditions, the second theorem provides a guarantee regarding the regeneration. Since they analysis was conducted using TV, which upper bounds the Jenson-Shannon divergence, the final corollary transfers the regeneration guarantees to WAEs. Fairly gentle assumptions and interesting theoretical results are discussed in detail.