arnauqb/gumbelsoftmax.jl

Julia implementation of the Gumbel-Softmax reparametrization trick compatible with Zygote and ForwardDiff

GumbelSoftmax

This package implements:

1. Gumbel-Softmax
2. Rao-Blackwellized Gumbel-Softmax

in the Julia programming language. The package supports both forward-mode automatic differentiation (AD) and reverse-mode AD through Zygote and ForwardDiff respectively.

1. Installation

The easiest way is to get the package directly from the Julia repository

using Pkg
Pkg.add("GumbelSoftmax")

2. Usage

The expected input shape is (latent_dimension, categorical_dimension, batch_dimension). As an example, let's suppose we have 4 Categorical distributions with 3 classes and we want to sample 10 times. In this case, latent_dimension=4, categorical_dimension=3, and batch_dimension=10.

using GumbelSoftmax, Random
logits = randn(4, 3, 10)
samples = sample_gumbel_softmax(logits=logits, tau=0.1, hard=true)
# or with Rao-Blackwellization
k = 10 # number of Monte-Carlo samples
samples = sample_rao_gumbel_softmax(logits=logits, tau=0.1, k=k)

3. Example: Discrete Variational Autoencoder (VAE)

We include an example of using the Gumbel-Softmax trick to implement a discrete VAE. The example can be found in examples/vae.jl and it can be run with

julia examples/vae.jl

Here are some results:

VAE loss:

VAE reconstructions and generated samples:

gumbelsoftmax.jl's Issues

TagBot trigger issue

This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers.
Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue.
I'll open a PR within a few hours, please be patient!

Help with VAE and gumbel-softmax trick

Thanks for the code! I'm trying to learn Julia and Flux.jl but I'm having trouble finding an example of a VAE with gumbel-softmax trick. Chances are that you have already done this. Could you provide an example for me?