Machine learning models involving discrete latent variables require gradient estimators to facilitate backpropagation in a computationally efficient manner. The most recent addition to the Straight-Through family of estimators, ReinMax, can be viewed from a numerical ODE perspective as incorporating an approximation via Heun's method to reduce bias, but at the cost of high variance. In this work, we introduce the ReinMax-Rao and ReinMax-CV estimators which incorporate Rao-Blackwellisation and control variate techniques into ReinMax to reduce its variance. Our estimators demonstrate superior performance on training variational autoencoders with discrete latent spaces. Furthermore, we investigate the possibility of leveraging alternative numerical methods for constructing more accurate gradient approximations and present an alternative view of ReinMax from a simpler numerical integration perspective.
Figures & tables
ReinMax
Straight-Through
ReinMax-Argmax
Std
7.400
2.733
2.873
Table 1: Sample standard deviation of different estimators after 50 epochs of training a discrete VAE with latent size 8×4 . The sample standard deviation is computed over 1024 samples of D while keeping the minibatch of input data fixed. We use τ=1.3 for all methods.
Figure 1: The sample bias and variance of the estimators over checkpoints of a discrete VAE trained with ReinMax, with 8 categorical dimensions and 4 latent dimensions. The bias is measured using the cosine similarity between the exact gradient and the sample mean of 1024 gradient estimates with a fixed batch of size 100 and fixed model parameters. The variance is simply the sample variance of 1024 gradient estimates. All the temperatures here are 1, except τ=0.1 for ReinMax-CV.
Method
8×4
4×24
8×16
16×12
64×8
10×30
Gumbel
127.59 ± 0.23
101.49 ± 0.09
99.20 ± 0.14
100.42 ± 0.15
103.87 ± 0.13
99.80 ± 0.11
Gumbel-Rao
125.65 ± 0.14
100.14 ± 0.12
99.63 ± 0.23
109.72 ± 6.83
112.48 ± 0.18
102.28 ± 0.19
ST
136.40 ± 0.41
112.21 ± 0.19
113.32 ± 0.09
113.31 ± 0.08
113.97 ± 0.07
136.73 ± 0.40
GST-1.0
128.31 ± 0.21
101.64 ± 0.11
98.31 ± 0.16
98.37 ± 0.17
102.41 ± 0.16
98.66 ± 0.14
ReinMax
125.08 ± 0.20
99.88 ± 0.08
97.80 ± 0.09
97.98 ± 0.13
101.09 ± 0.11
98.50 ± 0.13
ReinMax-Rao
125.31 ± 0.28
100.41 ± 0.12
98.03 ± 0.14
97.62 ± 0.09
100.24 ± 0.16
98.61 ± 0.16
Table 2: Train ELBO ( ↓ ) across different configurations of categorical dimension × latent dimension. Best results are in bold , second best are underlined , third best are italic .
Method
8×4
4×24
8×16
16×12
64×8
10×30
Gumbel
128.72 ± 0.22
103.80 ± 0.11
101.37 ± 0.16
102.71 ± 0.13
106.03 ± 0.15
101.74 ± 0.11
Gumbel-Rao
127.62 ± 0.12
102.88 ± 0.15
102.07 ± 0.23
111.28 ± 6.70
113.71 ± 0.17
103.97 ± 0.21
ST
137.01 ± 0.45
113.63 ± 0.20
114.81 ± 0.10
114.46 ± 0.09
115.45 ± 0.10
136.45 ± 0.61
GST-1.0
129.15 ± 0.15
104.04 ± 0.11
101.27 ± 0.14
101.67 ± 0.20
105.34 ± 0.18
100.77 ± 0.12
ReinMax
127.27 ± 0.16
102.26 ± 0.06
100.73 ± 0.10
100.86 ± 0.13
103.35 ± 0.10
100.80 ± 0.13
ReinMax-Rao
126.60 ± 0.28
103.31 ± 0.13
100.74 ± 0.14
100.32 ± 0.11
102.76 ± 0.19
101.01 ± 0.18
Table 3: Test ELBO ( ↓ ) across different configurations of categorical dimension × latent dimension. Best results are in bold , second best are underlined , third best are italic .
Figure 2: The ELBO on the train set for the 8×4 , 8×16 and 4×24 VAEs after 50 epochs trained with ∇\mboxReinMax−RK2,β as a function of β , spaced evenly from -0.2 to 1.2 in increments of 0.05. The minimum is achieved at approximately β=0.5 which corresponds to the original ReinMax estimator. We use τ=1 here.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
8×4
4×24
8×16
16×12
64×8
10×30
STGS
Optimizer
Adam
RAdam
RAdam
RAdam
RAdam
RAdam
Learning Rate
0.0003
0.0005
0.0005
0.0007
0.0007
0.0005
Temperature
0.5
0.3
0.5
0.7
0.7
0.5
GR
Optimizer
Adam
RAdam
RAdam
Adam
Adam
RAdam
Learning Rate
0.0005
0.00055
0.0007
0.0005
0.0007
0.0005
Temperature
0.5
0.3
0.7
1.0
2.0
1.0
Appendix
Table 4: Hyperparameters used for the VAEs of varying latent space sizes.