OVERCOMING THE DISENTANGLEMENT VS RECON- STRUCTION TRADE-OFF VIA JACOBIAN SUPERVISION
A major challenge in learning image representations is the disentangling of the factors of variation underlying the image formation. This is typically achieved with an autoencoder architecture where a subset of the latent variables is constrained to correspond to specific factors, and the rest of them are considered nuisance variables. This approach has an important drawback: as the dimension of the nuisance variables is increased, image reconstruction is improved, but the decoder has the flexibility to ignore the specified factors, thus losing the ability to condition the output on them. In this work, we propose to overcome this trade-off by progressively growing the dimension of the latent code, while constraining the Jacobian of the output image with respect to the disentangled variables to remain the same. As a result, the obtained models are effective at both disentangling and reconstruction. We demonstrate the applicability of this method in both unsupervised and supervised scenarios for learning disentangled representations. In a facial attribute manipulation task, we obtain high quality image generation while smoothly controlling dozens of attributes with a single model. This is an order of magnitude more disentangled factors than state-of-the-art methods, while obtaining visually similar or superior results, and avoiding adversarial training 1 .
Summary: The paper proposes a method to tackle the disentanglement-reconstruction tradeoff problem in many disentangling approaches. This is achieved by first training the teacher autoencoder (unsupervised or supervised) that learns to disentangle the factors of variation at the cost of poor reconstruction, and then distills these learned representations into a student model with extra latent dimensions, where these extra latents can be used to improve the reconstructions of the student autoencoder compared to the teacher autoencoder. The distillation of the learned representation is encouraged via a novel Jacobian loss term that encourages the change in reconstructions of the teacher and student to be similar when the latent representation changes. There is one experiment for progressive unsupervised disentangling (disentangling factor by factor) on MNIST data, and one experiment for semi-supervised disentangling on CelebA-HQ.
Pros:
- I think the idea of progressively capturing factors of variation one by one is neat, and this appears to be one of the first successful attempts at this problem.
- The distillation appears to work well on the MNIST data, and does indeed decrease the reconstruction loss of the student compared to the teacher.
- The qualitative results on CelebA-HQ look strong (especially apparent in the video), with the clear advantage over Fader Networks being that the proposed model is a single model that can manipulate the 40 different attributes, whereas Fader Nets can only deal with at most 3 attributes per model.
Cons:
- There are not enough quantitative results supporting the claim that the model is “effective at both disentangling and reconstruction.” The degree of disentanglement in the representations is only shown qualitatively via latent interpolation, and only for a single model. Such qualitative results are generally prone to cherry-picking and it is difficult to reliably compare different disentangling methods in this manner. This calls for quantitative measures of disentanglement. Had you used a dataset where you know the ground truth factors of variation (e.g. dSprites/2D Shapes data) for the unsupervised disentangling method, then the level of disentanglement in the learned representations could be quantified, and thus your method could be compared against unsupervised disentangling baselines. For the semi-supervised disentanglement example on CelebA, you could for example quantify how well the encoder predicts the different attributes (because there is ground truth here) e.g. report RMSE of the y_i’s on a held out test set with ground truth. A quantitative comparison with Fader Networks in this manner appears necessary. The qualitative comparison on a single face in Figure 5 is nowhere near sufficient.
- There is quantitative evidence that the reconstruction loss decreases when training the student, but here it’s not clear whether this quantitative difference makes a qualitative difference in the reconstructions. Getting higher fidelity images is one of the motivations behind improving reconstructions, so It would be informative to compare the reconstructions of the teacher and the student on the same image.
- In the CelebA experiments, the benefit of student training is not visible in the results. In Figure 5 you already show that the teacher model gives decent reconstructions, yet you don’t show the reconstruction for the student model (quantitatively you show that it improves in Figure 3b, but again it is worth checking if it makes a difference visually). Also it’s not clear whether Figure 4 are results from the student model or the teacher model. I’m guessing that they are from the student model.
- These quantitative results could form the basis of doing ablation studies for each of the different losses in the additive loss (for both unsupervised & semi-supervised tasks). Because there are many components in the loss, with a hyperparameter for each, it would be helpful to know what losses the results are sensitive to for the sake of tuning hyperparameters. This would be especially useful should I wish to apply the proposed method to a different dataset.
- I think the derivation of the Jacobian loss requires some more justification. The higher order terms in the Taylor expansion in (2) and (3) can only be ignored when ||y_2 - y_1|| is small compared to the coefficients, but there is no validation/justification regarding this.
Other Qs/comments:
- On page 5 in the last paragraph of section 3, you say that “After training of the student with d=1 is finished, we consider it as the new teacher”. Here do you append z to y when you form the new teacher?
- On page 6 in the paragraph for prediction loss, you say “This allows the decoder to naturally …. of the attributes”. I guess you mean this allows the model to give realistic interpolations between y=-1 and 1?
- bottom of page 6: “Here we could have used any random values in lieu of y_2” <- not sure I understand this?
- typo: conditionnning -> conditioning
- I would be inclined to boost the score up to 7 if the authors include some quantitative results along with more thorough comparisons to Fader Networks
************ Revision ***********
The authors' updates include further quantitative comparisons to Fader Networks and ablation studies for the different types of losses, addressing the concerns I had in the review. Hence I have boosted up my score to 7.