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: The submission aims to provide rigorous theoretical results for the Wasserstein Autoencoder (WAE) - A model that has achieved great empirical success yet remains severely understudied. In particular, their main contribution in showing learning theoretic bounds such as complexity and convergence guarantees. As a feature of the Autoencoder set-up, they introduce the notion of statistical regeneration and go out to prove this statement for their learning setting.