Uncertainty Autoencoders: Learning Compressed Representations via Variational Information Maximization
Aditya Grover Stefano Ermon


Uncertainty Autoencoders
[

Uncertainty Autoencoders: Learning Compressed Representations via Variational Information Maximization

 Aditya Grover Stefano Ermon 

 Stanford University Stanford University  
]



Compressed sensing techniques enable efficient acquisition and recovery of sparse, high-dimensional data signals via low-dimensional  projections. In this work, we propose Uncertainty Autoencoders, a learning framework for unsupervised representation learning inspired by compressed sensing. We treat the low-dimensional projections as noisy latent representations of an autoencoder and directly learn both the acquisition (, encoding) and amortized recovery (, decoding) procedures.
Our learning objective optimizes for a tractable  variational lower bound to the mutual information between the datapoints and the latent representations.
We show how our framework provides a unified treatment to several lines of research in dimensionality reduction, compressed sensing, and generative modeling.
Empirically, we demonstrate a 32% improvement on average over competing approaches for the task of statistical compressed sensing of high-dimensional datasets.


A preliminary version titled  "Variational Compressive Sensing using Uncertainty Autoencoders" appeared at the Uncertainty in Deep Learning Workshop at UAI, 2018.


§ INTRODUCTION

The goal of unsupervised representation learning is to learn transformations of the input data which succinctly capture the statistics of an underlying data distribution <cit.>.
In this work, we propose a learning framework for unsupervised representation learning inspired by compressed sensing. 
Compressed sensing is a class of techniques used to efficiently acquire and recover high-dimensional data using compressed measurements much fewer than the data dimensionality. 
The celebrated results in compressed sensing posit that sparse, high-dimensional datapoints can be acquired using much fewer measurements (roughly logarithmic) than the data dimensionality <cit.>.
The acquisition is done using certain classes of random matrices and the recovery procedure is based on LASSO <cit.>.


The assumptions of sparsity are fairly general and can be applied “out-of-the-box" for many data modalities, , images and audio are typically sparse in the wavelet and Fourier basis respectively.
However, such assumptions ignore the statistical nature of many real-world problems.
For representation learning in particular, we have access to a training dataset from an underlying domain.
In this work, we use this data to learn the acquisition and recovery procedures, thereby sidestepping generic sparsity assumptions.
In particular, we view the compressed measurements 
as the latent representations of an uncertainty autoencoder.

An uncertainty autoencoder (UAE) parameterizes both the acquisition and recovery procedures for compressed sensing.
The learning objective for a UAE is based on the InfoMax principle, which seeks to learn encodings that maximize the mutual information between the observed datapoints and noisy representations <cit.>. 
Since the mutual information is typically intractable in high-dimensions, we instead maximize tractable variational lower bounds <cit.>.
In doing so, we introduce a parameteric decoder that is trained to recover the original datapoint via its noisy representation.
Unlike LASSO-based recovery, a parametric decoder amortizes the recovery process, which requires only a forward pass through the decoder  at test time and thus enables scalability to large datasets <cit.>.

Notably, the framework of uncertainty autoencoders unifies and extends several lines of prior research in unsupervised representation learning.
First, we show theoretically under suitable assumptions that an uncertainty autoencoder is an implicit generative model of the underlying data distribution <cit.>, , a UAE permits sampling from the learned data distribution even though it does not specify an explicit likelihood function.
Hence, it directly contrasts with variational autoencoders (VAE) which specify a likelihood function (which is intractable and approximated by a tractable evidence lower bound) <cit.>. 
Unlike a VAE, a UAE does not require specifying a prior over the latent representations and hence offsets pathologically observed scenarios that cause the latent representations to be uninformative when used with expressive decoders <cit.>.


Next, we show that an uncertainty autoencoder, under suitable assumptions, is a generalization of principal component analysis (PCA). While earlier results connecting standard autoencoders with PCA assume linear encodings and decodings <cit.>, our result surprisingly holds even for non-linear decodings.
In practice, linear encodings learned jointly with non-linear decodings based on the UAE objective vastly outperform the linear encodings obtained via PCA. 
For dimensionality reduction on the MNIST dataset, we observed an average improvement of 5.33% over PCA when the low-dimensional representations are used for classification under a wide range of settings.

We evaluate UAEs for statistical compressed sensing of high-dimensional datasets. 
On the MNIST, Omniglot, and CelebA datasets, we observe average improvements of 38%, 31%, and 28% in recovery over the closest benchmark across all measurements considered. 
Finally, we demonstrate that uncertainty autoencoders demonstrate good generalization performance across domains in experiments where the encoder/decoder trained on a source dataset are transferred over for compressed sensing of another target dataset.



§ PRELIMINARIES

We use upper case to denote probability distributions and assume they admit absolutely continuous densities on a suitable reference measure, denoted by lower case notation. We also use upper and lower case for random variables and their realizations respectively.



  
Compressed sensing (CS). Let the datapoint and measurements be denoted with multivariate random variables X ∈ℝ^n and Y ∈ℝ^m respectively.
The goal is to recover X given the measurements Y.
For the purpose of compressed sensing, we assume m < n and relate these variables through a measurement matrix W ∈ℝ^m × l and a parameterized acquisition function f_ψ: ℝ^n →ℝ^l (for any integer l>0) such that:

    y    = W f_ψ(x) + ϵ

where ϵ is the measurement noise.
If we let f_ψ(·) be the identity function (i.e., f_ψ(x) = x for all x), then we recover a standard system of underdetermined linear equations where measurements are linear combinations of the datapoint corrupted by noise. 
In all other cases, the acquisition function transforms x such that f_ψ(x) is potentially more amenable for compressed sensing. 
For instance, f_ψ(·) could specify a change of basis that encourages sparsity, e.g., a Fourier basis for audio. 
Note that we allow the codomain of the mapping f_ψ(·) to be defined on a higher or lower dimensional space (i.e., l ≠ n in general).





  
Sparse CS.
To obtain nontrivial solutions to an underdetermined system, X is assumed to be sparse in some basis B.
We are not given any additional information about X. 
The measurement matrix W is a random Gaussian matrix and the recovery is done via LASSO <cit.>.
LASSO solves for a convex ℓ_1-minimization  
problem such that the reconstruction x for any datapoint x is given as:
x̂=argmin_x‖ Bx ‖_1
+ λ‖ y- Wx ‖_2^2
where λ>0 is a tunable hyperparameter. 



  
Statistical CS. 
In statistical compressed sensing <cit.>, we are additionally given access to a set of signals 𝒟,
 such that each x∈𝒟 is assumed to be sampled i.i.d. from a data distribution Q_data. 
 Using this dataset, we learn the the measurement matrix W and the acquisition function f_ψ(·) in Eq. (<ref>). 

