HEmile / storchastic

Stochastic Automatic Differentiation library for PyTorch.
GNU General Public License v3.0
177 stars 5 forks source link

comparison between gumbel softmax, rebar, and relax #102

Open fnzhan opened 2 years ago

fnzhan commented 2 years ago

Hi, @HEmile , I test the example in examples/vae/discrete_vae.py, but find the gumbel softmax performs much better than rebar, relax and reinforce (testing loss after 10 epochs: 98 for gumbel, 165 for rebar, 208 for relax, 209 for reinforce). The performance is not that consistent with the results in your guide, where gumbel softmax is slightly worse than reinforce in terms of training loss. Is there anything wrong with the example? I don't change any parameters. Looking forward to your response.

fnzhan commented 2 years ago

Besides, I also test the sample in examples/vae_reference.py, still find Gumbel outperform reinforce and rebar. (n_samples=10, final loss: 122 for Gumbel, 157 for reinforce, 152 for rebar).

HEmile commented 2 years ago

Hi @fnzhan . That is a good question! I wrote that tutorial a while ago, and am surprised by the results myself. I think the point here is that I used a latent space of 2 discrete latents with 10 categories. Apparently, back then, Gumbel softmax did rather poorly. In general I'd expect Gumbel to outperform REINFORCE, REBAR and RELAX if it is applicable. Did you test with a latent space of 10x10 discrete latents (default in the example code) or 2x10?

fnzhan commented 2 years ago

Hi @HEmile Thanks for your reply. In my test, the latent space size in examples/vae/discrete_vae.py and examples/vae_reference.py are 20×10 and 2×10 (i.e., the default setting), respectively. I run the examples/vae_reference.py with n_sample=10, while the n_sample is 1 in your tutorial. Thus, I guess the inconsistency for examples/vae_reference.py is caused by the different n_sample. For examples/vae/discrete_vae.py, I run with n_sample=1.

In all, it seems REINFORCE outperforms Gumbel only when there is a small latent space (e.g., 20) and small n_sample (e.g., 1). Otherwise, Gumbel may outperform REINFORCE and REBAR consistently. Above is only for discrete VAE.

HEmile commented 2 years ago

Yes, I think that is consistent with my results. I will update the documentation to make this clearer.