Computing hessian-vector products in tensorflow

poiss_tf.py

"""
Computing hessian-vector products in tensorflow.

For simplicity, we demonstrate the idea on a Poisson regression model.
"""

import tensorflow as tf
import numpy as np
from scipy.optimize import minimize

n = 1000   # number of datapoints
p = 5      # number of features

# Create data.
#   `X` is a matrix of independent/predictor variables.
#   `true_b` are the ground-truth log inputs to the Poisson distribution.
#   `y` is are the dependent/predicted variables (count data).
X = .3 * np.random.randn(n, p)
true_b = np.random.randn(p,)
y = np.random.poisson(np.exp(np.dot(X, true_b))).astype(float)

# Create placeholder for model coefficients.
b = tf.placeholder(shape=(p,), dtype=tf.float64)

# Compute loss and gradient.
Xb = tf.squeeze(tf.matmul(tf.constant(X), b[:, None]))
loss = tf.reduce_mean(tf.nn.log_poisson_loss(tf.constant(y), Xb))
grad = tf.gradients(loss, b)

# Create new placeholder for vector multiplied with Hessian.
z = tf.placeholder(shape=(p,), dtype=tf.float64)
hess_vec = tf.gradients(tf.reduce_sum(grad * z), b)

# Loss function at estimate b_
def f(b_):
    return sess.run(loss, feed_dict={b: b_})

# Gradient at estimate b_
def g(b_):
    return sess.run(grad, feed_dict={b: b_})[0]

# Hessian vector product at estimate b_
def hp(b_, z_):
    return sess.run(hess_vec, feed_dict={b: b_, z: z_})[0]

# Optimize
sess = tf.Session()
result = minimize(f, np.zeros(p), jac=g, hessp=hp, method='newton-cg')

print('True regression coeffs: {}'.format(true_b))
print('Estimated regression coeffs: {}'.format(result.x))

True regression coeffs: [-1.51245513 -0.82369068 -0.41878618  0.81168893 -1.40940972]
Estimated regression coeffs: [-1.48243638 -0.81270721 -0.3724834   0.93704033 -1.38154464]