At test time, we directly observe the measurements y_test that are assumed to satisfy Eq. (<ref>) for a target datapoint x_test∼ Q_data(X) and the task is to provide an accurate reconstruction x̂_test. Evaluation is based on the reconstruction error between x_test and x̂_test. 
Particularly relevant to this work, we can optionally learn a recovery function g_θ: ℝ^m →ℝ^n to reconstruct X given the measurements Y. 

This amortized approach <cit.> is in contrast to standard LASSO-based decoding which solves an optimization problem for every new datapoint at test time.
If we learned the recovery function g_θ(·) during training, then x̂_test=g_θ(y_test) and the ℓ_2 error is given by ‖ x_test -  g_θ(y_test)‖_2. 
Such a recovery process requires only a function evaluation at test time and permits scaling to large datasets <cit.>.



  
Autoencoders. An autoencoder is a pair of parameterized functions (e,d) designed to encode and decode datapoints. 
For a standard autoencoder, let e: ℝ^n →ℝ^m and d: ℝ^m →ℝ^n denote the encoding and decoding functions respectively for an n-dimensional datapoint and an m-dimensional latent space. 
The learning objective minimizes the l_2 reconstruction error over a dataset 𝒟:

    min_e,d∑_x ∈𝒟‖ x-d(e(x))‖_2^2

where the encoding and decoding functions are 
typically parameterized using neural networks.



§ UNCERTAINTY AUTOENCODER








Consider a joint distribution between the signals X and the measurements Y, which factorizes as Q_ϕ(X,Y) = Q_data(X)Q_ϕ(Y| X). Here, Q_data(X) is a fixed data distribution and Q_ϕ(Y| X) is a parameterized observation model that depends on the measurement noise ϵ, as given by Eq. (<ref>). In particular, ϕ corresponds to collectively the set of measurement matrix parameters W and the acquisition function parameters ψ.  For instance, for isotropic Gaussian noise ϵ with a fixed variance σ^2, we have Q_ϕ(Y| X) = 𝒩(Wf_ψ(X), σ^2 I_m).


In an uncertainty autoencoder, we wish to learn the parameters ϕ that permit efficient and accurate recovery of a signal X using the measurements Y. In order to do so, we propose to maximize the mutual information between X and Y:

    max_ϕ      I_ϕ(X,Y) = ∫ q_ϕ(x, y) logq_ϕ(x, y)/q_data(x) q_ϕ(y)dx dy
       = H(X) - H_ϕ(X | Y)

where H denotes differential entropy.
The intuition is simple: if the measurements preserve 
maximum information about the signal, we can hope that recovery will have low reconstruction error. We formalize this intuition by noting that this objective is equivalent to maximizing the average log-posterior probability of X given Y.
In fact, in Eq. (<ref>), we can omit the term corresponding to the data entropy (since it is independent of ϕ) to get the following equivalent objective:

    max_ϕ - H_ϕ(X | Y) = 𝔼_Q_ϕ(X,Y)[log q_ϕ(x | y)].


Even though the mutual information is maximized and equals the data entropy when Y=X, the dimensionality constraints on m ≪ n, the parametric assumptions on f_ψ(·),
and the noise model prohibit learning an identity mapping. 
Note that the properties of noise ϵ such as the distributional family and sufficient statistics are externally specified.
For example, these could be specified based on properties of the measurement device for compressed sensing.
More generally for unsupervised representation learning, we treat these properties as hyperparameters tuned based on the reconstruction loss on a held-out set, or any other form of available supervision.
It is not suggested to optimize for these statistics during learning as the UAE would tend to shrink this noise to zero to maximize mutual information, thus ignoring measurement uncertainty in the context of compressed sensing and preventing generalization to out-of-distribution examples for representation learning.
The theoretical results in Section <ref> analyze the effect of noise more formally.



Estimating mutual information between arbitrary high dimensional random variables can be challenging. 
However, we can lower bound the mutual information by introducing a variational approximation to the model posterior Q_ϕ(X | Y) <cit.>. Denoting this approximation as P_θ(X | Y), we get the following lower bound:

    I_ϕ(X,Y)     ≥ H(X) + 𝔼_Q_ϕ(X,Y)[log p_θ(x| y)].

Comparing Eqs. (<ref>, <ref>, <ref>), we can see that the second term in Eq. (<ref>) approximates the intractable negative conditional entropy, -H_ϕ(X | Y) with a variational lower bound.
Optimizing this bound leads to a decoding distribution given by P_θ(X| Y) with variational parameters θ.
The bound is tight when there is no distortion during recovery, or equivalently when the decoding distribution P_θ(X| Y) matches the true posterior Q_ϕ(X | Y) (i.e., the Bayes optimal decoder). 








  
Upper bounds. We can derive upper bounds to the mutual information by introducing a variational approximation to the marginal over the measurements Q_ϕ(Y) <cit.>. Denoting the marginal approximation as R_γ(Y), we get the following upper bound:

    I_ϕ(X,Y)     ≤𝔼_Q_ϕ(X,Y)[logq_ϕ(Y| X)/r_γ(Y)]
       =𝔼_Q_data(X)[KL(Q_ϕ(Y| X), R_γ(Y))].

Unlike the lower bound in Eq. (<ref>), the upper bound above does not specify an explicit decoder and instead it can be optimized for the variational marginal parameters γ (in addition to ϕ). The upper bound is tight when the variational marginal R_γ(Y) matches the true marginal Q_ϕ(Y).



  
Hybrid objective. Given upper and lower bounds on the mutual information, we consider a natural objectives that seek to maximize the lower bounds and minimize the upper bounds in Eq. (<ref>) and Eq. (<ref>). Formally, the uncertainty autoencoder (UAE) objective is given by:

    max_ϕ, θ, γ     𝔼_Q_ϕ(X,Y)[log p_θ(X| Y) - λlogq_ϕ(Y| X)/r_γ(Y)]

where λ≥ 0 is a hyperparameter that assigns a relative importance to the two bounds. Note that we have omitted the term corresponding to the data entropy in expressing the above objective since it is fixed with respect to the parameters being optimized. 

Assuming infinite model capacity, let θ=θ^∗, ϕ=ϕ^∗, γ=γ^∗ be the parameters that globally maximize the UAE objective for any λ≥ 0 in Eq. (<ref>). Then the following hold:



  * For all X, Y such that q_ϕ(Y)≠ 0, we have:

    p_θ^∗(X| Y) = q_ϕ^∗(X | Y)


  * For all Y, we have: 

    q_ϕ^∗(Y) = r_γ^∗(Y)


