Probability: TFP and TF2.0 incompatibility

[awav]

After using TFP with TF2.0, I've realized that I'm spending a lot of time on understanding why it doesn't work rather than making any progress. The TFP relies too much on TF1.0, even not realizing it. Some conceptions of the TF2.0 and the TFP are just incompatible, e.g. tf.Module.

class Model(tf.Module):
    def __init__(self):
        pass  # Create variables, set the prior and etc.

data = ... setup data
model = Model()

def run_chain():
    def posterior_log_prob_fn(*parameters):
        model.assign_parameters(*parameters)
        log_prob = -model.neg_log_marginal_likelihood(data)
        return log_prob

    # state_gradients_are_stopped=True
    variables = model.trainable_variables
    hmc = mcmc.HamiltonianMonteCarlo(target_log_prob_fn=posterior_log_prob_fn,
                                     num_leapfrog_steps=2,
                                     step_size=1.0,
                                     state_gradients_are_stopped=True)
    adaptive_hmc = mcmc.SimpleStepSizeAdaptation(hmc, num_adaptation_steps=2)
    return mcmc.sample_chain(num_results=2,
                             num_burnin_steps=0,
                             current_state=variables,
                             kernel=adaptive_hmc)

The HMC will not be able to compute the gradients as there are no links in a graph between states and the model (with or without tf.function).

Other examples you can find here: https://github.com/tensorflow/probability/issues/333, https://github.com/tensorflow/probability/issues/348, https://github.com/tensorflow/probability/issues/47

There is no easy fix for that and hope I'm wrong), as the roots of these problems are in TF2.0 itself. I've created another github issue at TF2.0: https://github.com/tensorflow/tensorflow/issues/29367, describing how and what makes TFP framework less effective.

[csuter]

Hi @awav, as mentioned just now on the parallel TF issue, we have a candidate solution for this, which we're working on getting checked in now. It's only a template, and we'll have to do a substantial refactoring of all bijectors and distributions to bring them up to date with the new approach. Stay tuned.

[awav]

@csuter, fantastic! Do you have any time estimates when it will be finished?

In example above the problem isn't exactly linked to the distribution's and linear operator's implementations. In my code snippet, the posterior_log_prob_fn connects model's parameters and the passed state though assign operation, therefore the gradient value is none. Wouldn't it better to make interface for HMC in a such way that it keeps variables for which the user wants to compute gradients and use them as a state, what do you think? Btw, I have a butchered version of TFP's HMC, where it works with tf.Module as I described.

[SiegeLordEx]

tfp.mcmc never worked with variables like that, so it's not a TF 2.0 specific problem. In general, tfp.mcmc is not designed to work with them because variables have a lot of limitations compared to immutable tensors. For example, with tensors you can run multiple MCMC chains in parallel. To do that with variables would require quite a bit of setup. If you really wanted to use variables, you can utilize tf.custom_gradient to decorate your posterior_log_prob_fn, something like this:

class Module(tf.Module):
  def __init__(self):
    self.v = tf.Variable(0.)

  def log_prob(self):
    return -tf.square(self.v)

def make_log_prob_fn(module):
  @tf.custom_gradient
  def log_prob_fn(*args):
    for v, val in zip(module.trainable_variables, args):
      v.assign(val)
    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(module.trainable_variables)
      lp = module.log_prob()

    def grad_fn(dlp, variables):
      return tape.gradient(lp, module.trainable_variables, dlp), [None] * len(variables)

    return lp, grad_fn
  return log_prob_fn

module = Module()
log_prob_fn = make_log_prob_fn(module)

with tf.GradientTape() as tape:
  x = tf.constant(-2.)
  tape.watch(x)
  lp = log_prob_fn(x)

print(tape.gradient(lp, x))

In the past, we had a way to convert a model using variables to one that used tensors by intercepting variable creators. That code got deleted due to a lack of use and incompatibility with TF2.0, but in principle it might still be possible to accomplish something like that again. The advantage of this latter method is that you wouldn't need to call assign anywhere, which is a pretty inefficient operation.

[awav]

@SiegeLordEx thanks a lot for this trick!

One question:

with tensors you can run multiple MCMC chains in parallel

I'd like to run multiple chains! How does it look in practice? Is there an option for sample_chain?

[brianwa84]

If your initial state is a batch of items, and your target density function
returns that same size batch, you will be running a batch of chains.

Brian Patton | Software Engineer | [email protected]

On Thu, Jun 6, 2019 at 11:29 AM Artem Artemev [e-mail]
wrote:

@SiegeLordEx https://github.com/SiegeLordEx thanks a lot for this trick!

One question:

with tensors you can run multiple MCMC chains in parallel

I'd like to run multiple chains! How does it look in practice? Is there an
option for sample_chain?

—
You are receiving this because you are subscribed to this thread.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/448?email_source=notifications&email_token=AFJFSIYAB3J6VTEAUUTHXQTPZEUN3A5CNFSM4HSPFNS2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXDHGZQ#issuecomment-499544934,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSIYFF33H4UOVE3HVV5TPZEUN3ANCNFSM4HSPFNSQ
.

[brianwa84]

"target density function returns that same size batch" >> target density
function returns a same-sized batch of densities
Brian Patton | Software Engineer | [email protected]

On Thu, Jun 6, 2019 at 11:35 AM Brian Patton 🚀 [e-mail] wrote:

If your initial state is a batch of items, and your target density
function returns that same size batch, you will be running a batch of
chains.

Brian Patton | Software Engineer | [email protected]

On Thu, Jun 6, 2019 at 11:29 AM Artem Artemev [e-mail]
wrote:

@SiegeLordEx https://github.com/SiegeLordEx thanks a lot for this
trick!

One question:

with tensors you can run multiple MCMC chains in parallel

I'd like to run multiple chains! How does it look in practice? Is there
an option for sample_chain?

—
You are receiving this because you are subscribed to this thread.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/448?email_source=notifications&email_token=AFJFSIYAB3J6VTEAUUTHXQTPZEUN3A5CNFSM4HSPFNS2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXDHGZQ#issuecomment-499544934,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSIYFF33H4UOVE3HVV5TPZEUN3ANCNFSM4HSPFNSQ
.

[awav]

@brianwa84 thanks! Let's nail it down: when we pass the current state as [[N, M], [N, K], [N, L, J], ...] and my log posterior returns [N,], where N is a batch size, then the multi-chain will run. And if we used instead the variables with the same setting (batches for the state and the log posterior are equal) we don't benefit from parallelization because the variables represent "synchronization" point for assign operation.

@SiegeLordEx , @brianwa84 example with custom gradient has a drawback: two gradient tapes, one in hmc (tape1) and another in log posterior (tape2), tape1 records the gradient for tape2.gradient(...) op! Is there a way to do this trick w/o this overhead? E.g. an option to pass a function which computes forward and backward steps?

[awav]

@SiegeLordEx , @brianwa84 example with custom gradient has a drawback: two gradient tapes, one in hmc (tape1) and another in log posterior (tape2), tape1 records the gradient for tape2.gradient(...) op! Is there a way to do this trick w/o this overhead? E.g. an option to pass a function which computes forward and backward steps?

Alright, that might not be a problem, as grad_fn runs outside of the gradient tapes. Although, two tapes still record the gradients for the posterior function and because tape1 is not used, it can be skipped.

[tom-schoener]

