Problem with "Regression with Probabilistic Layers in TensorFlow Probability"

user_5930903 0
Viewed 719

I'm having trouble using tfp.layers.DistributionLambda, I'm a TF newbie trying hard to make the tensors flow. Can someone please provide some insights into how to set up the output distribution's parameters?

Context:

TFP team wrote a tutorial on Regression with Probabilistic Layers in TensorFlow Probability, it set up the following model:

# Build model.
model = tfk.Sequential([
  tf.keras.layers.Dense(1 + 1),
  tfp.layers.DistributionLambda(
      lambda t: tfd.Normal(loc=t[..., :1],
                           scale=1e-3 + tf.math.softplus(0.05 * t[..., 1:]))),
])

My problem:

It outputs a normal distribution using tfp.layers.DistributionLambda, but I'm unclear how tfd.Normal's parameters (mean/loc and standard deviation/scale) were set up, so I'm having trouble changing the Normal to a Gamma Distribution. I tried the following, but didn't work (predicted distribution parameters are nan). 

def dist_output_layer (t, softplus_scale=0.05):
    """Create distribution with variable mean and variance
    """
    mean = t[..., :1]
    std_dev = 1e-3 + tf.math.softplus(softplus_scale * mean)

    alpha = (mean/std_dev)**2
    beta = alpha/mean

    return tfd.Gamma(concentration = alpha, 
                     rate = beta
                    )

# Build model.
model = tf.keras.Sequential([
    tf.keras.layers.Dense(20,activation="relu"), # "By using a deeper neural network and introducing nonlinear activation functions, however, we can learn more complicated functional dependencies!
    tf.keras.layers.Dense(1 + 1), #two neurons here b/c the output layer's distribution's mean and std. deviation
    tfp.layers.DistributionLambda(dist_output_layer)
])

Thanks a lot in advance.

tensorflow-probability tensorflow python
2 Answers

There is a a lot to say about the code snippet you pasted from Medium, to be honest.

I hope you will find my comments below somewhat useful, though.

# Build model.
model = tfk.Sequential([

    # The first layer is a Dense layer with 2 units, one for each of the parameters that will
    # be learnt (see next layer). Its implied shape is (batch_size, 2).
    # Note that this Dense layer has no activation function as we want are any real value that will be used
    # to parameterize the Normal distribution in the Normal distribution component of the following
    # layer
    tf.keras.layers.Dense(1 + 1),

    # The following layer is a DistributionLambda that encapsulates a Normal distribution. The
    # DistributionLambda takes a function in its constructor, and this function should take the output
    # tensor from the previous layer as its input (this is the Dense layer and the comments above).
    # The goal is to learn the 2 parameters of the distribution that is loc (the mean) and scale (the standard
    # deviation). For this, a lambda construct is used. The ellipsis you can see for the loc
    # and scale arguments (that is the 3 dots) are for the batch size. Also note that scale (the standard deviation)
    # cannot be negative. The softplus function was used to make sure that the learnt parameter scale doesn't get
    # negative.
    tfp.layers.DistributionLambda(
      lambda t: tfd.Normal(loc=t[..., :1],
                       scale=1e-3 + tf.math.softplus(0.05 * t[..., 1:]))),
])
user_11153770

Regarding the question about the .05 being added, it's a small offset to solve some gradient issues that can arise without it. Basically a prior saying we're confident that the real variability is NOT smaller than epsilon (here .05), so we're gonna make sure that the std dev is never smaller by just adding that.

See https://github.com/tensorflow/probability/issues/751

Money quote:

"If infinitesimal scales end up being a problem in practice on a given task, the fix we commonly use is a softplus-and-shift, e.g. scale = epsilon + tf.math.softplus(unconstrained_scale), where epsilon is some tiny value like 1e-5 that we are a priori confident is much smaller than the true scale."

EDIT: Actually what is added is 1e-3 for the reasons I described above. As for the multiplication.... might again just be a scaling or gradient adjustment. Or perhaps to make the scale parameter begin at a certain size.

user_159133810
Related