Further, the optimal value of the objective is given by (1-λ)I_θ^∗(X,Y) - H(X).







  
Stochastic optimization.
Formally, the uncertainty autoencoder (UAE) objective is given by:

    max_θ, ϕ𝔼_Q_ϕ(X,Y)[log p_θ(x| y)].

In practice, the data distribution Q_data(X) is unknown and accessible only via a finite dataset 𝒟. Hence, expectations with respect to Q_data(X) and its gradients can be estimated
using Monte Carlo methods. This allows us to express the UAE objective as:

    max_θ,ϕ     ∑_x ∈𝒟𝔼_Q_ϕ(Y| x)[log p_θ(x| y) 
    ]
    :=ℒ(ϕ, θ
    ; 
    𝒟).


Tractable evaluation of the above objective is closely tied to the distributional assumptions on the noise model. 
This could be specified externally based on, , properties of the sensing device in compressed sensing.
For the typical case of an isotropic Gaussian noise model, we know that Q_ϕ(Y| X) = 𝒩(Wf_ψ(X), σ^2 I_m), which is easy-to-sample. 

While Monte Carlo gradient estimates with respect to θ can be efficiently obtained via linearity of expectation,
gradient estimation with respect to ϕ is challenging since these parameters specify the sampling distribution Q_ϕ(Y | X). 
One solution is to evaluate score function gradient estimates along with control variates <cit.>. Alternatively, many continuous distributions (e.g., the isotropic Gaussian and Laplace distributions) can be reparameterized such that it is possible to obtain samples by applying a deterministic transformation to samples from a fixed distribution and typically leads to low-variance gradient estimates <cit.>. 



§ THEORETICAL ANALYSIS


In this section, we derive connections of uncertainty autoencoders with generative modeling and Principal Component Analysis (PCA). The proofs of all theoretical results in this section are in Appendix <ref>.


 §.§ Implicit generative modeling

Starting from an arbitrary point x^(0)∈ℝ^n, define a Markov chain over X,Y with the following transitions:

    y^(t)   ∼ Q_ϕ(Y | x^(t))
    
    x^(t+1)   ∼ P_θ(X | y^(t))


Let θ^∗,ϕ^∗ denote an optimal solution to the UAE objective in Eq. (<ref>).
If there exists a ϕ such that q_ϕ(x|y)=p_θ^∗(x|y)
and the Markov chain defined in Eqs. (<ref>, <ref>) is ergodic, then the stationary distribution of the chain for the parameters ϕ^∗ and θ^∗ is given by Q_ϕ^∗(X,Y).


The above theorem suggests an interesting insight into the behavior of UAEs. Under idealized conditions, the learned model specifies an implicit generative model for Q_ϕ^∗(X,Y). Further, ergodicity can be shown to hold for the isotropic Gaussian noise model.


Let θ^∗,ϕ^∗ denote an optimal solution to the UAE objective in Eq. (<ref>).
If there exists a ϕ such that q_ϕ(x|y)=p_θ^∗(x|y) and the noise model is Gaussian, then
the stationary distribution of the chain for the parameters ϕ^∗ and θ^∗ is given by Q_ϕ^∗(X,Y).


The marginal of the joint distribution Q_ϕ(X,Y) with respect to X corresponds to the data distribution. 
A UAE hence seeks to learn an implicit generative model of the data distribution <cit.>, i.e., even though we do not have a tractable estimate for the likelihood of the model, we can generate samples using the Markov chain transitions defined in Eqs. (<ref>, <ref>). 



 §.§ Optimal encodings



A UAE can also be viewed as a dimensionality reduction technique for the dataset 𝒟.
While in general the encoding performing this reduction can be nonlinear, the case of a linear encoding 
is one 
where the projection vectors are given as the rows of the measurement matrix W. 
The result below characterizes the optimal encoding of the dataset 𝒟 with respect to the UAE objective for an isotropic Gaussian noise model.


Assume a uniform data distribution over a finite dataset 𝒟. Further, we assume that expectations in the UAE objective exist, and the signals and measurement matrices are bounded in ℓ_2/Frobenius norms, i.e., ‖ x‖_2 ≤ k_1 for all x ∈𝒟, ‖ W‖_F ≤ k_2 for some positive constants k_1, k_2 ∈ℝ^+.
For a linear encoder 
and isotropic Gaussian noise ϵ∼𝒩(0, σ^2 I), the optimal measurement matrix W^∗ 
that maximizes the mutual information for an optimal decoder 
in the limit σ→∞ 
is given as:

    W^∗ = 
    eig_m ( 
    ∑_x_i,x_j∈𝒟[ (x_i-x_j)(x_i-x_j)^T] )

where eig_m(M) denotes the top-m eigenvectors of the matrix M with the largest eigenvalues (specified up to a positive scaling constant).














Under the stated assumptions, the above result suggests an interesting connection between UAE and PCA.
PCA seeks to find the directions that explain the most variance in the data. Theorem <ref> suggests that when the noise in the projected signal is very high, the optimal projection directions (i.e., the rows of W^∗) correspond to the principal components of the data signals. We note that this observation comes with a caveat; when the noise variance is high, it will dominate the contribution to the measurements Y in Eq. (<ref>) as one would expect. Hence, the measurements and the signal will have low mutual information even under the optimal measurement matrix W^∗. 

Our assumptions are notably different from prior results in autoencoding drawing connections with PCA. Prior results show that linear encoding and decoding in a standard autoencoder recovers the principal components of the data (Eq. (3) in <cit.>, Eq. (1) in <cit.>). In contrast, Theorem <ref> is derived from variational principles and does not assume linear decoding. 

In general, the behaviors of UAE and PCA can be vastly different. As noted in prior work <cit.>, the principal components may not be the the most informative low-dimensional projections for recovering the original high-dimensional data back from its projections. A UAE, on the other hand, is explicitly designed to preserve as much information as possible (see Eq. (<ref>)). We illustrate the differences in a synthetic experiment in Figure <ref>. The true data distribution is an equiweighted mixture of two Gaussians stretched along orthogonal directions. We sample 100 points (black) from this mixture and consider two dimensionality reductions. In the first case, we project the data on the first principal component (blue points on magenta line). This axis captures a large fraction of the variance in the data but collapses data sampled from the bottom right Gaussian in a narrow region. The projections of the data on the UAE axis (red points on green line) are more spread out.
This suggest that recovery is easier, even if doing so increases the total variance in the projected space compared to PCA.










