Optimise function for many pseudodata realisations in TensorFlow 2

Question

My end goal is to simulate likelihood ratio test statistics, however, the core problem I am having is that I do not understand how to get TensorFlow 2 to perform many optimi

Answer 1

Ok so here is what I came up with. The keys things I was missing were:

1. Define input variables as giant tensors so that all minimisations can occur at once.
2. Construct a single combined loss function for all minimisations at once
3. Construct intermediate variables for loss computation inside the loss function definition, so that TensorFlow can track the gradients (I think the minimize function wraps the loss function in a gradient tape or some such).
4. Define the loss function as part of a class so that intermediate variables can be stored.
5. minimize only does one step of the minimisation, so we need to loop over it lots of times until it converges according to some criterion.
6. I was running into some NaNs due to the invalid-ness of means less than zero for Poisson distributions. So I needed to add a constraint to the input variables.

With this, I can now do the equivalent of a million minimisations in like 10 seconds on my laptop, which is pretty nice!

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt

# Bunch of independent Poisson distributions that we want to combine
poises0 = [tfd.Poisson(rate = 10) for i in range(5)]

# Construct joint distributions
joint0 = tfd.JointDistributionSequential(poises0)

N = int(1e6)
samples0 = joint0.sample(N)

class Model(object):
  def __init__(self):
     self.mus = [tf.Variable(10*np.ones(N, dtype='float32'), name='mu{0}'.format(i),
                    constraint=lambda x: tf.clip_by_value(x, 0.000001, np.infty)) for i in range(5)]

  def loss(self):
     poises_free = [tfd.Poisson(rate = self.mus[i]) for i in range(5)]
     joint_free = tfd.JointDistributionSequential(poises_free)
     # Construct (half of) test statistic
     self.qM = -2*(joint_free.log_prob(samples0))
     self.last_loss = tf.math.reduce_sum(self.qM,axis=0)
     return self.last_loss

model = Model()

# Minimise
tol = 0.01 * N
delta_loss = 1e99
prev_loss = 1e99
i = 0
print("tol:", tol)
while delta_loss > tol:
    opt = tf.optimizers.SGD(0.1).minimize(model.loss,var_list=model.mus)
    delta_loss = np.abs(prev_loss - model.last_loss)
    print("i:", i," delta_loss:", delta_loss)
    i+=1
    prev_loss = model.last_loss

q0 =-2*joint0.log_prob(samples0)
q = q0 - model.qM

print("parameters:", model.mus)
print("loss:", model.last_loss)
print("q0:", q0)
print("qM:", model.qM)

fig = plt.figure()
ax = fig.add_subplot(111)
sns.distplot(q, kde=False, ax=ax, norm_hist=True)
qx = np.linspace(np.min(q),np.max(q),1000)
qy = np.exp(tfd.Chi2(df=5).log_prob(qx))
sns.lineplot(qx,qy)
plt.show()

Output:

tol: 10000.0
i: 0  delta_loss: inf
i: 1  delta_loss: 197840.0
i: 2  delta_loss: 189366.0
i: 3  delta_loss: 181456.0
i: 4  delta_loss: 174040.0
i: 5  delta_loss: 167042.0
i: 6  delta_loss: 160448.0
i: 7  delta_loss: 154216.0
i: 8  delta_loss: 148310.0
i: 9  delta_loss: 142696.0
i: 10  delta_loss: 137352.0
i: 11  delta_loss: 132268.0
i: 12  delta_loss: 127404.0
...
i: 69  delta_loss: 11894.0
i: 70  delta_loss: 11344.0
i: 71  delta_loss: 10824.0
i: 72  delta_loss: 10318.0
i: 73  delta_loss: 9860.0
parameters: [<tf.Variable 'mu0:0' shape=(1000000,) dtype=float32, numpy=
array([ 6.5849004, 14.81182  ,  7.506216 , ..., 10.       , 11.491933 ,
       10.760278 ], dtype=float32)>, <tf.Variable 'mu1:0' shape=(1000000,) dtype=float32, numpy=
array([12.881036,  7.506216, 12.881036, ...,  7.506216, 14.186232,
       10.760278], dtype=float32)>, <tf.Variable 'mu2:0' shape=(1000000,) dtype=float32, numpy=
array([16.01586  ,  8.378036 , 12.198007 , ...,  6.5849004, 12.198007 ,
        8.378036 ], dtype=float32)>, <tf.Variable 'mu3:0' shape=(1000000,) dtype=float32, numpy=
array([10.      ,  7.506216, 12.198007, ...,  9.207426, 10.760278,
       11.491933], dtype=float32)>, <tf.Variable 'mu4:0' shape=(1000000,) dtype=float32, numpy=
array([ 8.378036 , 14.81182  , 10.       , ...,  6.5849004, 12.198007 ,
       10.760278 ], dtype=float32)>]
loss: tf.Tensor(20760090.0, shape=(), dtype=float32)
q0: tf.Tensor([31.144037 31.440613 25.355555 ... 24.183338 27.195362 22.123463], shape=(1000000,), dtype=float32)
qM: tf.Tensor([21.74377  21.64162  21.526024 ... 19.488544 22.40428  21.08519 ], shape=(1000000,), dtype=float32)

Result is now chi-squared DOF=5! Or at least pretty close.