[BUG] Stacked Discrete Flows ocasionally fail at recovering the origin
- Dominant language
- Jupyter Notebook
- Stars
- 712
- Forks
- 77
- Avg merge
- 9h 39m
- Merged PRs (30d)
- 1
Description
I stack several autoregressive discrete flows on top of each other using tf.keras.Sequential and pass samples forward through the stacked flows. Then, the output I pass in the opposite direction using reverse operation. Occasionally, I find incorrectly restored samples, e.g., 1 appearing in different than original position. For example, it fails with the below enclosed code for seed=49, with the following output: "seed=49 num_errors=8 max-error=1.00". Other failing seeds: 683, 736, 900. The problem seems to be more apparent when more flows are used. I use TF ver. 2.1.0 and TFP ver. 0.9.0 (Python 3.7.4).
```
import tensorflow as tf
import tensorflow_probability as tfp
import edward2 as ed
import numpy as np
for s in range(1000): # test for various seeds
# fix seed
np.random.seed(s)
tf.random.set_seed(s)
# specify factorized base distribution
N, K = 3, 20
base_specs = np.arange(K)
base_specs = base_specs/sum(base_specs)
base = tfp.distributions.OneHotCategorical(base_specs)
base_samples = tf.cast( base.sample( (1000, N) ), 'float32')
# create a stack of flows
layers_specification = [[12, 16], [12, 16], [16, 12], [64, 32], [32, 16]]
layers = [ed.layers.DiscreteAutoregressiveFlow(
ed.layers.MADE(K, hidden_dims=hidden_dims),
tf.Variable(5.0, trainable=False))
for hidden_dims in layers_specification]
stacked_flows = tf.keras.Sequential(layers)
# pass samples both directions and check if the flow retrieves the origin
z = stacked_flows(base_samples) #forward
for f in stacked_flows.layers[-1::-1]: z = f.reverse(z) #reverse
worst_case_recovery_error = tf.reduce_max(z-base_samples).numpy()
num_errors = len( np.nonzero( tf.math.abs(z-base_samples).numpy() > 0.5 )[0] )
print("seed=%s num_errors=%i max-error=%.2f" % \
(s, num_errors, worst_case_recovery_error))
```
Contributor guide
Assessment
This issue has not been assessed yet.