§ EXPERIMENTS




 §.§ Statistical compressed sensing

We perform compressed sensing on three datasets:  MNIST <cit.>, Omniglot <cit.>, and CelebA dataset <cit.>, with extremely low number of measurements m∈{2,5,10,25,50, 100}. 
We discuss the MNIST and Omniglot datasets here since they have a similar setup. To save space, results on the CelebA dataset are deferred to Appendix <ref>. Every image in MNIST and Omniglot has a dimensionality of 28 × 28.
In all our experiments, we assume a Gaussian noise model with σ=0.1. We evaluated UAE against:



  * LASSO decoding with random Gaussian matrices. 
The MNIST and Omniglot datasets are reasonably sparse in the canonical pixel basis, and hence, we did not observe any gains after applying
Discrete Cosine Transform
and Daubechies-1 Wavelet Transform.


  * CS-VAE. This approach to compressed sensing was proposed by <cit.> and learns a latent variable generative model over the observed variables X and the latent variables Z. 
Such a model defines a mapping G: ℝ^k →ℝ^n from Z to X, which is given by either the mean function of the observation model for a VAE or the forward deterministic mapping to generate samples for a GAN. We use VAEs in our experiments.
Thereafter, using a classic acquisition matrix satisfying a generalized Restricted Eigenvalue Condition (say W) (, random Gaussian matrices), the reconstruction x̂ for any datapoint is given as:
x̂=G(argmin_z‖ y - W G(z) ‖_2). 
Intuitively, this procedure seeks the latent vector z such that the corresponding point on the range of G can best approximate the measurements y under the mapping W. 
We used the default parameter settings and architectures proposed in <cit.>.  


  * RP-UAE. To independently evaluate the effect of variational decoding, this ablation baseline encodes the data using Gaussian random projections (RP) and trains the decoder based on the UAE objective. Since LASSO and CS-VAE both use an RP encoding, the differences in performance would arise only due to the decoding procedures.


The UAE decoder and the CS-VAE encoder/decoder are multi-layer perceptrons consisting of two hidden layers with 500 units each.
For a fair comparison with random Gaussian matrices, the UAE encoder is linear. 
Further, we perform ℓ_2 regularization on the norm of W. This helps in generalization to test signals outside the train set and is equivalent to solving the Lagrangian of a constrained UAE objective:

    max_θ, ϕ𝔼_Q_ϕ(X,Y)[log p_θ(x| y)] 
     subject to ‖ W ‖_F ≤ k.

The Lagrangian parameter is chosen by line search on the above objective. 
The constraint ensures that UAE does not learn encodings W that trivially scale the measurement matrix to overcome noise. For each m, we choose k to be the expected norm of a random Gaussian matrix of dimensions n × m for fair comparisons with other baselines. In practice, the norm of the learned W for a UAE is much smaller than those of random Gaussian matrices suggesting that the observed performance improvements are non-trivial. 

Results. The ℓ_2 reconstruction errors on the standard test sets
are shown in Figure <ref>. For both datasets, we observe that UAE drastically outperforms both LASSO and CS-VAE for all values of m considered. LASSO (blue curves) is unable to reconstruct with such few measurements. The CS-VAE (red) error decays much more slowly compared to UAE as m grows. Even the RP-UAE baseline (yellow), which trains the decoder keeping the encoding fixed to a random projection, outperforms CS-VAE. 
Jointly training the encoder and the decoder using the UAE objective (green) exhibits the best performance.  
These results are also reflected qualitatively for the reconstructed test signals shown in Figure <ref> for m=25 measurements. 




 §.§ Transfer compressed sensing


To test the generalization of the learned models to similar, unseen datasets, we consider the task of transfer compressed sensing task introduced in <cit.>.

Experimental setup. We train the models on a source domain that is related to a target domain. Since the dimensions of MNIST and Omniglot images match, transferring from one domain to another requires no additional processing.
For UAE, we consider two variants. 
In UAE-SE, we used the encodings from the source domain and retrain the decoder on the target domain. 
For UAE-SD, we use source decoder and retrain the encoder on the target domain.

Results. The ℓ_2 reconstruction errors
are shown in Figure <ref>. 
LASSO (blue curves) does not involve any learning, and hence its performance is same as Figure <ref>. 
The CS-VAE (red) performance degrades significantly in comparison, even performing worse than LASSO in some cases. 
The UAE based methods outperform these approaches and UAE-SE (green) fares better than UAE-SD (yellow). 
Qualitative differences are highlighted in Figure <ref> for 
m=25 measurements. 



 §.§ Dimensionality reduction

Dimensionality reduction is a common preprocessing technique for specifying features for classification. We compare PCA and UAE on this task. While Theorem <ref> posits that the two techniques are equivalent in the regime of high noise given optimal UAE decodings, we set the noise as a hyperparameter based on a validation set to enable out-of-sample generalization.


Setup. 
We learn the principal components and UAE projections on the MNIST training set for varying number of dimensions. We then learn classifiers based on the these projections. Again, we use a linear encoder for the UAE for a fair evaluation.
Since the inductive biases vary across different classifiers, we considered 8 commonly used classifiers: k-Nearest Neighbors (kNN), Decision Trees (DT), Random Forests (RF), Multilayer Perceptron (MLP), AdaBoost (AdaB), Gaussian Naive Bayes (NB), Quadratic Discriminant Analysis (QDA), and Support Vector Machines (SVM) with a linear kernel. 

Results. The performance of the PCA and UAE feature representations for different number of dimensions is shown in Table <ref>. We find that UAE outperforms PCA in a majority of the cases. Further, this trend is largely consistent across classifiers. The improvements are especially high when the number of dimensions is low, suggesting the benefits of UAE as a dimensionality reduction technique for classification. 



§ RELATED WORK


In this section, we contrast uncertainty autoencoders with related works in autoencoding, compressed sensing, and mutual information maximization. 



Autoencoders. To contrast uncertainty autoencoders with other commonly used autoencoding schemes, consider a UAE with a Gaussian observation model with fixed isotropic covariance for the decoder of all the autoencoding objectives we discuss subsequently. 
The UAE objective can be simplified as:

    min_θ, ϕ𝔼_x,y ∼ Q_ϕ(X,Y)[‖ x - g_θ(y)‖_2^2]


