Is there a tf.keras.optimizers implementation for L-BFGS?

Viewed 2246

Does anybody have a Tensorflow 2 tf.keras subclass for the L-BFGS algorithm? If one wants to use L-BFGS, one has currently two (official) options:

  1. TF Probability
  2. SciPy optimization

These two options are quite cumbersome to use, especially when using custom models. So I am planning to implement a custom subclass of tf.keras.optimizers to use L-BFGS. But before I start, I was curious, whether somebody already tackled this task?

2 Answers

I've implemented an interface between keras and SciPy optimize. https://github.com/pedro-r-marques/keras-opt

I'm using 'cg' by default but you should also be able to use 'l-bfgs'. Take a look at the unit tests for example usage. I will add documentation as soon as possible.

Does anybody have a Tensorflow 2 tf.keras subclass for the L-BFGS algorithm?

Yes, here's (yet another) implementation L-BFGS (and any other scipy.optimize.minimize solver) for your consideration in case it fits your use case:

This package has a similar goal to Pedro's answer above, but I would recommend it over the keras-opt package if you run into issues with memory consumption during training. I implemented kormos when trying to build a Rendle-type factorization machine and kept OOMing with other full-batch solver implementations.

These two options are quite cumbersome to use, especially when using custom models. So I am planning to implement a custom subclass of tf.keras.optimizers to use L-BFGS. But before I start, I was curious, whether somebody already tackled this task?

Agreed, it's a little cumbersome to fit the signatures of tfp and scipy into the parameter fitting procedure in keras, because of the way that keras steps in and out of an optimizer that has persistent state between calls, which is not how most [old school?] optimization libraries work.

This is addressed specifically in the kormos package since IMO during prototyping it's a pretty common workflow to alternate between either a stochastic optimizer and a full-batch deterministic optimizer, and this should be simple enough to do ad hoc in the python interpreter.

The package has models that extend keras.Model and keras.Sequential:

These can be compiled to be fit with either the standard or the scipy solvers; it would look something like this:

from tensorflow import keras
from kormos.models import BatchOptimizedSequentialModel

# Create an Ordinary Least Squares regressor
model = BatchOptimizedSequentialModel()
model.add(keras.layers.Dense(
  units=1,
  input_shape=(5,),
))

# compile the model for stochastic optimization
model.compile(loss=keras.losses.MeanSquaredError(), optimizer="sgd")
model.fit(...)

# compile the model for deterministic optimization using scipy.optimize.minimize
model.compile(loss=keras.losses.MeanSquaredError(), optimizer="L-BFGS-B")
model.fit(...)
Related