Probability: tfp.optimizer.lbfgs_minimize as Keras optimizer?

[Hoeze]

Hi, is there a way to use tfp.optimizer.lbfgs_minimize as Keras optimizer?
This would be quite useful in certain cases where the loss function is approximately quadratic.

A colleague of mine would very much need it since an autoencoder written in R with negative-binomial loss converges faster than its Keras counterpart.

[dynamicwebpaige]

Thanks for the recommendation, @Hoeze! Are you using stand-alone Keras for your work, or tf.keras?

[Hoeze]

@dynamicwebpaige Thanks for your answer.
We are usually using the default Keras API with Tensorflow as backend.
For special needs, we construct Keras layers with the Tensorflow API.

I would be cool if it was implemented as a Tensorflow optimizer similar to ADAM:
https://github.com/tensorflow/tensorflow/blob/r2.1/tensorflow/python/training/adam.py
This way it would be more easy to use in plain Tensorflow as well, since switching to L-BFGS would only incorporate changing a single line of code.

The big problem with tfp.optimizer.lbfgs_minimize was that we had to provide it with a function returning the loss. Then it does some magic and eventually the result has converged. One cannot record the single optimization steps.
I am not sure how this black-box optimization fits with Keras, since e.g. the EarlyStopping callback would not work this way.

[dario-passos]

@Hoeze I was looking for the same functionality and I found this blog that show how to use the lbfgs_minimize() with a tf.keras model: https://pychao.com/2019/11/02/optimize-tensorflow-keras-models-with-l-bfgs-from-tensorflow-probability/

However it would be very useful if the tf.keras development team could implement this as a Tensorflow optimizer in a future update, just like you suggest!!
Cheers