Standard Autoencoder. If we assume no measurement noise (i.e., ϵ = 0) and assume the observation model P_θ(X| Y) to be a Gaussian with mean g_θ(Y) and a fixed isotropic Σ, then the UAE objective reduces to minimizing the mean squared error between the true and recovered datapoint:

    min_θ, W, ψ𝔼_x ∼ Q_data(X)[‖ x - g_θ(Wf_ψ(x))‖_2^2]

This special case of a UAE corresponds to a standard autoencoder <cit.> where the measurements Y signify a hidden representation for X. However, this case lacks the interpretation of an implicit generative model since the assumptions of Theorem <ref> do not hold.

Denoising Autoencoders. A DAE <cit.> 
adds noise at the level of the input datapoint X to learn robust representations. For a UAE, the noise model is defined at the level of the compressed measurements. 
Again, with the assumptions of a Gaussian decoder, the DAE objective can be expressed as:

    min_θ, W, ψ𝔼_x ∼ Q_data(X), x̃∼ C(X̃| x)[‖ x - g(Wf_ψ(x̃))‖_2^2]

where C(·| X) is some predefined noise corruption model. 
Similar to Theorem <ref>, a DAE also learns an implicit model of the data distribution <cit.>.

Variational Autoencoders. A VAE <cit.> explicitly learns a latent variable model P_θ(X,Y) for the dataset. The learning objective is a variational lower bound to the marginal log-likelihood assigned by the model to the data 𝒳, which notationally  corresponds to  𝔼_ Q_data(X)[log P_θ(x)]. The variational objective that maximizes this quantity can be simplified as:

    min_θ, ϕ     𝔼_x,y ∼ Q_ϕ(X,Y)[‖ x - g_θ(y)‖_2^2] 
       + 𝔼_x ∼ Q_data[KL(Q_ϕ(Y | x), P(Y))]


The learning objective includes a reconstruction error term, akin to the UAE objective. Crucially, it also includes a regularization term to minimize the KL divergence of the variational posterior
over Y
with a prior distribution over Y.
A key difference is that a UAE does not explicitly need to model the prior distribution over Y. 
On the downside, a VAE can perform efficient ancestral sampling while a UAE requires running relatively expensive Markov Chains to obtain samples.

Recent works have attempted to unify the variants of variational autoencoders through the lens of mutual information <cit.>. These works also highlight scenarios where the VAE can learn to ignore the latent code in the presence of a strong decoder thereby affecting the reconstructions to attain a lower KL loss. 
One particular variant, the β-VAE, weighs the additional KL
regularization term with a positive factor β and can effectively learn disentangled representations <cit.>. Although <cit.> does not consider this case, the UAE can be seen as a β-VAE with β=0.




To summarize, our uncertainty autoencoding formulation provides a combination of unique desirable properties for representation learning that are absent in prior autoencoders. 
As discussed, a UAE defines an implicit generative model without specifying a prior (Theorem <ref>) even under realistic conditions (Corollary <ref>; unlike DAEs) and has rich connections with PCA even for non-linear decoders (Theorem <ref>; unlike any kind of existing autoencoder).




Generative modeling and compressed sensing. The closely related works of <cit.> also use generative models for compressed sensing. As highlighted in Section <ref>, their approach is radically different from UAE. Similar to <cit.>, a UAE learns a data distribution. However, in doing so, it additionally learns an acquisition/encoding function and a recovery/decoding function, unlike <cit.> which rely on generic random matrices and ℓ_2 decoding. The cost of implicit learning in a UAE is that some of its inference capabilities, such as likelihood evaluation and sampling, are intractable or require running Markov chains. However, these inference queries are orthogonal to compressed sensing. Finally, our decoding is amortized and scales to large datasets, unlike <cit.> which solve an independent optimization problem for each test datapoint. 




  
Mutual information maximization. The principle of mutual information maximization, often referred to as InfoMax in prior work, was first proposed for learning encodings for communication over a noisy channel <cit.>. 
The InfoMax objective has also been applied for statistical compressed sensing for learning both linear and non-linear encodings <cit.>. Our work differs from these existing frameworks in two fundamental ways. First, we optimize for a tractable variational lower bound to the MI that which allows our method to scale to high-dimensional data. Second, we learn an amortized <cit.> decoder in addition to the encoder that sidesteps expensive, per-example optimization for the test datapoints. 


Further, we improve upon the IM algorithm proposed originally for variational information maximization <cit.>. While the IM algorithm proposes to optimize the lower bound on the mutual information in alternating “wake-sleep” phases for optimizing the encoder (“wake") and decoder (“sleep") analogous to the expectation-maximization procedure used in <cit.>, we optimize the encoder and decoder jointly using a single consistent objective leveraging recent advancements in gradient based variational stochastic optimization.




§ CONCLUSION

In this work, we presented uncertainty autoencoders (UAE), a framework for unsupervised representation learning via variational maximization of mutual information between an input signal and its latent representation. 
We presented connections of our framework with many related threads of research, in particular with respect to implicit generative modeling and principal component analysis. 
 Empirically, we showed that UAEs are a natural candidate for statistical compressed sensing, wherein we can learn the acquisition and recovery functions jointly.

In the future, it would be interesting to incorporate advancements in compressed sensing based on complex neural network architectures <cit.> within the UAE framework for real world applications, e.g., medical imaging.
Unlike the rich theory surrounding the compressed sensing of sparse signals, a similar theory surrounding generative model-based priors on the signal distribution is lacking. 
Recent works have made promising progress in developing a theory of SGD based recovery methods for nonconvex inverse problems, which continues to be an exciting direction for future work <cit.>. 














§ ACKNOWLEDGEMENTS

