Low MIG values bug found & solution

[DianeBouchacourt]

I trained a beta TCVAE with the code from https://github.com/rtqichen/beta-tcvae which gives MIG for beta TCVAE of ~0.50. When computing MIG with your code with the same model (based on MLP), I had values close to 0.0008.

Differences with Chen's code I found important:

• MIG values are not computed on shape in Chen's code (not considered a factor of variation). I had to modify the dsprites dataset to remove shape from dSprites lat_names, and write a custom _estimate_H_zCv function. I can share if you want.

• Chen uses samples, not the mean as you do here

  disentangling-vae/disvae/models/vae.py

  Line 71 in 7b8285b

  return mean

  since self.training is False

• The most important change is I changed these lines

  disentangling-vae/disvae/evaluate.py

  Line 272 in 7b8285b

  samples_zCx = samples_zCx.expand(len_dataset, latent_dim, n_samples)

  for

    samples_zCx = samples_zCx.permute(1,0)
    samples_zCx = samples_zCx.index_select(1, samples_x).view(latent_dim, n_samples)
    samples_zCx = samples_zCx.view(1, latent_dim, n_samples).expand(len_dataset, latent_dim, n_samples)
    mean = params_zCX[0].view(len_dataset, latent_dim, 1).expand(len_dataset, latent_dim, n_samples)
    log_var = params_zCX[1].view(len_dataset, latent_dim, 1).expand(len_dataset, latent_dim, n_samples)
  

which are closer to Chen's code, and I get values of ~0.50 now too. I don't exactly know why the original lines where not expanding the correct way

[DianeBouchacourt]

After a bit more digging, I think the problem comes from the fact that indexes are selected but the view() operation

disentangling-vae/disvae/evaluate.py

Line 272 in 7b8285b

samples_zCx = samples_zCx.expand(len_dataset, latent_dim, n_samples)

is reordering in a wrong manner, permute is the correct way to do it. Indeed, I've checked by taking the first 10000 samples, which are also the first 10000 samples in mean. With the view() operation it does not return the same values as in mean (which it should since the samples are mean values). I replaced with permute and it works as expected.

Thus I think it should be:

    samples_zCx = samples_zCx.index_select(0, samples_x).permute(1,0)
    samples_zCx = samples_zCx.unsqueeze(0).expand(len_dataset, latent_dim, n_samples)

in place of

disentangling-vae/disvae/evaluate.py

Line 269 in 7b8285b

samples_zCx = samples_zCx.index_select(0, samples_x).view(latent_dim, n_samples)

and

disentangling-vae/disvae/evaluate.py

Line 272 in 7b8285b

samples_zCx = samples_zCx.expand(len_dataset, latent_dim, n_samples)

[YannDubs]

Awesome thanks Diane, a few others had open issues for that but there was no solutions (I don't actively work on that anymore). Do you think you could send a PR ?

ideally, it would use samples instead of the mean , but I think it's more urgent to have the correct permutations !

[DianeBouchacourt]

Let me check my code again, and I will ! I also need to do a PR just for these changes (I edited quite some stuff in the fork I created).