Summary

This post showcases a workaround to optimize a tf.keras.Model model with a TensorFlow-based L-BFGS optimizer from TensorFlow Probability. The complete code can be found at my GitHub Gist here.

Update (06/08/2020): I’ve updated the code on GitHub Gist to show how to save loss values into a list when using the @tf.function decorator. But I didn’t update the blog post here, so the code line numbers may not match the code on GitHub.

While SGD, Adam, and similar optimizers now dominate the training of deep neural networks, some people, including me, may still want to use second-order methods such as L-BFGS.

The problem is that TensorFlow 2.0 does not have L-BFGS. PyTorch provides L-BFGS, so using Keras with a PyTorch backend might sound like a possible workaround. But I don’t use original Keras. I use TensorFlow 2.0 and build Keras models with the tf.keras module. That means I cannot simply switch the backend. And it is difficult for me to migrate my code to original Keras because my code contains many customized parts that require TensorFlow 2.0.

Another workaround is to use the L-BFGS solver from SciPy to train a tf.keras.Model or its subclasses. We can find some example code for that approach through a quick search. The problem, however, is that standard SciPy is not GPU-capable. And many people like me use TensorFlow because we need GPU computing. In addition, I personally do not think SciPy is ideal for serious large-scale calculations. It is more of a prototyping tool than something I would want to run on an HPC cluster. That is just my personal opinion.

Fortunately, a TensorFlow-based L-BFGS solver exists in TensorFlow Probability. The API documentation for this solver is here. We can use it through something like import tensorflow_probability as tfp, and then call result = tfp.optimizer.lbfgs_minimize(...).

The returned object, result, contains several pieces of information, and the final optimized parameters are stored in result.position. If you are using a GPU version of TensorFlow, then this L-BFGS solver should also run on the GPU.

Apparently, the solver is not implemented as a subclass of tf.keras.optimizers.Optimizer. So we cannot use it directly with model.compile(...) and model.fit(...). The solver is just a function. We need some workaround or wrapper to use it.

Let’s first see the arguments of this function:

Arguments of tfp.optimizer.lbfgs_minimize

1
tfp.optimizer.lbfgs_minimize(
2
    value_and_gradients_function,
3
    initial_position,
4
    num_correction_pairs=10,
5
    tolerance=1e-08,
6
    x_tolerance=0,
7
    f_relative_tolerance=0,
8
    initial_inverse_hessian_estimate=None,
9
    max_iterations=50,
10
    parallel_iterations=1,
11
    stopping_condition=None,
12
    name=None
13
)

The value_and_gradients_function is a function, or a callable object, that returns the loss and the gradients with respect to parameters. This callable object should take in the parameters that we want to optimize, which in this case are the model’s trainable parameters, i.e., the kernels and biases of trainable layers.

The first notable thing is that value_and_gradients_function takes model parameters, not training data. The second thing is that the model parameters fed into value_and_gradients_function must be a 1D tf.Tensor.

But TensorFlow and Keras store trainable model parameters as a list of multidimensional tf.Variable objects. We can easily see this with print(model.trainable_variables), assuming model is an instance of tf.keras.Model or one of its subclasses. This means we need a way to transform a list of multidimensional tf.Variable objects into a single 1D tensor. This can be done with tf.dynamic_stitch.

We also need a way to convert a 1D tf.Tensor back to a list of multidimensional tf.Tensor, tf.Variable, or numpy.ndarray objects so that we can update the model parameters. Under most circumstances, we can treat tf.Variable and tf.Tensor similarly for this purpose.

There are two ways to update a model’s parameters:

1. We can use model.set_weights(params) to update the values of the parameters. In this case, params is a list of multidimensional numpy.ndarray. And we need to convert the aforementioned 1D tf.Tensor to params. In TensorFlow 2.0, this is not difficult because of the default eager execution behavior. We first partition the 1D tf.Tensor with tf.dynamic_partition to a list of tensors, convert the list of tensors to a list of numpy.ndarray with tensor.numpy(), and then reshape each array to the corresponding shape. Finally, we can call model.set_weights(params).
2. The other way is more computationally efficient and can work in graph mode without eager execution. Each element in model.trainable_variables is a tf.Variable, and it provides an assign method to update its value. So we first partition the 1D tf.Tensor into a list of tensors with tf.dynamic_partition. Next, we use a for loop to reshape each tensor and assign it to the corresponding tf.Variable in model.trainable_variables.

The third thing is that when returning the gradients, the gradients should also be a 1D tf.Tensor to match the expected format of value_and_gradients_function. This again can be done with tf.dynamic_stitch.

In addition, the argument initial_position of tfp.optimizer.lbfgs_minimize should contain the initial parameter values of the model. And of course, we should use tf.dynamic_stitch to convert the initial model parameters into a 1D tf.Tensor before passing them to initial_position.

In a nutshell, when we create a function for the value_and_gradients_function argument, that function should store the following information:

1. the tf.keras.Model model we want to use,
2. the loss function to use,
3. the training data that we want to evaluate the loss,
4. the information required by tf.dynamic_stitch to convert a list of multidimensional tf.Tensor to a 1D tf.Tensor, and
5. the information required by tf.dynamic_partition to convert a 1D tf.Tensor to a list of multidimensional tf.Tensor or numpy.ndarray.

We can define a function factory to create such a function for us. Here’s an example: Code snippet (tf_keras_tfp_lbfgs.py, lines 25-107)

And here’s an example of how to use this function factory together with tfp.optimizer.lbfgs_minimize to train a tf.keras.Model model: Code snippet (tf_keras_tfp_lbfgs.py, lines 131-152)

The complete example code can be found at my GitHub Gist here.

Finally, the example code is only meant to give a sense of how to use the L-BFGS solver from TensorFlow Probability. Using a function factory is not the only option. value_and_gradients_function can be any callable object. So we could also wrap the model in a Python class and implement __call__. It is entirely up to us.