Using Autograd for Maximum Likelihood Estimation

import warnings
warnings.filterwarnings('ignore')
import autograd.numpy as np
from autograd import grad, jacobian, hessian
from autograd.scipy.stats import norm
from scipy.optimize import minimize
import statsmodels.api as sm
import pandas as pd
import timeit

Here we define a simple function and show the basics of how autograd works. Note that the package is a "drop-in" replacement for numpy that supports many numpy methods and some scipy distributions. Define a simple Cobb-Douglas production function demonstrating constant return to scale in two variables \(x\) and \(y\) as:

Define the function we want to auto-differentiate:

def f(x, y):
    return x**(.8) * y**(.2)

For this problem, calculating the analytical derivatives isn't hard:

but we can now use the grad function to take the first derivative without having to program these equations in python. For functions with more than 1 variable, we can ask for each first derivative as shown below:

# first derivatives - for f(x,y), x is position 0 (default)  
#                         and y is position 1
dfdx = grad(f)
dfdy = grad(f, 1)

# suppose values for x and y are as follows
x, y = 2.0, 3.0

# evaluate the gradiant at x, y
derivs = np.array([dfdx(x,y), dfdy(x,y)])
print("Gradiant: ", derivs)

Gradiant:  [0.86757742 0.14459624]

As a sanity check, let's program in the analytical derivatives and compare to what autograd gives us:

# analytical derivatives
a_dfdx = lambda x, y: .8*x**(-.2)*y**(.2)
a_dfdy = lambda x, y: .2*x**(.8)*y**(-.8)

print("Analytical Gradiant: ", np.array([ a_dfdx(x,y), a_dfdy(x,y)])) 

Analytical Gradiant:  [0.86757742 0.14459624]

That is fantastic. The package autograd gives us analytical derivates without the pain and suffering of deriving them and programming them into python!

We can take derivatives of derivatives (to uncover elements of the matrix of 2^nd derivatives):

# second derivates
d2fdxdx = grad(dfdx)
d2fdydy = grad(dfdy, 1)
d2fdxdy = grad(dfdx, 1)
d2fdydx = grad(dfdy)

Next we turn to a more realistic example based on my previous post on maximum likelihood methods in python linked above. Consider the population regression equation

And we have a sample of \(N=5000\) observations, where \(\beta\) the matrix of parameters is dimension \(K \times 1\) having an intercept.

# number of observations
N = 5000
# number of parameters
K = 10
# true parameter values
beta = 2 * np.random.randn(K)
# true error std deviation
sigma =  2

def datagen(N, beta, sigma):
    """
    Generates data for OLS regression.
    Inputs:
    N: Number of observations
    beta: K x 1 true parameter values
    sigma: std dev of error
    """
    K = beta.shape[0]
    x_ = 10 + 2 * np.random.randn(N,K-1)
    # x is the N x K data matrix with column of ones
    #   in the first position for estimating a constant
    x = np.c_[np.ones(N),x_]
    # y is the N x 1 vector of dependent variables
    y = x.dot(beta) + sigma*np.random.randn(N)
    return y, x

y, x  = datagen(N, beta, sigma)

For using auto-differentiation and autograd, we need to define the negative log-likelihood function (the function we are minimizing). For ordinary least squares regression, the negative log-likelihood function is given by

def neg_loglike(theta):
    beta = theta[:-1]
    sigma = theta[-1]
    mu = np.dot(x,beta)
    ll = -N/2 * np.log(2*np.pi*sigma**2) - (1/(2*sigma**2)) * np.sum((y - mu)**2)
    return -1 * ll

Note, that the function is defined using only numpy fundamental methods making it autograd ready.

However, autograd has support for some scipy distributions, including the normal distribution. Using this built-in logpdf function is easier to work with and it is optimized for fast auto-differentiation so we will work with this version of our log-likelihood rather than the one above:

def neg_loglike(theta):
    beta = theta[:-1]
    # transform theta[-1]
    # so that sigma > 0
    sigma = np.exp(theta[-1])
    mu = np.dot(x,beta)
    ll = norm.logpdf(y,mu,sigma).sum()
    return -1 * ll

We can then use the jacobian and hessian functions from autograd to evaluate the gradiant and hessian at particular values of our parameters. Here we evaluate the gradiant at true parameter values \(\pmb{\theta} = \begin{bmatrix} \pmb{\beta} & \sigma \end{bmatrix}\)

