Provide seed to DenseReparameterization layer

[danijar]

I use DenseReparameterization for the transition function of a simple state-space model. To sample posterior sequences, I need to auto-regressively apply the layer to its own output inside of a symbolic loop. However, the samples are not consistent because the layer uses a different weight matrix at each iteration of the loop. Is it possible or are there plans to pass the seed into layer(inputs, seed=0) in order to sample consistent sequences?

[dustinvtran]

That's a good question. We mimicked the randomness approach from Keras' dropout layer. Do you know what their recommendation is?

[danijar]

I don't know. @fchollet?

[danijar]

A simple API for this would be to add the seed as optional argument to __call__.

[dustinvtran]

It looks like the seed is in the init of tf.keras.layers.Dropout (https://www.tensorflow.org/api_docs/python/tf/layers/Dropout).

That should also work for us. Carry around the seed by storing a tfp.distributions.SeedStream object. This can be written in the init of the abstract classes like _DenseVariational. Then generate a new seed at each random op call via seed=self._seed_stream().

See tfp.mcmc.MetropolisHastings for an example implementation. Contributions welcome.

[dustinvtran]

@srvasude

On second thought, this is actually difficult with the current design because any randomness is imposed by the user via the {kernel,bias}_posterior_tensor_fn arguments. Perhaps this is another example of why we need to redesign the API to such generic callables.

[danijar]

How would passing the seed into __init__ solve this? If I need to generate multiple sequences using the same Bayesian transition model, I'd have to reuse the weights but not the seed within the same session run.

[dustinvtran]

By passing a seed into __init__, you can initialize a seed stream which lets you use a different seed at each call().

[danijar]

But I couldn't use it to use the same seed for multiple calls? The whole point of this request was that I can sample a consistent autoregressive sequence p(x_1:T,theta)~p(theta)p(x_1)prod_t=2..T p(x_t|x_t-1,theta). I want to implement int p(x_t|x_t-1,theta)q(theta) dtheta as a probabilistic layer, but the seed for the parameters needs to stay fixed over all time steps when sampling a sequence.

[dustinvtran]

As I mentioned, you can use the exact same seed by passing {kernel,bias}_posterior_tensor_fn as lambda d: d.sample(seed=exact_same_seed) in the __init__.

By passing a seed into __init__ (and a necessary redesign of the sampling op's API), you can also use a different seed at each call().

In other words, both options are possible.

[danijar]

Oh, I missed that first point about providing kernel_posterior_tensor_fn. However, I still don't quite see how this solves the problem, because exact_same_seed will either be a constant or a tensor that can have only a single value per session run. I don't think this is possible to solve without the ability to set a separate seed for each call.

In other words, it would not be possible to generate multiple sequences from my Bayesian state space model within the same session run, since this cannot be overridden on a per-call basis. If I set exact_same_seed, I get consistent sequences but even multiple sequences will use the same weights. Or if I don't set it, I get different sequences but the weights are not consistent within each sequence.

[dustinvtran]

it would not be possible to generate multiple sequences from my Bayesian state space model within the same session run,

Yeah you can't do that unless you add seed to call; or otherwise you build a new layer and set its variables to be the previous layers.

I'm still not sure passing {kernel,bias}_posterior_tensor_fn during __init__ is the right way to instantiate tensors given a distribution. Any ideas on redesigning that are welcome—whether it be during the init or the call.

[danijar]

Can you remind me which use cases are currently covered by kernel_posterior_tensor_fn? It almost seems to me that in a reparameterization layer, you always wanted to just do dist.sample(). If you wanted multiple samples for a better estimate, you'd need to implement a new layer anyway, right? To me the straightforward change would be to allow passing the seed to __call__. If it's not set, it defaults to None and uses an internal seed stream of the layer.