I am also using the nightly build of tfp and the tf 2 beta version:
Tensorflow: 2.0.0-beta0 Tensorflow Probability: 0.8.0-dev20190617
The DenseFlipout layer is causing me some problems (see https://github.com/tensorflow/probability/issues/409).

What do you recommend? Should we (for now) use tf version 1.13 and update at a later point in time or apply some sort of temporary workaround if there is one?

[awav]

It's only a template, and we'll have to do a substantial refactoring of all bijectors and distributions to bring them up to date with the new approach. Stay tuned.

@csuter , @brianwa84 , @SiegeLordEx Hello everyone! Any updates on the topic?

[brianwa84]

I think we're at like 40/130 so far. Is there a particular subset you're
interested in?

Essentially the thing to look for is a method called
_parameter_control_dependencies. If the bijector/distribution has
overridden it, it is probably "tape friendly" (i.e. retains variable refs).
Otherwise, not yet.

Brian Patton | Software Engineer | [email protected]

On Thu, Jul 4, 2019 at 3:53 PM Artem Artemev [e-mail]
wrote:

It's only a template, and we'll have to do a substantial refactoring of
all bijectors and distributions to bring them up to date with the new
approach. Stay tuned.

@csuter https://github.com/csuter , @brianwa84
https://github.com/brianwa84 , @SiegeLordEx
https://github.com/SiegeLordEx Hello everyone! Any updates on the topic?

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/448?email_source=notifications&email_token=AFJFSI4I7UYKZIRTPGFGVMTP5ZIL7A5CNFSM4HSPFNS2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODZIBTHI#issuecomment-508565917,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSI7XIDKWR32IIHR2YQ3P5ZIL7ANCNFSM4HSPFNSQ
.

[awav]

@brianwa84,

The MultivariateNormalTriL has _parameter_control_dependencies, but looks like still there are no links between variables and the distribution:

In [73]: mean = tf.Variable([0., 0.])
In [74]: scale_tril = tf.linalg.eye(2)
In [75]: mvn = tfp.distributions.MultivariateNormalTriL(mean, scale_tril)
In [76]: mvn.sample()
Out[76]: <tf.Tensor: id=1788, shape=(2,), dtype=float32, numpy=array([-0.35811806,  0.4914775 ], dtype=float32)>
In [77]: mean.assign([100000., 200000.])
Out[77]: <tf.Variable 'UnreadVariable' shape=(2,) dtype=float32, numpy=array([100000., 200000.], dtype=float32)>
In [78]: mvn.sample()
Out[78]: <tf.Tensor: id=1830, shape=(2,), dtype=float32, numpy=array([0.5472931, 2.378771 ], dtype=float32)>

[brianwa84]

Generally the distribution needs to override that function.

MvnTriL._pcd == Distribution._pcd

The mvns depend on TF linear operator, so we are blocked for the moment on
those, but @langmore is working on it.

On Thu, Jul 11, 2019, 6:36 AM Artem Artemev [e-mail]
wrote:

@brianwa84 https://github.com/brianwa84,

The MultivariateNormalTriL has _parameter_control_dependencies, but looks
like still there are no links between variables and the distribution:

In [73]: mean = tf.Variable([0., 0.])
In [74]: scale_tril = tf.linalg.eye(2)
In [75]: mvn = tfp.distributions.MultivariateNormalTriL(mean, scale_tril)
In [76]: mvn.sample()
Out[76]: In [77]: mean.assign([100000., 200000.])
Out[77]: In [78]: mvn.sample()
Out[78]:

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/448?email_source=notifications&email_token=AFJFSIYE2S2S7UHICCHTBMLP64EKVA5CNFSM4HSPFNS2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODZWJCVY#issuecomment-510431575,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSI66WCWJRSWCQDTHKTLP64EKVANCNFSM4HSPFNSQ
.

[awav]

@brianwa84 , @csuter Hello guys! What is the status for this feature? All MultivariateNormal* still don't work with variables, but I checked linear operators and they work fine with variables (I didn't check all of them though):

In [66]: tril = tril.assign([[10., 0.], [10., 10.]])

In [67]: lo_tril = tf.linalg.LinearOperatorLowerTriangular(tril)

In [68]: lo_tril.to_dense()
Out[68]:
<tf.Tensor: id=883, shape=(2, 2), dtype=float32, numpy=
array([[10.,  0.],
       [10., 10.]], dtype=float32)>

In [69]: tril.assign([[99., 0.], [99., 99.]])
Out[69]:
<tf.Variable 'UnreadVariable' shape=(2, 2) dtype=float32, numpy=
array([[99.,  0.],
       [99., 99.]], dtype=float32)>

In [70]: lo_tril.to_dense()
Out[70]:
<tf.Tensor: id=889, shape=(2, 2), dtype=float32, numpy=
array([[99.,  0.],
       [99., 99.]], dtype=float32)>

[csuter]

Hi Artem, MVNLinearOperator was migrated this week, and we're working on other MVN's next.

[kristofgiber]

Is it safer to avoid using TFP with TF2-beta and stay with TFP+TF1.14 for now?

[csuter]

@kristofgiber we strongly recommend starting to move towards TF 2.0 idioms. Right now about 70% of TFP Distributions and Bijectors have been updated to work more naturally with TF2, and we're working on the remaining ones. You can always ensure things will work as expected by writing your loss explicitly as a function of any tf.Variables, creating distributions inside the loss, and summing up -log_probs. To avoid extra python object creation overhead in your training loop, you can decorate the loss with @tf.function (this is often a good idea anyway). For examples of the "new way" of doing things, check out, e.g., the unit tests for the Normal distribution (search for "gradient" and "variable").

[kristofgiber]

@csuter Great, is there a list where one can follow which distributions exactly are already updated as you're moving forward? Or at least to see the list of those distributions / bijectors as of now?

[lmartak]

we have a candidate solution for this, which we're working on getting checked in now. It's only a template, and we'll have to do a substantial refactoring of all bijectors and distributions to bring them up to date with the new approach. Stay tuned.

@csuter does this candidate solution also handle the compatibility issue of TFP with tf.keras described here (and assessed below)?

Thanks!

[awav]

