We propose the Information Maximization Autoencoder (IMAE), an information theoretic approach to simultaneously learn continuous and discrete representations in an unsupervised setting. Unlike the Variational Autoencoder framework, IMAE starts from a stochastic encoder that seeks to map each input data to a hybrid discrete and continuous representation with the objective of maximizing the mutual information between the data and their representations. A decoder is included to approximate the posterior distribution of the data given their representations, where a high fidelity approximation can be achieved by leveraging the informative representations. We show that the proposed objective is theoretically valid and provides a principled framework for understanding the tradeoffs regarding informativeness of each representation factor, disentanglement of representations, and decoding quality.
Summary: the paper proposes a method for unsupervised disentangling of both discrete and continuous factors of variation in image data. It uses an autoencoder learned by optimising an additive loss composed of Mutual Information (MI) I(x;y,z) between the image x and the discrete+cts latents (y,z) and the reconstruction error. The mutual information is shown to decompose into I(x,y), I(x,z) and TC(y;z), and the I(x,z) is treated in a different manner to I(x,y). With Gaussian p(z|x), and it is shown that I(x,z_k) is maximal when p(z_k) is Gaussian. So KL(p(z_k)||N(0,1)) is optimised in lieu of optimising I(x,z), and I(x,y) (and TC(y;z)) is optimised by using mini-batch estimates of marginal distributions of y (and z). The paper claims improved disentangling of discrete and continuous latents compared to methods such as JointVAE and InfoVAE.
Pros:
- The derivation of the loss shows a nice link between Mutual information and total correlation in the latents.
- It is a sensible idea to treat the MI terms of the discrete latents differently to the continuous latents
- The mathematical and quantitative analysis of MI and its relation to decoder means and variances are informative.
Cons:
- There is not enough quantitative comparison of the quality of disentanglement across the different methods. The only values for this are the accuracy scores of the discrete factor, but for the continuous latents there are only qualitative latent traversals of single models, and I think these aren’t enough for comparing different disentangling methods - this is too prone to cherry-picking. I think it’s definitely necessary to report some metrics for disentangling that are averaged across multiple models trained with different random seeds. I understand that there are no ground truth cts factors for Mnist/FashionMnist, but this makes me think that a dataset such as dSprites (aka 2D Shapes) where the factors are known and has a mix of discrete and continuous factors would have been more suitable. Here you can use various metrics proposed in Eastwood et al, Kim et al, Chen et al for a quantitative comparison of the disentangled representations.
- In figure 4, it says beta=lamda=5 for all models. Shouldn’t you be doing a hyperparameter sweep for each model and choose the best value of hyperparameters for each? It could well be that beta=5 works best for IMAE but other values of beta/lambda can work better for the other models.
- When comparing against JointVAE, the authors point out that the accuracy for JointVAE is worse than that of IMAE, a sign of overfitting. You also say that VAT helps maintain local smoothness so as to prevent overfitting. Then shouldn’t you also be comparing against JointVAE + VAT? Looking at Appendix D, it seems like VAT makes a big difference in terms of I(y;y_true), so I’m guessing it will also have a big impact on the accuracy. Thus JointVAE + VAT might beat IMAE in terms of accuracy as well, at which point it will be hard to argue that IMAE is superior in learning the discrete factor.
- In the first paragraph of Section 4, the authors claim results on CelebA, but these are missing from the paper. Testing the approach on datasets more complex than (Fashion)Mnist would have been desirable.
- There aren’t any latent traversals for the discrete latents - this would be a useful visualisation to complement the accuracy plots in Figure 3.
Qs and comments:
- It’s not clear why posterior approximation quality (used as a starting point for motivating the loss) is an important quantity for disentangling.
- I see that the upper bound to I(x;z_k) in (4) and the objective in (6) have the same optimum at p(z_k) being Gaussian, but it’s not clear that increasing one leads to increasing the other. Using (6) to replace (4) seems to require further justification, whether it be mathematical or empirical.
- In proposition 2, I’m sceptical as to how meaningful the derived bound is, especially when you set N to be the size of the minibatch (B) in practice. It also seems that for small delta (i.e. to ensure high probability on the bound) and large K_2 (less restrictive conditions on p(y) and \hat{p}(y)), the bound can be quite big.
- \mathcal{L}_theta(y) in equation (10) hasn’t been introduced yet.
- The z dimension indices in the latent traversal plots of Figure 2 don’t seem to match the x-axis of the left figure. It’s not clear which are the estimates of I(x;z_k) for k=8,3,1 in the figure.