How Jacobian matrix helps optimization faster?

Question

I tried some python optimization functions and some of them needed Jacobian matrix prior for faster convergence. I understand Jacobians are basically transformation matrices that data from one space to another or carrying gradients information. Can someone explain me with some literature references, how the speed up is achieved?

Answer 1

You haven't told us exactly what optimization routine you're using, so it's difficult to provide a very specific answer to your question.

However, if you don't supply your own Jacobian function then the optimization routine that you're using is presumably approximating the derivatives using a finite difference approximation scheme. Computing finite difference approximations to the derivative requires many function evaluations and this slows down the optimization process. Furthermore, the inaccuracy of such approximate derivatives can cause the algorithm to require more steps to converge to an optimal solution and thus run more slowly.

The advantages of using exact analytical derivatives rather than finite difference approximations are discussed in most textbooks on nonlinear programming. See for example Practical Optimization by Gill, Murray, and Wright.

Answer 2

even in very simple example you can see decrease of time execution, or speed up of the code's execution

import time
import numpy as np
from scipy.optimize import minimize

# Define the objective function
def objective(x):
    return np.sum((x - 1)**2)

# Define the initial guess
x0 = np.array([0.5, 0.5])

# Optimize without Jacobian matrix
tic = time.time()
res1 = minimize(objective, x0, method='Nelder-Mead')
tac = time.time()
print('>> It took %.3f s. to compute without the Jacobian matrix' % (tac-tic))
print(res1.x)


# Optimize with Jacobian matrix
def jacobian(x):
    return 2*(x - 1)
tic = time.time()
res2 = minimize(objective, x0, method='BFGS', jac=jacobian)
tac = time.time()
print(res2.x)
print('>> It took %.3f s. to compute with the Jacobian matrix' % (tac-tic))

P.S. can test (in Colab) your keras_model speed with jacobian use like here - jacobian is being calculated fast enough

import numpy as np
import time

N = 500  # Input size
H = 100  # Hidden layer size
M = 10   # Output size

w1 = np.random.randn(N, H)  # first affine layer weights
b1 = np.random.randn(H)     # first affine layer bias
w2 = np.random.randn(H, M)  # second affine layer weights
b2 = np.random.randn(M)     # second affine layer bias

## Using Keras, we implement our network as follows:
import tensorflow as tf
from tensorflow.keras.layers import Dense
tf.compat.v1.disable_eager_execution()

# Keras as a simplified interface to TensorFlow: tutorial
# https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html
from keras import backend as K
sess = K.get_session()

# sess = tf.Session()

model = tf.keras.Sequential()
model.add(Dense(H, activation='relu', use_bias=True, input_dim=N))
model.add(Dense(M, activation='softmax', use_bias=True, input_dim=H))
model.get_layer(index=0).set_weights([w1, b1])
model.get_layer(index=1).set_weights([w2, b2])

##Let’s now try to compute the Jacobian of this model. 
## Unfortunately, Tensorflow does not currently provide a method to compute Jacobian matrices out-of-the-box. 
## The method tf.gradients(ys, xs) returnssum(dy/dx) for each x in xs, which in our case would be N-dimensional vector containing the sums of the Jacobian rows; not quite what we are looking for (see this issue). However, we can still compute our Jacobian matrix, by computing the gradients vectors for each yi, and grouping the output into a matrix:
def jacobian_tensorflow(x):
    jacobian_matrix = []
    for m in range(M):
        # We iterate over the M elements of the output vector
        grad_func = tf.gradients(model.output[:, m], model.input)
        gradients = sess.run(grad_func, feed_dict={model.input: x.reshape((1, x.size))})
        jacobian_matrix.append(gradients[0][0,:])

    return np.array(jacobian_matrix)

##Let’s make sure the Jacobian we compute is actually correct by checking it using numerical differentiation. The function is_jacobian_correct() below takes in argument a function computing the Jacobian and the feedforward function f:
def is_jacobian_correct(jacobian_fn, ffpass_fn):
    # Check of the Jacobian using numerical differentiation

    x = np.random.random((N,))
    epsilon = 1e-5    # Check a few columns at random

    for idx in np.random.choice(N, 5, replace=False):
        x2 = x.copy()
        x2[idx] += epsilon
        num_jacobian = (ffpass_fn(x2) - ffpass_fn(x)) / epsilon
        computed_jacobian = jacobian_fn(x)

        if not all(abs(computed_jacobian[:, idx] - num_jacobian) < 1e-3):
            return False
    return True

def ffpass_tf(x):
    # The feedforward function of our neural net

    xr = x.reshape((1, x.size))
    return model.predict(xr)[0]

print(is_jacobian_correct(jacobian_tensorflow, ffpass_tf))
# >> True     ##Very good, it’s correct. 

# Let’s see how long this computation takes:
x0 = np.random.random((N,))
tic = time.time()
jacobian_tf = jacobian_tensorflow(x0)   # , verbose=False
tac = time.time()
print('It took %.3f s. to compute the Jacobian matrix' % (tac-tic))

even faster could be with autograd, or even jax libs for numerical differentiation (res of article's tests):

where our Tensorflow implementation took about 650 ms, Autograd needs 3.7 ms, giving us a ~170x speed increase in this case. Of course, it’s not always convenient to specify one’s model using Numpy, because Tensorflow and Keras provide a lot of useful functions and training facilities out-of-the-box…

still pure Numpy is even faster (115 µs) [as it is written in C/C++]