# derivates of neg_loglike
jacobian_  = jacobian(neg_loglike)
hessian_ = hessian(neg_loglike)

# evaluate the gradiant at true theta
# theta = [beta log(sigma)]
theta = np.append(beta,np.log(sigma))
print(jacobian_(theta))

[ -42.89535311 -476.79418971 -490.35924314 -492.17676195 -424.97595445
 -557.19375403 -471.38236042 -459.67877535 -601.13561765 -429.86108545
  154.86658456]

We will use the minimize function from scipy for finding the maximum likelihood estimates. Depending on our choice of optimization algorithm, the minimize function can accept a jacobian and sometimes a hessian. For example, the "BFGS" algorithm for unconstrained problems accepts a jacobian and we will use jacobian_ defined above using autograd. This code sets up the minimization problem and stores results in res1:

theta_start = np.append(np.zeros(beta.shape[0]),0.0)
res1 = minimize(neg_loglike, theta_start, method = 'BFGS', \
               options={'disp': False}, jac = jacobian_)
print("Convergence Achieved: ", res1.success)
print("Number of Function Evaluations: ", res1.nfev)

Convergence Achieved:  True
Number of Function Evaluations:  68

For looking at the results, we'll calculate standard errors using hessian_ also from autograd and assemble everything in a dataframe for viewing a little later:

# estimated parameters
theta_autograd = res1.x

# for std errors, calculate the information matrix
# using the autograd hessian
information1 = np.transpose(hessian_(theta_autograd))
se1 = np.sqrt(np.diagonal(np.linalg.inv(information1))) 

# Put Results in a DataFrame
results_a = pd.DataFrame({'Parameter':theta_autograd,'Std Err':se1})
names = ['beta_'+str(i) for i in range(K)]
names.append('log(Sigma)')
results_a['Variable'] = names
results_a['Model'] = "MLE Autograd"

Notice that at the estimated parameter vector, the gradiant is very close to zero for our \(K+1\) parameters:

print(jacobian_(theta_autograd))

[2.14390006e-07 1.68197130e-06 1.74329119e-06 1.78678499e-06
 1.96383104e-06 2.10519761e-06 2.08673862e-06 2.70411755e-06
 2.05865251e-06 2.90452056e-06 3.56323801e-07]

Another big benefit of using autograd is that gradiants don't need to be numerically approximated using finite differences. Therefore, the number of function evaluations will be much lower. In the example above with 11 parameters, estimating the model with autograd required 68 function evaluations whereas using finite differences required 1207. If the likelihood calculation is expensive (ours isn't for this problem), then these additional evaluations can really slow down estimation. To see how this translates into runtimes, we loop the estimation for \(K = \begin{bmatrix} 10&50&100\end{bmatrix}\) and examine execution times both with and without autograd.

for k in [10, 50, 100]:
     beta = np.random.randn(k)
     theta_start = np.append(beta,1.0)
     y, x = datagen(N, beta, 1)
     print("Time for " + str(k) +" Parameters")
     %timeit minimize(neg_loglike, theta_start, method = 'BFGS', \
                options={'disp': False}, jac = jacobian_)
     %timeit minimize(neg_loglike, theta_start, method = 'BFGS', \
                options={'disp': False})    

Putting the various run times together, the following tables shows average runtimes in seconds:

df = pd.DataFrame(.001*np.c_[[214, 573, 1190],[961, 12000, 46600]],columns=['With Autograd','Without Autograd'],index=[10, 50, 100])
df['Speed Up'] = df['Without Autograd']/df['With Autograd']
df.index.name = "K"
df.head()

     With Autograd  Without Autograd   Speed Up
K                                              
10           0.214             0.961   4.490654
50           0.573            12.000  20.942408
100          1.190            46.600  39.159664

As the number of function evaluations increases (more \(K\)), the execution time without autograd really increases. For the most expensive problem considered here, maximum likelihood estimation with autograd was nearly 40 times faster. It should be noted that even if we compare the "BFGS" results using the jacobian from autograd to gradiant free methods like "Nelder Mead" (results not reported here), we still see an approximate 10x speed up using autograd that remained roughly constant even as \(K\) increased.

We should all be using autograd or some other method of auto-differentiation. It is probably going to be faster, definitely more accurate, and only requires a little bit of additional coding.