Flux.jl: LBFGS Optimizer

[kczimm]

Are there any plans on implementing the LBFGS optimizer? It seems to perform better and faster than ADAM for small sample sizes in my experience with sklearn python. Having this implemented will help me switch my current application to Julia/FluxML.

[thebhatman]

@kczimm Implementation of LBFGS Optimizer can be found here. https://github.com/JuliaNLSolvers/Optim.jl

[kczimm]

@thebhatman thanks. I did see that. Any idea how I can train a network using that algorithm?

[kczimm]

It would be great if the optimizer could plugin like the Flux optimizers,

using Optim
using Flux

opt = Optim.LBFGS()
Flux.train!(loss, params(m), data, opt)

[improbable22]

I suspect it's easier to hook these up the other way around:

using Flux, Optim, CatViews

m = Chain(Dense(5,5,relu), softmax)
X = randn(5,10);
Y = Flux.onehotbatch(rand(1:5, 10), 1:5)
loss() = Flux.crossentropy(m(X),Y)

catm(f) = CatView(f.(Flux.params(m).order)...)

function fg!(F,G,v)
    copyto!(catm(Tracker.data), v)
    fwd = loss()
    Flux.back!(fwd)
    if G != nothing
        copyto!(G, catm(Tracker.grad))
    end
    return Tracker.data(fwd)
end

@show loss()
res = Optim.optimize(Optim.only_fg!(fg!), copy(catm(Tracker.data)), LBFGS(), Optim.Options(iterations=100))
@show loss()

But this doesn't work very well, right now. See also discussion at https://discourse.julialang.org/t/two-questions-on-flux/22307

[kczimm]

@improbable22 thanks for the information. I wonder if Flux could implement/expose in its api the LBFGS optimizer but just call out to Optim as you suggest in your code example. This is basically what Scikit-learn does. They implement their own SGD optimizers but call out to scipy optimize for the LBFGS optimizer.

[improbable22]

That looks pretty similar, if I'm reading it right -- fmin_l_bfgs_b(... is like Optim.optimize(... and gets handed the loss function, max iterations, etc. And X is the full input data, like mine?

Maybe the big difference is that Flux's optimisers typically expect at each step a different batch of the data. This seems quite unlike what LBFGS is designed for, e.g. I think the Optim version runs a linesearch alg to decide the step size, by some kind of bisection procedure, and this will not work so well if each evaluation is on different data. The discussion above linked to 1802.04310, which seems to be about finding ways around such issues. So my guess is that it's more than an API problem, but I could be wrong.

[kczimm]

That looks pretty similar, if I'm reading it right -- fmin_l_bfgs_b(... is like Optim.optimize(... and gets handed the loss function, max iterations, etc. And X is the full input data, like mine?

That's my understanding, as well.

Unfortunately, my understanding of the detailed differences between SGD type optimizers and BFGS is limited. Perhaps you're right, and this is more than an API problem. I'm fairly new to Flux and even Julia for that matter. I wish I knew more so I could try to plan a PR for adding LBFGS to Flux because it seems like a good optimizer to have available in a ML library like Flux.

[MikeInnes]

cc @pkofod. We may be able to unify the Optim and Flux optimisers to some extent.

[pkofod]

cc @pkofod. We may be able to unify the Optim and Flux optimisers to some extent.

I think that is quite feasible. We "just" need to agree on an interface for the action to "apply one iteration". It can be called many things, and can take many types of input, so we just need to agree on that.

[kczimm]

What can I do to help continue this conversation?

[DoktorMike]

I'm also quite interested in this, since Adam or indeed any other SGD method is not suitable for low dimensional full batch optimizations. Pytorch handles this internally by passing a complete closure to the LBFGS optimizer. Looks like this in their code:

for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        return loss
    optimizer.step(closure)

Also they've made it such that these two conditions need to hold to use the LBFGS:

This optimizer doesn’t support per-parameter options and parameter groups (there can be only one).

Right now all parameters have to be on a single device. This will be improved in the future.

There's more information here: https://pytorch.org/docs/stable/optim.html

[improbable22]

There's a package by @baggepinnen which ought to be linked here: https://github.com/baggepinnen/FluxOptTools.jl