@csuter , @brianwa84 some bijectors still cannot handle variables properly:

In [5]: a = tf.Variable(1.)
In [6]: affine = tfp.bijectors.Affine(shift=a)
In [7]: affine.trainable_variables
Out[7]: ()

[SiegeLordEx]

@awav Affine is now deprecated as it was too tough to make it TF2 compatible: it got replaced by separate Scale and Shift, and a whole bunch of ScaleMatvec* bijectors. According to our internal checklist, all non-deprecated bijectors are compatible, and all but 7 distributions are compatible.

[yenicelik]

@SiegeLordEx just to make sure: so instead of using the tfb.Affine(), which two other tfb.Bijector modules should we use instead? (i.e. I assume chaining them together will work)

[yenicelik]

Also, which other 7 distributions are not compatible?

[csuter]

@yenicelik in place of Affine, we recommend you use (using ScaleMatvecTriL, for example)

shift = Shift(loc=...)
scale = ScaleMatvecTriL(scale_tril=...)
affine = shift(scale)

or, more succinctly, Shift(loc)(ScaleMatvecTriL(scale_tril)). These are both equivalent to writing Chain([Shift(loc), ScaleMatvecTriL(scale_tril)]).

The as-yet unmigrated Distributions are

• Autoregressive
• BatchReshape
• HiddenMarkovModel
• TransformedDistribution
• VariationalGaussianProcess
• VectorDiffeomixture

[SiegeLordEx]

Here's a rosetta stone between the arguments of tfb.Affine and the new bijectors (when used individually).

tfb.Affine(shift=shift)
# ===
tfb.Shift(shift)

tfb.Affine(scale_identity_multiplier=scale_identity_multiplier)
# ===
tfb.Scale(scale_identity_multiplier)

tfb.Affine(scale_diag=scale_diag)
# ===
tfb.ScaleMatvecDiag(scale_diag)

tfb.Affine(scale_tril=scale_tril)
# ===
tfb.ScaleMatvecTriL(scale_tril)

tfb.Affine(
  scale_perturb_factor=scale_perturb_factor,
  scale_perturb_diag=scale_perturb_diag,
)
# ===
linop = tf.linalg.LinearOperatorLowRankUpdate(
  tf.linalg.LinearOperatorIdentity(),
  u=scale_perturb_factor,
  diag_update=scale_perturb_diag,
  is_diag_update_positive=scale_perturb_diag is None,
  is_non_singular=True,
  is_self_adjoint=True,
  is_positive_definite=True,
  is_square=True)
tfb.ScaleMatvecLinearOperator(linop)

Now combining them is... complicated. tfb.Shift composes with all the other args (tfb.Affine does scale and then shift), but the general strategy for combining the other args is to use tfb.ScaleMatvecLinearOperator and construct the linear-operator doing the low rank updates you did before. Here's a way to do it (it doesn't quite do everything right, but it should let the gradients flow between the bijector outputs and the arguments).

tfb.Affine(
  scale_identity_multiplier=scale_identity_multiplier,
  scale_diag=scale_diag,
  scale_tril=scale_tril,
  scale_perturb_factor=scale_perturb_factor,
  scale_perturb_diag=scale_perturb_diag,
)
# ===
def make_tril(args):
  scale_identity_multiplier, scale_diag, scale_tril = args
  return scale_tril + scale_identity_multiplier * tf.eye(tf.shape(scale_tril)[-1]) + tf.diag(scale_diag)

tril = tfp.util.DeferredTensor((scale_identity_multiplier, scale_diag, scale_tril),
  make_tril, shape=tf.shape(scale_tril), dtype=scale_tril.dtype)
linop = tf.linalg.LinearOperatorLowerTriangular(
  tril,
  is_non_singular=True,
  is_self_adjoint=False)

linop = tf.linalg.LinearOperatorLowRankUpdate(
  linop,
  u=scale_perturb_factor,
  diag_update=scale_perturb_diag,
  is_diag_update_positive=scale_perturb_diag is None,
  is_non_singular=True,
  is_self_adjoint=True,
  is_positive_definite=True,
  is_square=True)
tfb.ScaleMatvecLinearOperator(linop)

[giovp]

@csuter @SiegeLordEx just wanted to report that tfb.AffineScalar is still lying around in e.g. tfb.MaskedAutoregressiveFlow and tfb.RealNVP.
Any plan for update?

[csuter]

Good catch @giovp. We will try to tidy that up shortly, but feel free to send a quick PR changing it to use tfb.Shift(shift=...)(tfb.Scale(scale=...)) -- it would be much appreciated!