This research was supported by NSF (#1651565, #1522054, #1733686), ONR (N00014-19-1-2145), AFOSR (FA9550-19-1-0024), and FLI.
AG is supported by a Microsoft Research Ph.D. fellowship and a Stanford Data Science scholarship.
We are thankful to Daniel Levy and Yang Song for helpful discussions on a proof and Kristy Choi, Manik Dhar, Neal Jean, and Ben Poole for helpful comments. 

unsrt



§ PROOFS OF THEORETICAL RESULTS



 §.§ Proof of Theorem <ref>



We can rewrite the UAE objective in Eq. (<ref>) as:

    𝔼_Q_ϕ(X,Y)[log p_θ(x| y) 
    ]
       = 𝔼_Q_ϕ(Y)[∫ q_ϕ(x| y)log p_θ(x| y) dx] 
       = -H_ϕ(X | Y) - 𝔼_Q_ϕ(Y)[KL(Q_ϕ(X| y)‖  P_θ(X| y))].

The KL-divergence is non-negative and minimized when its argument distributions are identical. Hence, for a fixed optimal value of θ=θ^∗, if there exists a ϕ in the space of encoders being optimized that satisfies:

    p_θ^∗(X| Y) = q_ϕ(X| Y)

for all X,Y with p_θ^∗(Y) ≠ 0, then it corresponds to the optimal encoder, ,

    ϕ=ϕ^∗.



For any value of ϕ, we know the following Gibbs chain converges to Q_ϕ(X,Y) if the chain is ergodic:

    y^(t)   ∼ Q_ϕ(Y | x^(t))
    
    x^(t+1)   ∼ Q_ϕ(X | y^(t)).

Substituting the results from Eqs. (<ref>-<ref>) in the Markov chain transitions in Eqs. (<ref>, <ref>) finishes the proof.




 §.§ Proof of Corollary <ref>


By using earlier results (Proposition 2 in <cit.>), we need to show that the Markov chain defined in Eqs. (<ref>)-(<ref>) is Φ-irreducible with a Gaussian noise model.[Note that the symbol Φ here is different from the parameters denoted by little ϕ used in the rest of the paper.] That is, there exists a measure such that there is a non-zero probability of transitioning from every set of non-zero measure to every other such set defined on the same measure using this Markov chain. 

Consider the Lebesgue measure. Formally, given any (x, y) and (x',y') such that the density q(x,y) > 0 and q(x',y')>0 for the Lebesgue measure, we need to show that the probability density of transitioning q(x',y'| x,y)>0. 


(1) Since q(y| x)>0 for all x ∈ℝ^n, y ∈ℝ^m (by Gaussian noise model assumption), we can use Eq. (<ref>) to transition from (x,y) to (x, y') with non-zero probability.  
 
(2) Next, we claim that the transition probability q(x' | y) is non-negative for all x', y. By Bayes rule, we have:

    q(x'| y) = q(y| x') q(x')/q(y).


Since q(x,y)>0 and q(x',y')>0, the marginals q(y) and q(x') are positive. Again, q(y| x')>0 for all x' ∈ℝ^n, y ∈ℝ^m by the Gaussian noise model assumption.  Hence, q(x'| y)  is positive. Finally, using the optimality assumption for the posteriors p(x'| y) matching q(x'| y) for all x', y', we can use Eq. (<ref>) to transition from (x,y') to (x',y') with non-zero probability.

From (1) and (2), we see that there is a non-zero probability of transitioning from (x,y) to (x',y'). Hence, under the assumptions of the corollary the Markov chain in Eqs. (<ref>, <ref>) is ergodic.






 §.§ Proof of Theorem <ref>













Under an optimal decoder, the model posterior P_θ(X | Y) matches the true posterior Q_ϕ(X | Y) and hence, the UAE objective can be simplified as:

    𝔼_Q_ϕ(X,Y)[log q_ϕ(x | y)]   =𝔼_Q_ϕ(X,Y)[log q_ϕ(x, y) - log q_ϕ(y)]
       =-H(X) - 𝔼_Q_data(X)[H(Y | x)] - 𝔼_Q_ϕ(X,Y)[log q_ϕ(y)].


The first term corresponds to the negative of the data entropy, is independent of ϕ and σ, and hence it can be removed. 
For the second term, note that Y| x is a normal distributed random variable and hence its entropy is given by a constant 1/2log 2π eσ^2.
Only the third term depends on ϕ. 

Removing the data entropy term since it is a constant independent of both ϕ and σ, we can define a modified objective M(W, 𝒟, σ) as:


    M(W, 𝒟, σ) := 𝔼_Q_ϕ(X,Y)[log q_ϕ(y)] + 1/2log 2π e σ^2.

As σ→∞, the optimal encodings maximizing the mutual information can be specified as:

    W^∗ =  lim_σ→∞max_W - M(W, 𝒟, σ).


We can lower-bound M(W,𝒟, σ) using Jensen's inequality:



    M(W, 𝒟, σ)
       = 𝔼_Q_ϕ(X,Y)[log𝔼_x_j ∼ Q_data(X)[q_ϕ(y | x_j )]] + 1/2log 2π e σ^2
       = 1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[log1/|𝒟|∑_x_j∈𝒟 q_ϕ(y | x_j )]+1/2log 2π e σ^2
       ≥1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[∑_x_j∈𝒟1/|𝒟|log  q_ϕ(y | x_j ) ] +1/2log 2π e σ^2
       := C(W, 𝒟, σ)

where we have used the fact that the data distribution is uniform over the entire dataset (by assumption).

Finally, we denote the non-negative slack term for the above inequality as S(W, 𝒟, σ) such that:

    M(W, 𝒟, σ)    = C(W, 𝒟, σ) + S(W, 𝒟, σ).


Overview of proof strategy: We will first simplify expressions for the lower bound C(W, 𝒟, σ) and slack term S(W, 𝒟, σ).
Then, we will show that as σ→∞, the ratio of the slack term and the lower bound converges pointwise to 0 and hence, the lower bound is arbitrarily close to M(W, 𝒟, σ) in this regime for a fixed W.
Further, we will show that the convergence is uniform in W.
Finally, we will note that the optimal encodings W^∗ for the lower bound correspond to the stated expressions for W in the proof statement.


As a first step, we consider simplifications of the lower bound and the slack term.






Lower bound:  C(W, 𝒟, σ)

    C(W, 𝒟, σ)   = 1/|𝒟|∑_x_i∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[1/|𝒟|∑_x_j∈𝒟[log q_ϕ(Wx_i + ϵ| x_j )]]+1/2log 2π e σ^2
       = 1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[log q_ϕ(Wx_i + ϵ| x_j )]+1/2log 2π e σ^2
       = -1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[ (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2 +1/2log 2 πσ^2 ]+1/2log 2π e σ^2
       = -1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[ (Wx_i-Wx_j)^T(Wx_i -Wx_j) + 2 ϵ^T(Wx_i -Wx_j) + ϵ^Tϵ/2σ^2] +  1/2
       = -1/|𝒟|^2∑_x_i, x_j∈𝒟((Wx_i-Wx_j)^T(Wx_i -Wx_j)/2σ^2 +  1/2) +  1/2
       = -1/|𝒟|^2∑_x_i, x_j∈𝒟((Wx_i-Wx_j)^T(Wx_i -Wx_j)/2σ^2).



Slack:  S(W, 𝒟, σ)


    S(W, 𝒟, σ)    = -C(W, 𝒟, σ) + M(W, 𝒟, σ)
       =-1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[1/|𝒟|∑_x_j∈𝒟log q_ϕ(y | x_j )] +1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[log1/|𝒟|∑_x_j∈𝒟 q_ϕ(y | x_j )]
       = -1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[𝔼_Q_data(X)[logq_ϕ(y , x_j )/q_data(x_j)]] +1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[log q_ϕ(y)]
       = 1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[𝔼_Q_data(X)[logq_data(x_j)/q_ϕ(x_j | y)q_ϕ(y)]] +1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[log q_ϕ(y)]
       = 1/|𝒟|∑_x_i∈𝒟𝔼_Q_ϕ(Y | X)[KL(Q_data(X),Q_ϕ(X | y))] 
       =- 1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_Q_ϕ(Y | X)[  log q_ϕ(x_j | y) + log|𝒟|]
       = - log|𝒟| -1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[log q_ϕ(x_j | W x_i + ϵ)]



We can simplify the posteriors Q_ϕ(x_j | W x_i + ϵ) as:

    Q_ϕ(x_j | W x_i + ϵ) 
    = Q_ϕ(x_j , W x_i + ϵ)/Q_ϕ(W x_i + ϵ)
       = Q_ϕ(x_j) Q_ϕ(W x_i + ϵ|  x_j)/∑_x_k ∈𝒟Q_ϕ(W x_i + ϵ| x_k)Q_ϕ(x_k)
    = exp(-(W(x_i-x_j) + ϵ)^T(W(x_i-x_j) + ϵ)/2 σ^2)/∑_x_k ∈𝒟exp(-(W(x_i-x_k) + ϵ)^T(W(x_i-x_k) + ϵ)/2 σ^2)

where we have used the fact that the data distribution is uniform and the decoder is isotropic Gaussian. 

Substituting the above expression for the slack term:

    S(W, 𝒟, σ)    = - log|𝒟| -1/|𝒟|^2∑_x_i, x_j∈𝒟𝔼_ϵ∼𝒩(0, σ^2I)[ (-(W(x_i-x_j) + ϵ)^T(W(x_i-x_j) + ϵ)/2 σ^2) 
        - log∑_x_k ∈𝒟exp(-(W(x_i-x_k) + ϵ)^T(W(x_i-x_k) + ϵ)/2 σ^2)] 
       =  1/|𝒟|^2∑_x_i, x_j∈𝒟[ (Wx_i-Wx_j)^T(Wx_i -Wx_j)/2σ^2 +  1/2
       + 𝔼_ϵ∼𝒩(0, σ^2I)(log∑_x_k ∈𝒟exp(-(W(x_i-x_k) + ϵ)^T(W(x_i-x_k) + ϵ)/2 σ^2) - log|𝒟|)_∫γ(σ, ϵ, W,𝒟, x_i)dϵ].




For any fixed x_i, W, the integrand γ(σ, ϵ, W,𝒟, x_i) in Eq. (<ref>) can be seen as a sequence of functions indexed by σ. 
We next make the claim that dominated convergence in ϵ holds for this sequence for all x_i, W.
We show so by first observing that γ(σ, ϵ, W,𝒟, x_i) converges pointwise to a constant (=0) as σ→∞ and thereafter deriving integrable upper and lower bounds for γ(σ, ϵ, W,𝒟, x_i) below that are independent of σ.

For the upper bound, we note that:

    log∑_x_j∈𝒟exp(- (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2)    ≤max_x_j∈𝒟 - (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2 + log|𝒟|
       ≤log|𝒟|.

This gives an upper bound on the integrand γ(σ, ϵ, W,𝒟, x_i):

    γ(σ, ϵ, W,𝒟, x_i)    ≤ - 1/√(2 πσ^2)exp(-ϵ^Tϵ/2σ^2)  [log|𝒟|-log|𝒟| ] 
       = 0.


For the lower bound, we note that:

    log∑_x_j∈𝒟exp(- (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2)   ≥max_x_j∈𝒟(- (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2)
       = - min_x_j∈𝒟((Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2).
.

Hence, we have the following lower bound:

    γ(σ, ϵ, W,𝒟, x_i)    ≥ - 1/√(2 πσ^2)exp(-ϵ^Tϵ/2σ^2) (min_x_j∈𝒟(Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)/2σ^2-log|𝒟|)
       ≥ -1/√(πϵ^Tϵ)(1/ϵ^Tϵ[min_x_j∈𝒟 (Wx_i + ϵ-Wx_j)^T(Wx_i + ϵ-Wx_j)] -2log|𝒟|)
       = -1/√(πϵ^Tϵ)(1/ϵ^Tϵ[min_x_j∈𝒟 (Wx_i -Wx_j)^T(Wx_i -Wx_j) + 2(Wx_i -Wx_j)^Tϵ + ϵ^T ϵ] -2log|𝒟|)
       ≥ -1/√(πϵ^Tϵ)(4k_1^2 k_2^2/ϵ^Tϵ + 4k_1k_2/√(ϵ^Tϵ) + 1-2log|𝒟|)

where we used the inequalities exp(-1/z) ≤ z^3/2, exp(-1/z) ≤ z^1/2 for any z>0 in the second step (with z=2σ^2/ϵ^Tϵ) and Cauchy-Schwarz for the last step (since ‖ x‖_2 ≤ k_1 for all x ∈𝒟, ‖ W‖_F ≤ k_2 for some positive constants k_1, k_2 ∈ℝ^+ by assumption). 


Since both the upper and lower bounds for the integrand are independent of σ, dominated convergence holds for the third term in Eq. (<ref>).







Consequently, we can evaluate limits to obtain a limiting ratio between the slack term and the lower bound:

    lim_σ→∞S(W, 𝒟, σ)/C(W, 𝒟, σ) = 0

using the expressions derived in Eq. (<ref>) and Eq. (<ref>), dominated convergence for interchanging limits and expectations, along with L'Hôpital's rule.




We can now rewrite Eq. (<ref>) as:

    M(W, 𝒟, σ)    = C(W, 𝒟, σ)(1 + S(W, 𝒟, σ)/C(W, 𝒟, σ)) .


By the (ϵ, δ) definition of limit, we know that for any fixed W that satisfies ‖ W ‖_F≤ k_2 and ∀ϵ>0, there exists a δ>0 such that ∀σ > δ, we have:

    | M(W, 𝒟, σ) - C(W, 𝒟, σ) | < ϵ.


Next, we note that the slack term S(W, 𝒟, σ) is monotonic in σ and converges pointwise for any fixed W that satisfies ‖ W ‖_F≤ k_2. 

    lim_σ→∞ S(W, 𝒟, σ)    = lim_σ→∞ M(W, 𝒟, σ) - C(W, 𝒟, σ) =0.


Using Dini's Theorem, this implies the convergence of the slack term is uniform in W as σ→∞. Hence, for all W that satisfy ‖ W ‖_F≤ k_2 and ∀ϵ>0, there exists a δ>0 such that ∀σ > δ, we have:

    | M(W, 𝒟, σ) - C(W, 𝒟, σ) | < ϵ.


Since the max operator preserves continuity (via Berge's maximum theorem) and is assumed to be identifiable, we conclude that ∀ W satisfying ‖ W ‖_F≤ k_2 and ∀ϵ>0, there exists a δ>0 such that ∀σ > δ, we have:

    | W^∗ - max_W∑_x_i, x_j∈𝒟(Wx_i-Wx_j)^T(Wx_i -Wx_j)) | < ϵ






which finishes the proof.





§ EXPERIMENTAL DETAILS

For MNIST, we use the train/valid/test split of 50,000/10,000/10,000 images. For Omniglot, we use train/valid/test split of 23,845/500/8,070 images. For CelebA, we used the splits as provided by  <cit.> on the dataset website. All images were scaled such that pixel values are between 0 and 1. We used the Adam optimizer with a learning rate of 0.001 for all the learned models. For MNIST and Omniglot, we used a batch size of 100. For CelebA, we used a batch size of 64. Further, we implemented early stopping based on the best validation bounds after 200 epochs for MNIST, 500 epochs for Omniglot, and 200 epochs for CelebA. 




 §.§ Hyperparameters for compressed sensing on MNIST and Omniglot


For both datasets, the UAE decoder used 2 hidden layers of 500 units each with ReLU activations. The encoder was a single linear layer with only weight parameters and no bias parameters. The encoder and decoder architectures for the VAE baseline are symmetrical with 2 hidden layers of 500 units each and 20 latent units. We used the LASSO baseline implementation from sklearn and tuned the Lagrange parameter on the validation sets. For the baselines, we do 10 random restarts with 1,000 steps per restart and pick the reconstruction with best measurement error as prescribed in <cit.>. Refer to <cit.> for further details of the baseline implementations.




Table <ref> shows the average norms for the random Gaussian matrices used in the baselines and the learned UAE encodings. The lower norms for the UAE encodings suggest that the UAE baseline is not trivially overcoming noise by increasing the norm of W.




 §.§ Hyperparameters for dimensionality reduction

For PCA and each of the classifiers, we used the standard implementations in sklearn with default parameters and the following exceptions: 


  * KNN: n_neighbors = 3

  * DT: max_depth = 5

  * RF: max_depth = 5,  n_estimators = 10, max_features = 1

  * MLP: alpha=1

  * SVC: kernel=linear, C=0.025




 §.§ Statistical compressed sensing on CelebA dataset






For the CelebA dataset, the dimensions of the images are 64 × 64 × 3 and σ=0.01. The naive pixel basis does not augur well for compressed sensing on such high-dimensional RGB datasets. Following <cit.>, we experimented with the Discrete Cosine Transform (DCT) and Wavelet basis for the LASSO baseline. Further, we used the DCGAN architecture <cit.> as in  <cit.> as our main baseline. For the UAE approach, we used additional convolutional layers in the encoder to learn a 256 dimensional feature space for the image before projecting it down to m dimensions. 

Encoder architecture:

    Signal   →Conv[Kernel: 4x4, Stride: 2, Filters: 32, Padding: Same, Activation: Relu]
       →Conv[Kernel: 4x4, Stride: 2, Filters: 32, Padding: Same, Activation: Relu]
       →Conv[Kernel: 4x4, Stride: 2, Filters: 64, Padding: Same, Activation: Relu]
       →Conv[Kernel: 4x4, Stride: 2, Filters: 64, Padding: Same, Activation: Relu]
       →Conv[Kernel: 4x4, Stride: 1, Filters: 256, Padding: Valid, Activation: Relu]
       →Fully_Connected[Units: m, Activation: None]


Decoder architecture:

    Measurements   →Fully_Connected[Units: 256, Activation: Relu]
       →Conv_transpose[Kernel: 4x4, Stride: 1, Filters: 256, Padding: Valid, Activation: Relu]
       →Conv_transpose[Kernel: 4x4, Stride: 2, Filters: 64, Padding: Same, Activation: Relu]
       →Conv_transpose[Kernel: 4x4, Stride: 2, Filters: 64, Padding: Same, Activation: Relu]
       →Conv_transpose[Kernel: 4x4, Stride: 2, Filters: 32, Padding: Same, Activation: Relu]
       →Conv_transpose[Kernel: 4x4, Stride: 2, Filters: 3, Padding: Valid, Activation: Sigmoid]



We consider m={20, 50, 100, 200, 500} measurements. The results are shown in Figure <ref>. While the performance of DCGAN is comparable with that of UAE for m=500, UAE outperforms DCGAN significantly when m is low. The LASSO baselines do not perform well, consistent with the observations made for the experiments on the MNIST and Omniglot datasets. Qualitative evaluations are shown in Figure <ref> for m=50 measurements.







§ ADDITIONAL RELATED WORK


Dictionary learning. An uncertainty autoencoder can be also seen as a more flexible, generalized form of undercomplete dictionary learning with non-linear encoding and decoding. To see the connection, consider the simplified noise-free setting where the decoding distribution is a Gaussian with fixed variance and the mean of the decoding function is linear in a linear function of the measurements. That is, we are considering a standard linear autoencoder with Y=WX and P_θ(X | Y) = 𝒩(WY, Σ), where W is some decoding matrix. Under these assumptions, the UAE objective simplifies to:

    min_W, W𝔼_x ∼ Q_data[‖ x - W W x  ‖_2^2].

If we think of the decoding W as a dictionary and the encoding WX as the representation then we arrive at an undercomplete dictionary learning. A large body of prior research has focussed on overcomplete dictionary learning for compressed sensing. Here, the goal is to learn an encoding dictionary and an overcomplete basis in which the original signal is sparse. This basis allows us to leverage algorithms for compressed sensing that are designed based on sparsity assumptions over the signals (see <cit.> and references therein). An uncertainty autoencoder makes no sparsity assumptions, and it crucially learns a decoding dictionary and an encoding basis. By adding more (non-linear) layers to the encoder, one could also learn an (over/under) complete basis for the dataset with desired properties such as sparsity.


Further applications of variational information maximization. The variational information maximization principle underlies many recent algorithms and tasks, such as feature selection <cit.>, interpretable representation learning in generative adversarial networks <cit.>, and intrinsic motivation in reinforcement learning <cit.>. 



