Sufficient and Disentangled Representation Learning

1 INTRODUCTION . Representation learning is a fundamental problem in machine learning and artificial intelligence ( Bengio et al. , 2013 ) . Certain deep neural networks are capable of learning effective data representation automatically and achieve impressive prediction results . For example , convolutional neural networks , which can encode the basic characteristics of visual observations directly into the network architecture , is able to learn effective representations of image data ( LeCun et al. , 1989 ) . Such representations in turn can be subsequently used for constructing classifiers with outstanding performance . Convolutional neural networks learn data representation with a simple structure that captures the essential information through the convolution operator . However , in other application domains , optimizing the standard cross-entropy and least squares loss functions do not guarantee that the learned representations enjoy any desired properties ( Alain & Bengio , 2016 ) . Therefore , it is imperative to develop general principles and approaches for constructing effective representations for supervised learning . There is a growing literature on representation learning in the context deep neural network modeling . Several authors studied the internal mechanism of supervised deep learning from the perspective of information theory ( Tishby & Zaslavsky , 2015 ; Shwartz-Ziv & Tishby , 2017 ; Saxe et al. , 2019 ) , where they showed that training a deep neural network that optimizes the information bottleneck ( Tishby et al. , 2000 ) is a trade-off between the representation and prediction at each layer . To make the information bottleneck idea more practical , deep variational approximation of information bottleneck ( VIB ) is considered in Alemi et al . ( 2016 ) . Information theoretic objectives describing conditional independence such as mutual information are utilized as loss functions to train a representation-learning function , i.e. , an encoder in the unsupervised setting ( Hjelm et al. , 2018 ; Oord et al. , 2018 ; Tschannen et al. , 2019 ; Locatello et al. , 2019 ; Srinivas et al. , 2020 ) . There are several interesting extensions of variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) in the form of VAE plus a regularizer , including beta-VAE ( Higgins et al. , 2017 ) , Annealed-VAE ( Burgess et al. , 2018 ) , factor-VAE ( Kim & Mnih , 2018 ) , beta-TC-VAE ( Chen et al. , 2018 ) , DIP-VAE ( Kumar et al. , 2018 ) . The idea of using a latent variable model has also been used in adversarial auto- encoders ( AAE ) ( Makhzani et al. , 2016 ) and Wasserstein auto-encoders ( WAE ) ( Tolstikhin et al. , 2018 ) . However , these existing works focus on the unsupervised representation learning . A challenge of supervised representation learning that distinguishes it from standard supervised learning is the difficulty in formulating a clear and simple objective function . In classification , the objective is clear , which is to minimize the number of misclassifications ; in regression , a least squares criterion for model fitting error is usually used . In representation learning , the objective is different from the ultimate objective , which is typically learning a classifier or a regression function for prediction . How to establish a simple criterion for supervised presentation learning has remained an open question ( Bengio et al. , 2013 ) . We propose a sufficient and disentangled representation learning ( SDRL ) approach in the context of supervised learning . With SDRL , we seek a data representation with two characteristics : sufficiency and disentanglement . In the context of representation learning , sufficient means that a good representation should preserve all the information in the data about the supervised learning task . This is a basic requirement and a long-standing principle in statistics . This is closely related to the fundamental concept of sufficient statistics in parametric statistical models ( Fisher , 1922 ) . A sufficient representation can be naturally characterized by the conditional independence principle , which stipulates that , given the representation , the original input data does not contain any additional information about the response variable . In addition to the basic sufficiency property , the representation should have a simple statistical structure . Disentangling is based on the general notion that some latent causes underlie data generation process : although the observed data are typically high-dimensional , complex and noisy , the underlying factors are low-dimensional , independent and have a relatively simple statistical structure . There is a range of definitions of disentangling ( Higgins et al. , 2018 ; Eastwood & Williams , 2018 ; Ridgeway & Mozer , 2018 ; Do & Tran , 2020 ) . Several metrics have been proposed for the evaluation of disentangling . However , none of these definitions and metrics have been turned into empirical criterions and algorithms for learning disentangled representations . We adopt a simple definition of disentangling which defines a representation to be disentangled if its components are independent ( Achille & Soatto , 2018 ) . This definition requires the representation to be maximally disentangled in the sense that the total correlation is zero , where the total correlation is defined as the KL divergence between the joint distribution of g ( x ) and the product of the marginal distributions of its components ( Watanabe , 1960 ) . In the rest of the paper , we first discuss the motivation and the theoretical framework for learning a sufficient and disentangled representation map ( SDRM ) . This framework leads to the formulation of an objective function based on the conditional independence principle and a metric for disentanglement and invariance adopted in this work . We estimate the target SDRM based on the sample version of the objective function using deep neural networks and develop an efficient algorithm for training the SDRM . We establish an upper error bound on the measure of conditional independence and disentanglement and show that it reaches the nonparametric minimax rate under mild regularity conditions . This result provides strong statistical guarantees for the proposed method . We validate the proposed SDRL via numerical experiments and real data examples . 2 SUFFICIENT AND DISENTANGLED REPRESENTATION . Consider a pair of random vectors ( x , y ) ∈ Rp×Rq , where x is a vector of input variables and y is a vector of response variables or labels . Our goal is to find a sufficient and disentangled representation of x. Sufficiency We say that a measurable map g : Rp → Rd with d ≤ p is a sufficient representation of x if y x|g ( x ) , ( 1 ) that is , y and x are conditionally independent given g ( x ) . This condition holds if and only if the conditional distribution of y given x and that of y given g ( x ) are equal . Therefore , the information in x about y is completely encoded by g ( x ) . Such a g always exists , since if we simply take g ( x ) = x , then ( 1 ) holds trivially . This formulation is a nonparametric generalization of the basic condition in sufficient dimension reduction ( Li , 1991 ; Cook , 1998 ) , where it is assumed g ( x ) = BTx with B ∈ Rp×d belonging to the Stiefel manifold , i.e. , BTB = Id . Denote the class of sufficient representations satisfying ( 1 ) by F = { g : Rp → Rd , g measurable and satisfies y x|g ( x ) } . We refer toF as a Fisher class because of its close connection with the concept of sufficient statistics ( Fisher , 1922 ; Cook , 2007 ) . For an injective measurable transformation T : Rd → Rd and g ∈ F , T ◦ g ( x ) is also sufficient by the basic property of conditional probability . Therefore , the Fisher class F is invariant in the sense that T ◦ F = F , provided T is injective , where T ◦ F = { T ◦ g : g ∈ F } . An important class of transformations is the class of affine transformations , T ◦ g = Ag + b , whereA is a d× d nonsingular matrix and b ∈ Rd . Disentanglement We focus on the disentangled representations among those that are sufficient . Therefore , we start from the functions of the input data that are sufficient representations in the Fisher class F . For any sufficient and disentangled representation g ( x ) , let Σg = Var ( g ( x ) ) . Since the components of g ( x ) are disentangled in the sense that they are independent , Σg is a diagonal matrix , thus Σ−1/2g g ( x ) also has independent components . Therefore , we can always rescale g ( x ) such that it has identity covariance matrix . To further simplify the statistical structure of a representation g , we also require it to be rotation invariant in distribution , that is , Qg ( x ) = g ( x ) in distribution for any orthogonal matrix Q ∈ Rd×d . The Fisher class F is rotation invariant in terms of conditional independence , but not all its members are rotation invariant in distribution . By the Maxwell characterization of the Gaussian distributions ( Maxwell , 1860 ; Bartlett , 1934 ; Bryc , 1995 ; Gyenis , 2017 ) , a random vector of dimension two or more with independent components is rotation invariant in distribution if and only if it is Gaussian with zero mean and a spherical covariance matrix . Therefore , after absorbing the scaling factor , for a sufficient representation map to be disentangled and rotation invariant , it is necessarily distributed as Nd ( 0 , Id ) . Let M be the Maxwell class of functions g : Rd → Rd , where g ( x ) is disentangled and rotation invariant in distribution . By the Maxwell characterization , we can write M = { g : Rp → Rd , g ( x ) ∼ N ( 0 , Id ) } . ( 2 ) Now our problem becomes that of finding a representation in F ∩M , the intersection of the Fisher class and the Maxwell class . The first question to ask is whether such a representation exists . The following result from optimal transport theory provides an affirmative answer and guarantees that F ∩M is nonempty under mild conditions ( Brenier , 1991 ; McCann , 1995 ; Villani , 2008 ) . Lemma 2.1 . Let µ be a probability measure on Rd . Suppose it has finite second moment and is absolutely continuous with respect to the standard Gaussian measure , denoted by γd . Then it admits a unique optimal transportation map T : Rd → Rd such that T # µ = γd ≡ N ( 0 , Id ) , where T # µ denotes the pushforward distribution of µ under T . Moreover , T is injective µ-almost everywhere . Denote the law of a random vector z by µz . Lemma 2.1 implies that , for any g ∈ F with E‖g ( x ) ‖2 < ∞ and µg ( x ) absolutely continuous with respect to γd , there exists a map T ∗ transforming the distribution of g ( x ) to N ( 0 , Id ) . Therefore , R∗ : = T ∗ ◦ g ∈ F ∩M , that is , x y|R∗ ( x ) and R∗ ( x ) ∼ N ( 0 , Id ) , ( 3 ) i.e. , R∗ is a sufficient and disentangled representation map ( SDRM ) .

Summary: The authors present a new representation learning algorithm that trades off between a sufficiency condition (that is, the label should be independent of the input conditioned on the representation) and what they call a "disentangling" condition - that the representation vectors should be independent of one another and rotationally invariant. While the first condition has been used to define disentangled representations, the second is not standard. From the condition of rotational invariance, they require that the distribution over representations is isomorphic to a uniform Gaussian. They arrive at a loss with two terms, the first is a distance correlation between labels and representation, and the second is a divergence between the representation and a uniform Gaussian. In this sense, the regularization term looks quite similar to a VAE while the loss term looks quite similar to standard classification losses. The regularization is represented as a maximum over another loss, leading to a GAN-like coupled optimization problem.